feat(refit): add a ModelExpress-brokered nccl_reshard transport

docs/design-docs/nccl-reshard-refit.md

@@ -233,6 +233,43 @@ names/attributes change):
    contracts, never your layout.
 4. The misc producer/consumer — reuses the conventional packed-broadcast path.
 
+## ModelExpress-brokered rendezvous (optional)
+
+`refit_transport=mx_nccl_reshard` keeps everything above unchanged and replaces only
+how the NCCL communicators are bootstrapped. The bulk path still ends in
+`nccl.m2n.reshard` with the same meshes and placements, and the misc path still rides
+the packed broadcast, so the two transports move identical bytes and can be compared
+directly.
+
+What changes is the rendezvous. The native path uses `StatelessProcessGroup`, a
+`TCPStore` whose only job is to move 128 bytes of `ncclUniqueId` from rank 0 to
+everyone else, at an IP and port the driver allocates per PP stage. Routing that
+through a ModelExpress server adds three things that store cannot provide:
+
+* **Admission** against an expected participant set, so a worker that never joins
+  becomes a bounded failure naming the missing slot rather than a hang.
+* **Fencing** of worker generations, so a worker that dies and restarts is admitted as
+  a new generation instead of silently rejoining a formed group.
+* **Observable readiness**, so the trainer can confirm the generators it is about to
+  push into have prepared their destinations before it enters the collective.
+
+It also removes the per-stage port allocation from the driver, and lets a refit target
+a *subset* of the generators rather than all of them, because membership is explicit.
+
+Configuration:
+
+```bash
+uv run ./examples/run_grpo.py \
+  --config <your_config>.yaml \
+  policy.generation.colocated.enabled=false \
+  policy.generation.refit_transport=mx_nccl_reshard \
+  policy.generation.mx_server_url=<modelexpress-server>:8001
+```
+
+Requires the ModelExpress collective client (`modelexpress_rl.collective`). Every worker
+on both sides must be given the same server address; two different addresses form two
+groups, neither of which reaches READY.
+
 ## `xferdtensor` Transports
 
 `xferdtensor()` (in `nemo_rl/weight_sync/xferdtensor.py`) is the single entry point both

examples/run_grpo.py

@@ -17,6 +17,7 @@
 import pprint
 import time
 
+import ray
 from omegaconf import OmegaConf
 
 from nemo_rl.algorithms.grpo import (
@@ -54,6 +55,44 @@ def _select_trainer(master_config: MasterConfig):
     return grpo_train
 
 
+def _shutdown_runtime(
+    policy,
+    policy_generation,
+    cluster,
+    teacher_worker_groups,
+) -> None:
+    """Drain every Ray owner before disconnecting the driver.
+
+    Leaving policy or cluster cleanup to ``__del__`` races Ray's atexit hook:
+    a late ``ray.get``/``ray.kill`` then auto-initializes a second CoreWorker
+    while Python is finalizing. Explicit, idempotent shutdown keeps teardown
+    ordered and gives the launcher a truthful zero exit after a successful run.
+    """
+    resources = [("generation", policy_generation)]
+    resources.extend(
+        (f"teacher {name}", worker_group)
+        for name, worker_group in (teacher_worker_groups or {}).items()
+    )
+    resources.append(("policy", policy))
+    resources.extend(
+        (f"cluster {index}", virtual_cluster)
+        for index, virtual_cluster in enumerate(cluster or ())
+    )
+
+    seen = set()
+    for label, resource in resources:
+        if resource is None or id(resource) in seen:
+            continue
+        seen.add(id(resource))
+        try:
+            resource.shutdown()
+        except Exception as error:
+            print(f"Error shutting down {label}: {error}", flush=True)
+
+    if ray.is_initialized():
+        ray.shutdown()
+
+
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(description="Run GRPO training with configuration")
@@ -247,11 +286,15 @@ def _make_policy(**kwargs):
                     master_config,
                 )
     finally:
-        shutdown_environments(task_to_env, val_task_to_env)
         try:
-            policy_generation.shutdown()
-        except Exception as error:
-            print(f"Error shutting down generation: {error}", flush=True)
+            shutdown_environments(task_to_env, val_task_to_env)
+        finally:
+            _shutdown_runtime(
+                policy,
+                policy_generation,
+                cluster,
+                teacher_worker_groups,
+            )
 
 
 if __name__ == "__main__":

nemo_rl/algorithms/grpo.py

@@ -1311,7 +1311,10 @@ 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 is not None and refit_transport not in (
+            "nccl_reshard",
+            "mx_nccl_reshard",
+        ):
             # nccl_reshard is handled below via nccl_reshard_refit_enabled,
             # not via checkpoint-engine.
             checkpoint_engine_config = checkpoint_engine_refit_config(generation_config)
@@ -1511,7 +1514,8 @@ def init_dynamo():
     policy.print_node_ip_and_gpu_id()
 
     nccl_reshard_refit_enabled = (
-        generation_config.get("refit_transport") == "nccl_reshard"
+        generation_config.get("refit_transport")
+        in ("nccl_reshard", "mx_nccl_reshard")
     )
     if nccl_reshard_refit_enabled:
         from nemo_rl.weight_sync.nccl_reshard_utils import (

nemo_rl/models/generation/vllm/config.py

@@ -176,13 +176,16 @@ def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None:
     transport = config.get("refit_transport")
     if transport is None:
         return None
-    if transport == "nccl_reshard":
-        # nccl_reshard doesn't takes refit_cfg.
+    if transport in ("nccl_reshard", "mx_nccl_reshard"):
+        # Neither takes refit_cfg. mx_nccl_reshard is the same wire path with
+        # the communicator bootstrap brokered by ModelExpress, so it carries no
+        # transport-scoped config of its own either.
         return 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', 'mx_nccl_reshard', 'vllm_s3_sparse', 'vllm_zmq_sparse', "
+            "'nixl', 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

@@ -284,6 +284,56 @@ def init_nccl_reshard_comm_group(
             group.init_nccl_communicator(device=self.device)
             self.pp_comm_groups[stage] = group
 
+    def init_mx_reshard_comm_group(
+        self,
+        rank_prefix: int,
+        mx_server_url: str,
+        model_name: str,
+        trainer_slots: list,
+        generator_slots: list,
+        source_partition_count: int,
+        plan_digest: str = "",
+        phase: object = "all",
+    ) -> None:
+        """Bootstrap this gen worker's comm groups through ModelExpress.
+
+        Phased for the same reason the trainer side is: one lane at a time,
+        cluster-wide, with the driver barriering between them.
+        """
+        from nemo_rl.weight_sync.mx_collective_bootstrap import (
+            mx_init_lane,
+            mx_rendezvous,
+        )
+
+        if phase == "rendezvous":
+            import uuid
+
+            local_rank = torch.distributed.get_rank()
+            index_in_role = rank_prefix + local_rank
+            if not hasattr(self, "_mx_worker_id"):
+                self._mx_worker_id = (  # pyrefly: ignore[implicitly-defined-attribute]
+                    f"gen-{index_in_role}-{uuid.uuid4().hex}"
+                )
+            self._mx_state = mx_rendezvous(
+                mx_server_url=mx_server_url,
+                model_name=model_name,
+                role="GENERATOR",
+                index_in_role=index_in_role,
+                slot_id=f"gen/{index_in_role}",
+                worker_id=self._mx_worker_id,
+                trainer_slots=trainer_slots,
+                generator_slots=generator_slots,
+                source_partition_count=source_partition_count,
+                plan_digest=plan_digest,
+                device=self.device,
+            )
+            return
+        if phase == "finish":
+            self.pp_comm_groups = self._mx_state.reshard_groups
+            self.model_update_group = self._mx_state.broadcast_group
+            return
+        mx_init_lane(self._mx_state, int(phase))
+
     def report_device_id(self) -> str:
         """Retrieve the UUID of the current CUDA device."""
         from nemo_rl.utils.nvml import get_device_uuid

nemo_rl/models/generation/vllm/vllm_generation.py

@@ -1143,6 +1143,47 @@ def init_nccl_reshard_comm_group(
         # co-works with lm_policy; wait for all futures to complete outside
         return futures
 
+    def init_mx_reshard_comm_group(
+        self,
+        mx_server_url: str,
+        model_name: str,
+        trainer_slots: list,
+        generator_slots: list,
+        source_partition_count: int,
+        plan_digest: str = "",
+        phase: object = "all",
+    ) -> list[ray.ObjectRef]:
+        """Initialize the ModelExpress-brokered comm groups on all gen workers."""
+        if not self.worker_group or not self.worker_group.workers:
+            raise RuntimeError("Worker group is not initialized")
+
+        method_name = (
+            "init_mx_reshard_comm_group_async"
+            if self.cfg["vllm_cfg"]["async_engine"]
+            else "init_mx_reshard_comm_group"
+        )
+
+        total_workers = len(self.worker_group.workers)
+        workers_per_group = total_workers // self.dp_size
+        rank_prefix_list = list(range(0, total_workers, workers_per_group))
+
+        futures = self.worker_group.run_all_workers_multiple_data(
+            method_name,
+            rank_prefix=rank_prefix_list,
+            run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
+            common_kwargs={
+                "mx_server_url": mx_server_url,
+                "model_name": model_name,
+                "trainer_slots": trainer_slots,
+                "generator_slots": generator_slots,
+                "source_partition_count": source_partition_count,
+                "plan_digest": plan_digest,
+                "phase": phase,
+            },
+        )
+        # co-works with lm_policy; wait for all futures to complete outside
+        return futures
+
     def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None:
         """Forward per-layer param metadata to vLLM workers for nccl_reshard refit."""
         method_name = (

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -1231,6 +1231,32 @@ def init_nccl_reshard_comm_group(
             ),
         )
 
+    def init_mx_reshard_comm_group(
+        self,
+        rank_prefix: int,
+        mx_server_url: str,
+        model_name: str,
+        trainer_slots: list,
+        generator_slots: list,
+        source_partition_count: int,
+        plan_digest: str = "",
+        phase: object = "all",
+    ) -> None:
+        """Forward ModelExpress comm-group init to vLLM backend workers."""
+        self.llm.collective_rpc(
+            "init_mx_reshard_comm_group",
+            args=(
+                rank_prefix,
+                mx_server_url,
+                model_name,
+                trainer_slots,
+                generator_slots,
+                source_partition_count,
+                plan_digest,
+                phase,
+            ),
+        )
+
     def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None:
         """Forward refit info to vLLM backend workers."""
         self.llm.collective_rpc("prepare_nccl_reshard_refit_info", args=(refit_info,))

nemo_rl/models/policy/lm_policy.py

@@ -1152,6 +1152,34 @@ def init_nccl_reshard_comm_group(
         # co-works with vllm; wait for all futures to complete outside
         return futures
 
+    def init_mx_reshard_comm_group(
+        self,
+        mx_server_url: str,
+        model_name: str,
+        trainer_slots: list,
+        generator_slots: list,
+        source_partition_count: int,
+        pp_stages: list,
+        plan_digest: str = "",
+        phase: object = "all",
+    ) -> list[ray.ObjectRef]:
+        """Initialize the ModelExpress-brokered comm groups on all train workers."""
+        futures = self.worker_group.run_all_workers_multiple_data(
+            "init_mx_reshard_comm_group",
+            my_pp_stage=pp_stages,
+            common_kwargs={
+                "mx_server_url": mx_server_url,
+                "model_name": model_name,
+                "trainer_slots": trainer_slots,
+                "generator_slots": generator_slots,
+                "source_partition_count": source_partition_count,
+                "plan_digest": plan_digest,
+                "phase": phase,
+            },
+        )
+        # co-works with vllm; wait for all futures to complete outside
+        return futures
+
     def prepare_nccl_reshard_refit_info(
         self,
         train_parallelism,

nemo_rl/models/policy/workers/base_policy_worker.py

@@ -83,6 +83,65 @@ def init_nccl_reshard_comm_group(
         self.pp_comm_group.init_nccl_communicator(device=device)
         self.my_pp_stage = my_pp_stage
 
+    def init_mx_reshard_comm_group(
+        self,
+        mx_server_url: str,
+        model_name: str,
+        trainer_slots: list,
+        generator_slots: list,
+        source_partition_count: int,
+        my_pp_stage: int = 0,
+        plan_digest: str = "",
+        phase: object = "all",
+    ) -> None:
+        """Bootstrap this train worker's comm groups through ModelExpress.
+
+        Same groups the TCPStore path produces, same downstream refit loop; only
+        the ncclUniqueId's provenance differs.
+
+        Driven in phases: ``"rendezvous"``, then one call per lane id, then
+        ``"finish"``. The driver barriers between them because creating two
+        different communicators concurrently across overlapping rank sets
+        deadlocks NCCL -- the same reason the TCPStore path separates its
+        all-ranks group from its per-stage groups.
+        """
+        from nemo_rl.weight_sync.mx_collective_bootstrap import (
+            mx_init_lane,
+            mx_rendezvous,
+        )
+
+        if phase == "rendezvous":
+            import uuid
+
+            if not hasattr(self, "_mx_worker_id"):
+                self._mx_worker_id = (  # pyrefly: ignore[implicitly-defined-attribute]
+                    f"train-{self.rank}-{uuid.uuid4().hex}"
+                )
+            self._mx_state = mx_rendezvous(
+                mx_server_url=mx_server_url,
+                model_name=model_name,
+                role="TRAINER",
+                index_in_role=self.rank,
+                slot_id=f"train/{self.rank}",
+                worker_id=self._mx_worker_id,
+                trainer_slots=trainer_slots,
+                generator_slots=generator_slots,
+                source_partition_count=source_partition_count,
+                source_partition=my_pp_stage,
+                plan_digest=plan_digest,
+                device=torch.cuda.current_device(),
+            )
+            return
+        if phase == "finish":
+            state = self._mx_state
+            # A trainer belongs to exactly one source partition, so it holds one
+            # reshard lane; generators hold all of them.
+            self.pp_comm_group = state.reshard_groups[my_pp_stage]
+            self.model_update_group = state.broadcast_group
+            self.my_pp_stage = my_pp_stage
+            return
+        mx_init_lane(self._mx_state, int(phase))
+
     def prepare_nccl_reshard_refit_info(
         self,
         train_parallelism: dict[str, int],