feat(refit): integrate ModelExpress weight synchronization

nemo_rl/algorithms/grpo.py

@@ -1311,7 +1311,7 @@ def initialize_generation_with_policy(
             )
             assert remote_transport is not None
             remote_synchronizer_cls = VllmRemoteSparseWeightSynchronizer
-        elif refit_transport is not None and refit_transport != "nccl_reshard":
+        elif refit_transport not in {None, "nccl_reshard", "model_express"}:
             # nccl_reshard is handled below via nccl_reshard_refit_enabled,
             # not via checkpoint-engine.
             checkpoint_engine_config = checkpoint_engine_refit_config(generation_config)
@@ -1513,6 +1513,9 @@ def init_dynamo():
     nccl_reshard_refit_enabled = (
         generation_config.get("refit_transport") == "nccl_reshard"
     )
+    model_express_refit_enabled = (
+        generation_config.get("refit_transport") == "model_express"
+    )
     if nccl_reshard_refit_enabled:
         from nemo_rl.weight_sync.nccl_reshard_utils import (
             check_nccl_reshard_refit_support,
@@ -1551,6 +1554,7 @@ def init_dynamo():
         not colocated_inference
         and remote_transport is None
         and checkpoint_engine_config is None
+        and not model_express_refit_enabled
     ):
         t0 = time.perf_counter()
         # init collective
@@ -1582,7 +1586,21 @@ def init_dynamo():
             ray.get(futures_train + futures_inference)
         setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0
 
-    if remote_transport is not None:
+    if model_express_refit_enabled:
+        t0 = time.perf_counter()
+        policy_generation.weight_synchronizer = create_weight_synchronizer(
+            policy=policy,
+            generation=policy_generation,
+            generation_backend=backend,
+            colocated=colocated_inference,
+            train_cluster=train_cluster,
+            inference_cluster=inference_cluster,
+        )
+        policy_generation.weight_synchronizer.init_communicator()
+        setup_timing_metrics.extras["model_express_init_time_s"] = (
+            time.perf_counter() - t0
+        )
+    elif remote_transport is not None:
         t0 = time.perf_counter()
         assert isinstance(policy_generation, VllmGeneration)
         assert remote_synchronizer_cls is not None

nemo_rl/models/generation/vllm/config.py

@@ -19,7 +19,13 @@
 from nemo_rl.models.generation.interfaces import GenerationConfig
 
 VllmRefitTransportName = Literal["s3", "zmq"]
-VllmRefitSelector = Literal["vllm_s3_sparse", "vllm_zmq_sparse", "nixl", "nccl_reshard"]
+VllmRefitSelector = Literal[
+    "vllm_s3_sparse",
+    "vllm_zmq_sparse",
+    "nixl",
+    "nccl_reshard",
+    "model_express",
+]
 VLLM_SPARSE_REFIT_TRANSPORTS = frozenset({"vllm_s3_sparse", "vllm_zmq_sparse"})
 
 
@@ -134,6 +140,10 @@ class VllmNixlRefitConfig(BaseModel, extra="forbid"):
     shard_expert_weights: bool = False
 
 
+class VllmModelExpressRefitConfig(BaseModel, extra="forbid"):
+    server_url: str | None = None
+
+
 class VllmCheckpointEnginePluginConfig(BaseModel, extra="allow"):
     update_weights_bucket_memory_ratio: Annotated[float, Field(gt=0, lt=1)] = 0.05
     release_after_refit: bool = False
@@ -142,6 +152,9 @@ class VllmCheckpointEnginePluginConfig(BaseModel, extra="allow"):
 class VllmRefitConfig(BaseModel, extra="allow"):
     sparse: VllmSparseRefitConfig = Field(default_factory=VllmSparseRefitConfig)
     nixl: VllmNixlRefitConfig = Field(default_factory=VllmNixlRefitConfig)
+    model_express: VllmModelExpressRefitConfig = Field(
+        default_factory=VllmModelExpressRefitConfig
+    )
 
 
 class VllmConfig(GenerationConfig):
@@ -182,7 +195,8 @@ def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None:
     if transport not in get_args(VllmRefitSelector) and ":" not in transport:
         raise ValueError(
             f"Unknown vLLM refit transport {transport!r}: expected null, "
-            "'nccl_reshard', 'vllm_s3_sparse', 'vllm_zmq_sparse', 'nixl', or a "
+            "'nccl_reshard', 'vllm_s3_sparse', 'vllm_zmq_sparse', 'nixl', "
+            "'model_express', or a "
             "'module:ClassName' checkpoint-engine path."
         )
     # The encoder-cache reset is implemented only on the collective/IPC and

nemo_rl/models/generation/vllm/vllm_backend.py

@@ -11,17 +11,23 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 import gc
 import logging
+import os
 import re
 import socket
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from contextlib import contextmanager
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import torch
 import zmq
 
+if TYPE_CHECKING:
+    from modelexpress_rl import ModelExpressGeneratorClient, WeightVersionRef
+
 from nemo_rl.models.generation.vllm.checkpoint_engine import (
     VllmCheckpointEngineMixin,
     preinit_nixl_from_vllm_config,
@@ -184,6 +190,7 @@ class VllmInternalWorkerExtension:
     _mtp_drafter_from_disk: bool = False
     _sparse_delta_applier: Any = None
     _nrl_named_parameters: dict[str, torch.nn.Parameter]
+    _model_express: ModelExpressGeneratorClient | None = None
 
     def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]:
         params = getattr(self, "_nrl_named_parameters", None)
@@ -192,6 +199,38 @@ def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]:
             self._nrl_named_parameters = params
         return params
 
+    def initialize_model_express(self, server_url: str | None = None) -> None:
+        """Initialize ModelExpress inside the vLLM rank that owns live weights."""
+        if self._model_express is not None:
+            return
+        from modelexpress_rl import (
+            ModelExpressGeneratorClient,
+            ModelExpressGeneratorConfig,
+            VllmGeneratorContext,
+        )
+
+        self._model_express = ModelExpressGeneratorClient.initialize(
+            ModelExpressGeneratorConfig(
+                engine_context=VllmGeneratorContext(
+                    model=self.model_runner.model,
+                    vllm_config=self.model_runner.vllm_config,
+                ),
+                model_name=self.model_runner.model_config.model,
+                server_url=server_url,
+            )
+        )
+
+    def update_weights_from_model_express(self, version: WeightVersionRef) -> bool:
+        """Stage, verify, and install an exact MX version at a safe point."""
+        if self._model_express is None:
+            raise RuntimeError("ModelExpress generator client is not initialized")
+        staged = self._model_express.stage_weight(version=version)
+        try:
+            self._model_express.apply_weight(staged)
+        finally:
+            staged.release()
+        return True
+
     def _load_full_hf_weights(
         self, policy_weights: list[tuple[str, torch.Tensor]]
     ) -> None:
@@ -1108,6 +1147,9 @@ def _receive_and_load_misc_params(self) -> None:
 
     def cleanup(self) -> None:
         """Shutdown and cleanup resources."""
+        if self._model_express is not None:
+            self._model_express.close()
+            self._model_express = None
         # Close ZMQ socket and context if they exist
         if hasattr(self, "zmq_socket"):
             self.zmq_socket.close()

nemo_rl/models/generation/vllm/vllm_generation.py

@@ -1082,6 +1082,34 @@ def update_weights_via_ipc_zmq(self) -> list[ray.ObjectRef]:
         # this function should co-work with lm_policy, so we should wait for all futures to complete outside
         return futures
 
+    def initialize_model_express(self, *, server_url: str | None = None) -> None:
+        """Initialize an MX generator client in every vLLM rank."""
+        method_name = (
+            "initialize_model_express_async"
+            if self.cfg["vllm_cfg"]["async_engine"]
+            else "initialize_model_express"
+        )
+        futures = self.worker_group.run_all_workers_single_data(
+            method_name, server_url=server_url
+        )
+        ray.get(futures)
+
+    def update_weights_from_model_express(self, version: Any) -> None:
+        """Install an exact MX version in every vLLM rank."""
+        method_name = (
+            "update_weights_from_model_express_async"
+            if self.cfg["vllm_cfg"]["async_engine"]
+            else "update_weights_from_model_express"
+        )
+        futures = self.worker_group.run_all_workers_single_data(
+            method_name, version=version
+        )
+        results = ray.get(futures)
+        if not results or any(result is not True for result in results):
+            raise RuntimeError(
+                f"one or more vLLM engines failed ModelExpress refit: {results}"
+            )
+
     def update_weights_from_collective(self) -> list[ray.ObjectRef]:
         """Update weights of the policy using collective communication."""
         if not self.worker_group or not self.worker_group.workers:

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -212,10 +212,16 @@ def configure_worker(
                 engine_index_on_node = (
                     local_bundle_indices[0] % num_gpus_per_node
                 ) // mp_size
-            env_vars["VLLM_PORT"] = str(
+            vllm_port = (
                 DEFAULT_VLLM_PORT_RANGE_LOW
                 + engine_index_on_node * DEFAULT_VLLM_PORTS_PER_ENGINE
             )
+            env_vars["VLLM_PORT"] = str(vllm_port)
+            # MX inference derives rank-local listeners from these configured
+            # bases plus the CUDA device id. Keep them in unused portions of
+            # this engine's deterministic 100-port reservation.
+            env_vars["MX_METADATA_PORT"] = str(vllm_port + 64)
+            env_vars["MX_WORKER_GRPC_PORT"] = str(vllm_port + 80)
 
         # Check if this worker is part of a parallel group (TP or TP+PP).
         # A worker is part of a parallel group if it's a secondary member (local_bundle_indices is None)
@@ -320,6 +326,7 @@ def _init_config(
         self.precision = self.cfg["vllm_cfg"]["precision"]
         self.fraction_of_gpus = fraction_of_gpus
         self.is_model_owner = bundle_indices is not None
+        self._bundle_indices = bundle_indices
         self._extra_env_vars = extra_env_vars
 
         # Store the Python executable being used by this worker
@@ -1146,6 +1153,27 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None:
         """Prepare the info for refit."""
         self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,))
 
+    def initialize_model_express(self, *, server_url: str | None = None) -> None:
+        """Initialize the rank-local MX clients owned by this vLLM engine."""
+        assert self.llm is not None, "vLLM must be initialized before ModelExpress"
+        self.llm.collective_rpc(
+            "initialize_model_express",
+            args=(server_url,),
+        )
+
+    def update_weights_from_model_express(self, *, version: Any) -> bool:
+        """Apply one exact MX version on every internal vLLM rank."""
+        assert self.llm is not None, "vLLM must be initialized before ModelExpress"
+        results = cast(
+            list[bool],
+            self.llm.collective_rpc(
+                "update_weights_from_model_express", args=(version,)
+            ),
+        )
+        if not results or not all(results):
+            raise RuntimeError(f"ModelExpress update failed: {results}")
+        return True
+
     @wrap_with_nvtx_name("vllm_genertion_worker/update_weights_via_ipc_zmq")
     def update_weights_via_ipc_zmq(self) -> bool:
         """Update weights from IPC handles via ZMQ socket."""

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -1407,6 +1407,30 @@ async def prepare_refit_info_async(self, state_dict_info: dict[str, Any]) -> Non
         """Async version of prepare_refit_info."""
         await self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,))
 
+    async def initialize_model_express_async(
+        self, *, server_url: str | None = None
+    ) -> None:
+        """Initialize the rank-local MX clients owned by this vLLM engine."""
+        assert self.llm is not None, "vLLM must be initialized before ModelExpress"
+        await self.llm.collective_rpc(
+            "initialize_model_express",
+            args=(server_url,),
+        )
+
+    async def update_weights_from_model_express_async(self, *, version: Any) -> bool:
+        """Apply one exact MX version on every internal vLLM rank."""
+        assert self.llm is not None, "vLLM must be initialized before ModelExpress"
+        results = await self.llm.collective_rpc(
+            "update_weights_from_model_express", args=(version,)
+        )
+        if asyncio.iscoroutine(results):
+            results = await results
+        worker_results = cast(list[bool], results)
+        if not worker_results or not all(worker_results):
+            raise RuntimeError(f"ModelExpress update failed: {worker_results}")
+        await self._reset_encoder_cache_after_weight_update()
+        return True
+
     async def _reset_encoder_cache_after_weight_update(self) -> None:
         """Invalidate weight-dependent multimodal encoder outputs when enabled."""
         if not self.cfg["vllm_cfg"].get(

nemo_rl/models/policy/lm_policy.py

@@ -400,6 +400,28 @@ def run_all_workers_multiple_data(self, method_name: str, *args, **kwargs) -> An
         results = ray.get(futures)
         return results
 
+    def initialize_model_express(self, *, server_url: str | None = None) -> list[str]:
+        """Initialize rank-local MX trainer clients and return their source slots."""
+        return [
+            slot
+            for slot in self.run_all_workers_single_data(
+                "initialize_model_express", server_url=server_url
+            )
+            if slot is not None
+        ]
+
+    def publish_model_express_version(self, version: Any) -> None:
+        """Publish one global version from every trainer rank."""
+        self.run_all_workers_single_data(
+            "publish_model_express_version", version=version
+        )
+
+    def release_model_express_version(self, version: Any) -> None:
+        """Withdraw every trainer shard after its version is retired."""
+        self.run_all_workers_single_data(
+            "release_model_express_version", version=version
+        )
+
     def init_collective(
         self,
         ip: str,