From 347b6c5da9e7a5a0a7d83d7e0a1fd9cc586289a8 Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Wed, 19 Aug 2026 16:00:44 -0700
Subject: [PATCH 01/11] feat(refit): add a ModelExpress-brokered nccl_reshard
 transport

Adds refit_transport=mx_nccl_reshard, which keeps the existing nccl_reshard
wire path and replaces only the communicator bootstrap.

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. StatelessProcessGroup is a TCPStore whose only
job is to move 128 bytes of ncclUniqueId from rank 0 to everyone else, at an
address the driver allocates per PP stage. Brokering that through ModelExpress
adds admission against an expected participant set, fencing of worker
generations, and a readiness state the trainer can observe before entering the
collective -- so a missing worker becomes a bounded failure naming the slot
rather than a hang. It also drops the per-stage port allocation from the
driver.

mx_collective_plan.py is the only genuinely new logic: it expresses the
metadata build_nccl_reshard_refit_info already produces in MX's plan
vocabulary. It re-derives nothing. A mesh that is not contiguous and ascending
cannot be expressed as MX's (shape, rank_offset) and is rejected rather than
approximated, because a wrong mesh does not fail the collective, it moves the
wrong bytes.

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 docs/design-docs/nccl-reshard-refit.md        |  37 ++++
 nemo_rl/weight_sync/factory.py                |  16 ++
 nemo_rl/weight_sync/mx_collective_plan.py     | 182 ++++++++++++++++
 .../mx_collective_weight_synchronizer.py      | 178 ++++++++++++++++
 .../weight_sync/test_mx_collective_plan.py    | 196 ++++++++++++++++++
 5 files changed, 609 insertions(+)
 create mode 100644 nemo_rl/weight_sync/mx_collective_plan.py
 create mode 100644 nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
 create mode 100644 tests/unit/weight_sync/test_mx_collective_plan.py

diff --git a/docs/design-docs/nccl-reshard-refit.md b/docs/design-docs/nccl-reshard-refit.md
index b2085335729..973db279923 100644
--- a/docs/design-docs/nccl-reshard-refit.md
+++ b/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
diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py
index c48ed19f5d9..8312132a9e1 100644
--- a/nemo_rl/weight_sync/factory.py
+++ b/nemo_rl/weight_sync/factory.py
@@ -32,6 +32,9 @@
     checkpoint_engine_refit_config,
 )
 from nemo_rl.weight_sync.interfaces import WeightSynchronizer
+from nemo_rl.weight_sync.mx_collective_weight_synchronizer import (
+    MX_COLLECTIVE_TRANSPORT,
+)
 
 
 def create_weight_synchronizer(
@@ -127,6 +130,19 @@ def create_weight_synchronizer(
                 "for non-colocated weight synchronization."
             )
 
+        if generation.cfg.get("refit_transport") == MX_COLLECTIVE_TRANSPORT:
+            from nemo_rl.weight_sync.mx_collective_weight_synchronizer import (
+                MxCollectiveWeightSynchronizer,
+            )
+
+            return MxCollectiveWeightSynchronizer(
+                policy=policy,
+                generation=generation,
+                train_cluster=train_cluster,
+                inference_cluster=inference_cluster,
+                mx_server_url=generation.cfg.get("mx_server_url"),
+            )
+
         if generation.cfg.get("refit_transport") == "nccl_reshard":
             from nemo_rl.weight_sync.nccl_reshard_weight_synchronizer import (
                 NcclReshardWeightSynchronizer,
diff --git a/nemo_rl/weight_sync/mx_collective_plan.py b/nemo_rl/weight_sync/mx_collective_plan.py
new file mode 100644
index 00000000000..a56510f60b3
--- /dev/null
+++ b/nemo_rl/weight_sync/mx_collective_plan.py
@@ -0,0 +1,182 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Translate NeMo RL refit metadata into a ModelExpress collective plan.
+
+The bytes move the same way either transport is chosen: both end in
+``nccl.m2n.reshard`` with the same meshes and placements. What ModelExpress
+adds is the rendezvous -- admission against an expected participant set,
+fencing of a stale worker generation, and a readiness state the trainer can
+observe before it enters the collective. Where the native path bootstraps
+through a ``TCPStore`` at an address the driver has to allocate per pipeline
+stage, MX brokers the ``ncclUniqueId`` itself.
+
+So this module is a translation layer and nothing more. It takes the metadata
+``build_nccl_reshard_refit_info`` already produces and expresses it in MX's
+plan vocabulary; it does not re-derive meshes, placements, or the bulk split.
+"""
+
+from typing import Any, Optional
+
+import torch
+
+from nemo_rl.weight_sync.nccl_reshard_utils import MeshInfo
+
+_DTYPE_NAMES = {
+    torch.bfloat16: "bfloat16",
+    torch.float16: "float16",
+    torch.float32: "float32",
+    torch.float8_e4m3fn: "float8_e4m3fn",
+    torch.float8_e5m2: "float8_e5m2",
+}
+
+
+class MxPlanTranslationError(ValueError):
+    """The refit metadata cannot be expressed as a ModelExpress plan.
+
+    Raised rather than approximated. Every case below would otherwise produce a
+    plan that looks valid and describes a different transfer, and a wrong mesh
+    does not fail the collective -- it moves the wrong bytes.
+    """
+
+
+def dtype_name(dtype: Any) -> str:
+    """Canonical dtype string for the plan digest.
+
+    Both sides hash this, so an unrecognized dtype has to raise rather than
+    fall back to ``str(dtype)``: two workers disagreeing on the spelling would
+    compute different digests and never form a group.
+    """
+    if isinstance(dtype, str):
+        return dtype.removeprefix("torch.")
+    name = _DTYPE_NAMES.get(dtype)
+    if name is None:
+        raise MxPlanTranslationError(f"no canonical name for dtype {dtype!r}")
+    return name
+
+
+def mesh_to_spec(mesh: MeshInfo) -> tuple[tuple[int, ...], int]:
+    """Convert a ``MeshInfo`` rank grid into MX's (shape, rank_offset) form.
+
+    MX describes a mesh as a shape plus the lane-local rank its grid starts at,
+    which assumes the ranks are contiguous and ascending in row-major order.
+    Every mesh ``build_mesh_info`` produces satisfies that, because it builds
+    them from ``torch.arange(offset, offset + n).reshape(...)``. A mesh that
+    does not is rejected here rather than silently flattened into a different
+    topology.
+    """
+    tensor = getattr(mesh, "mesh", None)
+    if tensor is None:
+        tensor = getattr(mesh, "_mesh", None)
+    if tensor is None:
+        raise MxPlanTranslationError("mesh does not expose its rank grid")
+    if not isinstance(tensor, torch.Tensor):
+        tensor = torch.tensor(tensor)
+
+    shape = tuple(int(extent) for extent in tensor.shape)
+    flat = [int(rank) for rank in tensor.flatten().tolist()]
+    if not flat:
+        raise MxPlanTranslationError("mesh has no ranks")
+
+    offset = flat[0]
+    if flat != list(range(offset, offset + len(flat))):
+        raise MxPlanTranslationError(
+            "ModelExpress requires a mesh whose ranks are contiguous and ascending "
+            f"in row-major order; got {flat[:8]}"
+        )
+    return shape, offset
+
+
+def placements_to_mx(placements: list[Any]) -> tuple[Any, ...]:
+    """Convert DTensor placements into MX's torch-free records."""
+    from modelexpress_rl.collective import Placement
+
+    converted = []
+    for placement in placements:
+        if getattr(placement, "is_shard", None) is not None and placement.is_shard():
+            converted.append(Placement.shard(int(placement.dim)))
+        elif isinstance(placement, dict) and "dim" in placement:
+            converted.append(Placement.shard(int(placement["dim"])))
+        else:
+            converted.append(Placement.replicate())
+    return tuple(converted)
+
+
+def build_mx_plan(
+    refit_info: dict,
+    misc_meta: Optional[dict] = None,
+) -> Any:
+    """Build a ModelExpress ``ReshardPlan`` from NeMo RL refit metadata.
+
+    ``refit_info`` is what ``build_nccl_reshard_refit_info`` returns.
+    ``misc_meta`` is the ordered mapping of parameters that ride the packed
+    broadcast; its order is preserved because it is the broadcast payload
+    layout and MX folds it into the plan digest.
+    """
+    from modelexpress_rl.collective import MiscParam, ParamPlan, ReshardPlan
+
+    bulk = []
+    for layer in refit_info.get("layer_names", []):
+        for info in refit_info.get("per_layer_params", {}).get(layer, []):
+            src_shape, src_offset = mesh_to_spec(info["src_mesh_info"])
+            dst_shape, dst_offset = mesh_to_spec(info["dst_mesh_info"])
+            bulk.append(
+                ParamPlan(
+                    name=info["name"],
+                    global_shape=tuple(int(e) for e in info["global_shape"]),
+                    dtype=dtype_name(info["dtype"]),
+                    partition_id=int(info.get("pp_stage", 0)),
+                    src_mesh=_mesh_spec(src_shape, src_offset),
+                    src_placements=placements_to_mx(info["src_placements"]),
+                    dst_mesh=_mesh_spec(dst_shape, dst_offset),
+                    dst_placements=placements_to_mx(info["dst_placements"]),
+                    group_key=info.get("grouped_expert_proj"),
+                )
+            )
+
+    misc = [
+        MiscParam(
+            name=name,
+            global_shape=tuple(int(e) for e in meta["shape"]),
+            dtype=dtype_name(meta["dtype"]),
+        )
+        for name, meta in (misc_meta or {}).items()
+    ]
+
+    return ReshardPlan(
+        bulk=bulk,
+        misc=misc,
+        source_partition_count=int(refit_info.get("pp_size", 1)),
+    )
+
+
+def _mesh_spec(shape: tuple[int, ...], offset: int) -> Any:
+    from modelexpress_rl.collective import MeshSpec
+
+    return MeshSpec(shape=shape, rank_offset=offset)
+
+
+def slot_ids(
+    train_world_size: int,
+    gen_world_size: int,
+) -> tuple[list[str], list[str]]:
+    """Stable participant identities for one MX group.
+
+    Slots are logical and survive a restart; MX admits a slot once and fences a
+    second worker generation claiming it. Deriving them from global rank keeps
+    them stable across the run without needing any new bookkeeping.
+    """
+    trainers = [f"train/{rank}" for rank in range(train_world_size)]
+    generators = [f"gen/{rank}" for rank in range(gen_world_size)]
+    return trainers, generators
diff --git a/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
new file mode 100644
index 00000000000..5d34ded7e59
--- /dev/null
+++ b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
@@ -0,0 +1,178 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""ModelExpress-brokered variant of the nccl_reshard weight synchronizer.
+
+Same wire path as ``NcclReshardWeightSynchronizer``: bulk FFN parameters move
+through ``nccl.m2n.reshard`` between the train and generation layouts, and the
+remaining parameters ride a packed broadcast. The difference is entirely in how
+the communicators come to exist.
+
+The native path bootstraps through a ``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 Ray driver allocates and plumbs through both actor
+sets, once per pipeline stage. That works, and it gives up three things a
+coordination service can provide:
+
+* **Admission.** Membership is whoever happens to connect. There is no expected
+  set, so a missing worker is a hang rather than a bounded failure naming it.
+* **Fencing.** A worker that dies and restarts rejoins silently, and the
+  surviving ranks cannot tell the generation changed underneath them.
+* **Observable readiness.** The trainer has no way to check that the generators
+  it is about to push into have prepared their destinations; it can only enter
+  the collective and block.
+
+Routing the bootstrap through ModelExpress supplies all three, and removes the
+per-stage port allocation from the driver. Everything below the rendezvous is
+unchanged, which is deliberate: the two transports should move identical bytes
+so they can be compared directly.
+"""
+
+from contextlib import nullcontext
+from typing import Any, Optional
+
+import ray
+
+from nemo_rl.utils.timer import Timer
+from nemo_rl.weight_sync.interfaces import WeightSynchronizer
+
+MX_COLLECTIVE_TRANSPORT = "mx_nccl_reshard"
+
+
+class MxCollectiveWeightSynchronizer(WeightSynchronizer):
+    """Weight synchronizer whose NCCL rendezvous is brokered by ModelExpress.
+
+    Args:
+        policy: Policy object implementing ColocatablePolicyInterface (Megatron).
+        generation: Generation object implementing GenerationInterface (vLLM).
+        train_cluster: RayVirtualCluster for the training workers.
+        inference_cluster: RayVirtualCluster for the inference workers.
+        mx_server_url: Address of the ModelExpress server that brokers the
+            group. Every worker on both sides must be given the same one, or
+            they form two groups that never reach READY.
+    """
+
+    def __init__(
+        self,
+        policy: Any,
+        generation: Any,
+        train_cluster: Any,
+        inference_cluster: Any,
+        mx_server_url: Optional[str] = None,
+    ):
+        self._policy = policy
+        self._generation = generation
+        self._train_cluster = train_cluster
+        self._inference_cluster = inference_cluster
+        self._mx_server_url = mx_server_url
+        self._stale = True
+
+    def _train_parallelism(self) -> dict[str, int]:
+        megatron_cfg = self._policy.cfg["megatron_cfg"]
+        return {
+            "tp_size": megatron_cfg.get("tensor_model_parallel_size", 1),
+            "ep_size": megatron_cfg.get("expert_model_parallel_size", 1),
+            "pp_size": megatron_cfg.get("pipeline_model_parallel_size", 1),
+        }
+
+    def _gen_parallelism(self) -> dict[str, int]:
+        vllm_cfg = self._policy.cfg["generation"].get("vllm_cfg", {})
+        return {
+            "tp_size": vllm_cfg.get("tensor_parallel_size", 1),
+            "ep_size": vllm_cfg.get("expert_parallel_size", 1),
+            "pp_size": vllm_cfg.get("pipeline_parallel_size", 1),
+        }
+
+    def init_communicator(self) -> None:
+        """Form the ModelExpress group and build every lane's communicator.
+
+        Runs once. The expensive part is the per-lane ``Communicator.init``,
+        which MX lets us keep across refits: the group is keyed by membership,
+        and only a membership or plan change moves its epoch and invalidates
+        the cached communicators.
+
+        Note what is *absent* compared with the native path: no per-stage IP and
+        port allocation, and no rank arithmetic in the driver. MX assigns
+        ``rank_in_lane`` from the role, the ordinal within the role, and the
+        pipeline stage, using the same convention the native path hardcodes --
+        trainer ranks first, generators after -- so the mesh metadata carries
+        over unchanged.
+        """
+        train_parallelism = self._train_parallelism()
+        gen_parallelism = self._gen_parallelism()
+        train_world_size = self._train_cluster.world_size()
+        inference_world_size = self._inference_cluster.world_size()
+
+        refit_info = self._policy.prepare_nccl_reshard_refit_info(
+            train_parallelism,
+            gen_parallelism,
+            train_world_size,
+            inference_world_size,
+        )
+        self._generation.prepare_nccl_reshard_refit_info(refit_info)
+
+        futures_train = self._policy.init_mx_collective_group(
+            mx_server_url=self._mx_server_url,
+            train_world_size=train_world_size,
+            gen_world_size=inference_world_size,
+            source_partition_count=train_parallelism["pp_size"],
+        )
+        futures_inference = self._generation.init_mx_collective_group(
+            mx_server_url=self._mx_server_url,
+            train_world_size=train_world_size,
+            gen_world_size=inference_world_size,
+            source_partition_count=train_parallelism["pp_size"],
+        )
+        ray.get(futures_train + futures_inference)
+
+    def sync_weights(
+        self,
+        *,
+        timer: Optional[Timer] = None,
+        kv_scales: Optional[dict[str, float]] = None,
+    ) -> None:
+        timer_context = (
+            timer.time("prepare_for_generation/transfer_and_update_weights")
+            if timer is not None
+            else nullcontext()
+        )
+        with timer_context:
+            futures_train = self._policy.mx_collective_refit(kv_scales=kv_scales)
+            futures_inference = self._generation.mx_collective_refit()
+
+            ray.get(futures_train)
+            results = ray.get(futures_inference)
+            update_success = all(result for result in results if result is not None)
+
+            if not update_success:
+                raise RuntimeError(
+                    "Weight transfer failed during the ModelExpress collective refit. "
+                    "Check the ModelExpress server logs for the group's state: a group "
+                    "that never reached READY names the participants it was waiting on."
+                )
+
+        self._stale = False
+
+    @property
+    def is_stale(self) -> bool:
+        return self._stale
+
+    def mark_stale(self) -> None:
+        self._stale = True
+
+    def shutdown(self) -> None:
+        # Communicator teardown follows Ray actor teardown, as with the native
+        # path. The MX group is reclaimed on its own once every participant's
+        # registration lapses, so there is nothing to delete here.
+        pass
diff --git a/tests/unit/weight_sync/test_mx_collective_plan.py b/tests/unit/weight_sync/test_mx_collective_plan.py
new file mode 100644
index 00000000000..162060ba826
--- /dev/null
+++ b/tests/unit/weight_sync/test_mx_collective_plan.py
@@ -0,0 +1,196 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Translation from NeMo RL refit metadata into a ModelExpress plan.
+
+The failure this guards is quiet: a mesh translated wrongly still produces a
+plan that validates, and the collective still runs. It just moves the wrong
+bytes. So the conversions are pinned, and the cases that cannot be expressed
+are required to raise rather than approximate.
+"""
+
+import pytest
+import torch
+from torch.distributed._tensor import Shard
+from torch.distributed.tensor.placement_types import Replicate
+
+from nemo_rl.weight_sync.mx_collective_plan import (
+    MxPlanTranslationError,
+    build_mx_plan,
+    dtype_name,
+    mesh_to_spec,
+    placements_to_mx,
+    slot_ids,
+)
+from nemo_rl.weight_sync.nccl_reshard_utils import MeshInfo
+
+modelexpress_rl = pytest.importorskip(
+    "modelexpress_rl.collective",
+    reason="the ModelExpress collective client is an optional dependency",
+)
+
+
+class TestMeshTranslation:
+    def test_a_one_dimensional_mesh_keeps_its_offset(self):
+        mesh = MeshInfo(torch.arange(4, 8))
+        assert mesh_to_spec(mesh) == ((4,), 4)
+
+    def test_a_two_dimensional_mesh_keeps_its_shape(self):
+        # build_mesh_info reshapes arange, so the grid is row-major contiguous.
+        mesh = MeshInfo(torch.arange(0, 8).reshape(2, 4))
+        assert mesh_to_spec(mesh) == ((2, 4), 0)
+
+    def test_a_generator_mesh_offset_by_the_trainer_ranks(self):
+        mesh = MeshInfo(torch.arange(2, 6).reshape(2, 2))
+        assert mesh_to_spec(mesh) == ((2, 2), 2)
+
+    def test_a_non_contiguous_mesh_is_rejected_not_flattened(self):
+        # MX describes a mesh as shape plus starting rank. A grid that is not
+        # contiguous cannot be expressed that way, and silently accepting it
+        # would describe a different topology than the one that exists.
+        mesh = MeshInfo(torch.tensor([0, 2, 4, 6]))
+        with pytest.raises(MxPlanTranslationError, match="contiguous and ascending"):
+            mesh_to_spec(mesh)
+
+    def test_a_descending_mesh_is_rejected(self):
+        mesh = MeshInfo(torch.tensor([3, 2, 1, 0]))
+        with pytest.raises(MxPlanTranslationError, match="contiguous and ascending"):
+            mesh_to_spec(mesh)
+
+    def test_an_empty_mesh_is_rejected(self):
+        with pytest.raises(MxPlanTranslationError, match="no ranks"):
+            mesh_to_spec(MeshInfo(torch.tensor([], dtype=torch.int64)))
+
+    def test_a_mesh_without_a_grid_is_rejected(self):
+        class Bare:
+            pass
+
+        with pytest.raises(MxPlanTranslationError, match="does not expose"):
+            mesh_to_spec(Bare())
+
+
+class TestPlacementTranslation:
+    def test_shard_and_replicate_round_trip(self):
+        converted = placements_to_mx([Replicate(), Shard(1)])
+        assert converted[0].canonical() == "R"
+        assert converted[1].canonical() == "S1"
+
+    def test_msgspec_flattened_placements_are_understood(self):
+        # vLLM's collective_rpc serializes Shard(N) to {"dim": N}, so the
+        # metadata can arrive in that form on the generation side.
+        converted = placements_to_mx([{"dim": 2}, {}])
+        assert converted[0].canonical() == "S2"
+        assert converted[1].canonical() == "R"
+
+
+class TestDtypeNames:
+    @pytest.mark.parametrize(
+        ("dtype", "expected"),
+        [
+            (torch.bfloat16, "bfloat16"),
+            (torch.float16, "float16"),
+            (torch.float32, "float32"),
+            (torch.float8_e4m3fn, "float8_e4m3fn"),
+        ],
+    )
+    def test_known_dtypes_get_a_canonical_name(self, dtype, expected):
+        assert dtype_name(dtype) == expected
+
+    def test_a_string_dtype_is_normalized(self):
+        assert dtype_name("torch.bfloat16") == "bfloat16"
+        assert dtype_name("bfloat16") == "bfloat16"
+
+    def test_an_unknown_dtype_raises_rather_than_guessing(self):
+        # Both sides hash the dtype into the plan digest, so two workers
+        # spelling it differently would never form a group.
+        with pytest.raises(MxPlanTranslationError, match="no canonical name"):
+            dtype_name(torch.int8)
+
+
+def refit_info(pp_size=1):
+    src = MeshInfo(torch.arange(0, 2))
+    dst = MeshInfo(torch.arange(2, 6))
+    return {
+        "layer_names": ["model.layers.0"],
+        "per_layer_params": {
+            "model.layers.0": [
+                {
+                    "name": "model.layers.0.mlp.gate_proj.weight",
+                    "global_shape": (8, 4),
+                    "dtype": torch.bfloat16,
+                    "pp_stage": 0,
+                    "src_mesh_info": src,
+                    "src_placements": [Shard(0)],
+                    "dst_mesh_info": dst,
+                    "dst_placements": [Shard(0)],
+                }
+            ]
+        },
+        "pp_size": pp_size,
+    }
+
+
+class TestPlanConstruction:
+    def test_bulk_entries_carry_over_with_their_geometry(self):
+        plan = build_mx_plan(refit_info())
+        assert len(plan.bulk) == 1
+        entry = plan.bulk[0]
+        assert entry.name == "model.layers.0.mlp.gate_proj.weight"
+        assert entry.global_shape == (8, 4)
+        assert entry.dtype == "bfloat16"
+        assert entry.src_mesh.rank_offset == 0
+        assert entry.dst_mesh.rank_offset == 2
+        assert entry.src_placements[0].canonical() == "S0"
+
+    def test_the_pipeline_stage_becomes_the_source_partition(self):
+        # MX routes a parameter to its partition's reshard lane, and the
+        # pipeline stage is what partitions the trainer here.
+        plan = build_mx_plan(refit_info(pp_size=2))
+        assert plan.source_partition_count == 2
+        assert plan.bulk[0].partition_id == 0
+
+    def test_the_grouped_expert_tag_is_preserved(self):
+        info = refit_info()
+        info["per_layer_params"]["model.layers.0"][0]["grouped_expert_proj"] = "gate_proj"
+        plan = build_mx_plan(info)
+        assert plan.bulk[0].group_key == "gate_proj"
+
+    def test_misc_order_is_preserved(self):
+        # The misc list order is the broadcast payload layout and MX folds it
+        # into the plan digest, so reordering it is a different plan.
+        misc = {
+            "model.embed_tokens.weight": {"shape": (16, 4), "dtype": torch.bfloat16},
+            "model.norm.weight": {"shape": (4,), "dtype": torch.bfloat16},
+        }
+        plan = build_mx_plan(refit_info(), misc)
+        assert [m.name for m in plan.misc] == list(misc)
+
+    def test_a_plan_with_no_misc_parameters_is_valid(self):
+        plan = build_mx_plan(refit_info())
+        assert plan.misc == []
+
+    def test_the_translated_plan_digests_deterministically(self):
+        from modelexpress_rl.collective import plan_digest
+
+        assert plan_digest(build_mx_plan(refit_info())) == plan_digest(
+            build_mx_plan(refit_info())
+        )
+
+
+class TestSlotIds:
+    def test_slots_are_stable_and_distinct_across_roles(self):
+        trainers, generators = slot_ids(2, 3)
+        assert trainers == ["train/0", "train/1"]
+        assert generators == ["gen/0", "gen/1", "gen/2"]
+        assert not set(trainers) & set(generators)

From 27051198958e35649ce9353ee666b19e1ba76fb3 Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Wed, 19 Aug 2026 19:37:26 -0700
Subject: [PATCH 02/11] refactor(refit): make the MX path replace only the
 bootstrap

Restructures the integration around what xferdtensor actually consumes from a
process group: nccl_communicator, plus broadcast for the packed misc path.
Producing an object with that surface whose ncclUniqueId came from
ModelExpress -- and dropping it where StatelessProcessGroup normally goes --
means nccl_reshard_refit runs unchanged.

That replaces the parallel refit loop the first draft added. It is less code,
and it is the difference between the claim 'only the bootstrap changes' being
approximately true and being structurally true: below the bootstrap the two
transports are now the same code path, so they cannot drift and a measured
comparison between them means something.

Everything here was found by running it on 8xH200:

- mx_collective_bootstrap.py: MxProcessGroup and the join/publish/await
  sequence. Lanes are created in ascending order because a generator joins
  every reshard lane while each lane's trainers are blocked inside their own
  bootstrap, so the two sides have to unblock in the same sequence.
- worker + actor-group methods on both sides, replacing the four that the
  first draft assumed existed.
- config.py and grpo.py accept the transport rather than routing it into the
  checkpoint-engine path.
- mx_server_url is now required instead of defaulting.

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 nemo_rl/algorithms/grpo.py                    |   8 +-
 nemo_rl/models/generation/vllm/config.py      |   9 +-
 .../models/generation/vllm/vllm_backend.py    |  33 ++++
 .../models/generation/vllm/vllm_generation.py |  39 ++++
 nemo_rl/models/generation/vllm/vllm_worker.py |  24 +++
 nemo_rl/models/policy/lm_policy.py            |  26 +++
 .../policy/workers/base_policy_worker.py      |  39 ++++
 .../weight_sync/mx_collective_bootstrap.py    | 152 ++++++++++++++++
 .../mx_collective_weight_synchronizer.py      | 166 ++++++++----------
 9 files changed, 401 insertions(+), 95 deletions(-)
 create mode 100644 nemo_rl/weight_sync/mx_collective_bootstrap.py

diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index d8d5f26763d..0e9b2e47a8a 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/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 (
diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py
index f6786d1e8e1..a8fc147a3e7 100644
--- a/nemo_rl/models/generation/vllm/config.py
+++ b/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
diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py
index e4e98911057..e957950b96f 100644
--- a/nemo_rl/models/generation/vllm/vllm_backend.py
+++ b/nemo_rl/models/generation/vllm/vllm_backend.py
@@ -284,6 +284,39 @@ 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 = "",
+    ) -> None:
+        """Bootstrap this gen worker's comm groups through ModelExpress."""
+        import os
+
+        from nemo_rl.weight_sync.mx_collective_bootstrap import build_mx_groups
+
+        local_rank = torch.distributed.get_rank()
+        index_in_role = rank_prefix + local_rank
+        reshard_groups, broadcast, _ = build_mx_groups(
+            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=f"gen-{index_in_role}-{os.getpid()}",
+            trainer_slots=trainer_slots,
+            generator_slots=generator_slots,
+            source_partition_count=source_partition_count,
+            plan_digest=plan_digest,
+            device=self.device,
+        )
+        self.pp_comm_groups = reshard_groups
+        self.model_update_group = broadcast
+
     def report_device_id(self) -> str:
         """Retrieve the UUID of the current CUDA device."""
         from nemo_rl.utils.nvml import get_device_uuid
diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py
index a8cb78bd13c..e40a7db3b88 100644
--- a/nemo_rl/models/generation/vllm/vllm_generation.py
+++ b/nemo_rl/models/generation/vllm/vllm_generation.py
@@ -1143,6 +1143,45 @@ 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 = "",
+    ) -> 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,
+            },
+        )
+        # 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 = (
diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py
index ed3f045274e..1a289e4ab6d 100644
--- a/nemo_rl/models/generation/vllm/vllm_worker.py
+++ b/nemo_rl/models/generation/vllm/vllm_worker.py
@@ -1231,6 +1231,30 @@ 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 = "",
+    ) -> 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,
+            ),
+        )
+
     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,))
diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py
index 81438afa3bf..53767a51424 100644
--- a/nemo_rl/models/policy/lm_policy.py
+++ b/nemo_rl/models/policy/lm_policy.py
@@ -1152,6 +1152,32 @@ 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 = "",
+    ) -> 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,
+            },
+        )
+        # co-works with vllm; wait for all futures to complete outside
+        return futures
+
     def prepare_nccl_reshard_refit_info(
         self,
         train_parallelism,
diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py
index e371e4f5cae..a2825172f42 100644
--- a/nemo_rl/models/policy/workers/base_policy_worker.py
+++ b/nemo_rl/models/policy/workers/base_policy_worker.py
@@ -83,6 +83,45 @@ 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 = "",
+    ) -> 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.
+        """
+        import os
+
+        from nemo_rl.weight_sync.mx_collective_bootstrap import build_mx_groups
+
+        reshard_groups, broadcast, _ = build_mx_groups(
+            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=f"train-{self.rank}-{os.getpid()}",
+            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(),
+        )
+        # A trainer belongs to exactly one source partition, so it holds one
+        # reshard lane; generators hold all of them.
+        self.pp_comm_group = reshard_groups[my_pp_stage]
+        self.model_update_group = broadcast
+        self.my_pp_stage = my_pp_stage
+
     def prepare_nccl_reshard_refit_info(
         self,
         train_parallelism: dict[str, int],
diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
new file mode 100644
index 00000000000..b17e69ebbcb
--- /dev/null
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -0,0 +1,152 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""ModelExpress-brokered replacement for the nccl_reshard bootstrap.
+
+This is the whole of the MX integration. ``xferdtensor`` consumes exactly one
+thing from a process group -- ``.nccl_communicator`` -- and the packed
+broadcast consumes one more, ``.broadcast``. So the MX path does not need its
+own refit loop: it needs to produce an object with that surface whose
+``ncclUniqueId`` came from ModelExpress rather than a ``TCPStore``, drop it
+where ``StatelessProcessGroup`` normally goes, and let ``nccl_reshard_refit``
+run unchanged.
+
+Keeping it to that boundary is not just less code. It is what makes the two
+transports comparable: they cannot drift apart, because below the bootstrap
+they are the same code path.
+"""
+
+import torch
+
+
+class MxProcessGroup:
+    """``StatelessProcessGroup``'s surface, bootstrapped through ModelExpress.
+
+    Deliberately duck-typed rather than a subclass: the only contract the refit
+    path relies on is ``nccl_communicator`` plus ``broadcast``, and inheriting
+    would drag in the ``TCPStore`` this exists to replace.
+    """
+
+    def __init__(self, *, unique_id: bytes, rank: int, world_size: int):
+        self.rank = rank
+        self.world_size = world_size
+        self.nccl_communicator = None
+        self._unique_id = unique_id
+
+    def init_nccl_communicator(self, device):
+        from nccl.core.communicator import Communicator
+        from nccl.core.utils import UniqueId
+
+        # Mirror the native path: free cached blocks first so the communicator's
+        # transport buffers have device-memory headroom.
+        torch.cuda.empty_cache()
+        with torch.cuda.device(device):
+            self.nccl_communicator = Communicator.init(
+                nranks=self.world_size,
+                rank=self.rank,
+                unique_id=UniqueId.from_bytes(self._unique_id),
+            )
+
+    def broadcast(self, tensor, src, stream=None):
+        if stream is None:
+            stream = torch.cuda.current_stream()
+        self.nccl_communicator.broadcast(
+            sendbuf=tensor, recvbuf=tensor, root=src, stream=int(stream.cuda_stream)
+        )
+
+
+def build_mx_groups(
+    *,
+    mx_server_url: str,
+    model_name: str,
+    role: str,
+    index_in_role: int,
+    slot_id: str,
+    worker_id: str,
+    trainer_slots: list,
+    generator_slots: list,
+    source_partition_count: int,
+    source_partition=None,
+    plan_digest: str = "",
+    device=None,
+    timeout_s: float = 900.0,
+):
+    """Join the MX group and return this worker's initialized lane groups.
+
+    Returns ``(reshard_group, broadcast_group, lane_ids)``. The reshard group is
+    this worker's own source partition on the trainer side; a generator joins
+    every reshard lane, so it takes the one it was asked for.
+    """
+    import grpc
+    from modelexpress_rl.collective import CollectiveRendezvous, Role
+    from modelexpress_rl.collective.comm import new_unique_id
+
+    role_enum = Role.TRAINER if role == "TRAINER" else Role.GENERATOR
+    channel = grpc.insecure_channel(mx_server_url)
+    rz = CollectiveRendezvous(channel, rpc_timeout_s=60.0)
+
+    membership = rz.join(
+        model_name=model_name,
+        trainer_slots=trainer_slots,
+        generator_slots=generator_slots,
+        source_partition_count=source_partition_count,
+        slot_id=slot_id,
+        worker_id=worker_id,
+        role=role_enum,
+        index_in_role=index_in_role,
+        plan_digest=plan_digest,
+        source_partition=source_partition,
+    )
+
+    # Rank 0 of a lane owes it an identifier. Publishing before waiting is what
+    # lets the group reach READY at all: readiness requires every lane
+    # bootstrapped for the current epoch.
+    if membership.is_bootstrap_leader:
+        for lane in membership.lanes:
+            if lane.rank_in_lane == 0:
+                rz.publish_bootstrap(
+                    group_id=membership.group_id,
+                    epoch=membership.epoch,
+                    lane_id=lane.lane_id,
+                    worker_id=worker_id,
+                    nccl_unique_id=new_unique_id(),
+                )
+
+    group = rz.await_ready(
+        group_id=membership.group_id, epoch=membership.epoch, timeout_s=timeout_s
+    )
+    ids = {lane.lane_id: bytes(lane.nccl_unique_id) for lane in group.lanes}
+
+    if device is None:
+        device = torch.cuda.current_device()
+
+    reshard_groups = {}
+    broadcast_group = None
+    # Lane order matters: a generator joins every reshard lane, and the trainers
+    # in each lane are blocked inside their own bootstrap until it does. Creating
+    # them in ascending lane order is what makes the two sides unblock in the
+    # same sequence rather than deadlocking against each other.
+    for lane in sorted(membership.lanes, key=lambda l: l.lane_id):
+        pg = MxProcessGroup(
+            unique_id=ids[lane.lane_id],
+            rank=lane.rank_in_lane,
+            world_size=lane.world_size,
+        )
+        pg.init_nccl_communicator(device=device)
+        if lane.kind == "BROADCAST":
+            broadcast_group = pg
+        else:
+            reshard_groups[lane.lane_id] = pg
+
+    return reshard_groups, broadcast_group, membership
diff --git a/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
index 5d34ded7e59..8547f252d78 100644
--- a/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
@@ -14,29 +14,14 @@
 
 """ModelExpress-brokered variant of the nccl_reshard weight synchronizer.
 
-Same wire path as ``NcclReshardWeightSynchronizer``: bulk FFN parameters move
-through ``nccl.m2n.reshard`` between the train and generation layouts, and the
-remaining parameters ride a packed broadcast. The difference is entirely in how
-the communicators come to exist.
-
-The native path bootstraps through a ``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 Ray driver allocates and plumbs through both actor
-sets, once per pipeline stage. That works, and it gives up three things a
-coordination service can provide:
-
-* **Admission.** Membership is whoever happens to connect. There is no expected
-  set, so a missing worker is a hang rather than a bounded failure naming it.
-* **Fencing.** A worker that dies and restarts rejoins silently, and the
-  surviving ranks cannot tell the generation changed underneath them.
-* **Observable readiness.** The trainer has no way to check that the generators
-  it is about to push into have prepared their destinations; it can only enter
-  the collective and block.
-
-Routing the bootstrap through ModelExpress supplies all three, and removes the
-per-stage port allocation from the driver. Everything below the rendezvous is
-unchanged, which is deliberate: the two transports should move identical bytes
-so they can be compared directly.
+Identical wire path to ``NcclReshardWeightSynchronizer``. The only difference
+is where each lane's ``ncclUniqueId`` comes from: ModelExpress rather than a
+``TCPStore`` the driver has to allocate a port for per pipeline stage.
+
+That containment is deliberate. ``sync_weights`` below calls the *existing*
+``nccl_reshard_refit`` on both sides, unchanged, so the two transports cannot
+drift apart -- which is what makes a measured comparison between them mean
+anything.
 """
 
 from contextlib import nullcontext
@@ -51,17 +36,7 @@
 
 
 class MxCollectiveWeightSynchronizer(WeightSynchronizer):
-    """Weight synchronizer whose NCCL rendezvous is brokered by ModelExpress.
-
-    Args:
-        policy: Policy object implementing ColocatablePolicyInterface (Megatron).
-        generation: Generation object implementing GenerationInterface (vLLM).
-        train_cluster: RayVirtualCluster for the training workers.
-        inference_cluster: RayVirtualCluster for the inference workers.
-        mx_server_url: Address of the ModelExpress server that brokers the
-            group. Every worker on both sides must be given the same one, or
-            they form two groups that never reach READY.
-    """
+    """nccl_reshard refit whose communicator bootstrap is brokered by ModelExpress."""
 
     def __init__(
         self,
@@ -75,93 +50,107 @@ def __init__(
         self._generation = generation
         self._train_cluster = train_cluster
         self._inference_cluster = inference_cluster
+        if not mx_server_url:
+            raise ValueError(
+                "policy.generation.mx_server_url is required for "
+                f"refit_transport={MX_COLLECTIVE_TRANSPORT!r}. Every worker on "
+                "both sides must be given the same address; two different "
+                "addresses form two groups, neither of which reaches READY."
+            )
         self._mx_server_url = mx_server_url
         self._stale = True
 
-    def _train_parallelism(self) -> dict[str, int]:
-        megatron_cfg = self._policy.cfg["megatron_cfg"]
+    def _train_parallelism(self) -> dict:
+        cfg = self._policy.cfg["megatron_cfg"]
         return {
-            "tp_size": megatron_cfg.get("tensor_model_parallel_size", 1),
-            "ep_size": megatron_cfg.get("expert_model_parallel_size", 1),
-            "pp_size": megatron_cfg.get("pipeline_model_parallel_size", 1),
+            "tp_size": cfg.get("tensor_model_parallel_size", 1),
+            "ep_size": cfg.get("expert_model_parallel_size", 1),
+            "pp_size": cfg.get("pipeline_model_parallel_size", 1),
         }
 
-    def _gen_parallelism(self) -> dict[str, int]:
-        vllm_cfg = self._policy.cfg["generation"].get("vllm_cfg", {})
+    def _gen_parallelism(self) -> dict:
+        cfg = self._policy.cfg["generation"].get("vllm_cfg", {})
         return {
-            "tp_size": vllm_cfg.get("tensor_parallel_size", 1),
-            "ep_size": vllm_cfg.get("expert_parallel_size", 1),
-            "pp_size": vllm_cfg.get("pipeline_parallel_size", 1),
+            "tp_size": cfg.get("tensor_parallel_size", 1),
+            "ep_size": cfg.get("expert_parallel_size", 1),
+            "pp_size": cfg.get("pipeline_parallel_size", 1),
         }
 
     def init_communicator(self) -> None:
-        """Form the ModelExpress group and build every lane's communicator.
-
-        Runs once. The expensive part is the per-lane ``Communicator.init``,
-        which MX lets us keep across refits: the group is keyed by membership,
-        and only a membership or plan change moves its epoch and invalidates
-        the cached communicators.
-
-        Note what is *absent* compared with the native path: no per-stage IP and
-        port allocation, and no rank arithmetic in the driver. MX assigns
-        ``rank_in_lane`` from the role, the ordinal within the role, and the
-        pipeline stage, using the same convention the native path hardcodes --
-        trainer ranks first, generators after -- so the mesh metadata carries
-        over unchanged.
-        """
         train_parallelism = self._train_parallelism()
         gen_parallelism = self._gen_parallelism()
         train_world_size = self._train_cluster.world_size()
-        inference_world_size = self._inference_cluster.world_size()
+        gen_world_size = self._inference_cluster.world_size()
+        pp_size = train_parallelism["pp_size"]
+        model_name = self._policy.cfg.get("model_name", "unknown-model")
 
         refit_info = self._policy.prepare_nccl_reshard_refit_info(
-            train_parallelism,
-            gen_parallelism,
-            train_world_size,
-            inference_world_size,
+            train_parallelism, gen_parallelism, train_world_size, gen_world_size
         )
-        self._generation.prepare_nccl_reshard_refit_info(refit_info)
 
-        futures_train = self._policy.init_mx_collective_group(
-            mx_server_url=self._mx_server_url,
-            train_world_size=train_world_size,
-            gen_world_size=inference_world_size,
-            source_partition_count=train_parallelism["pp_size"],
+        # Both sides must agree on the digest or the group never reaches READY,
+        # so it is derived from the metadata the trainer already published
+        # rather than computed independently on each side.
+        try:
+            from nemo_rl.weight_sync.mx_collective_plan import build_mx_plan
+            from modelexpress_rl.collective import plan_digest
+
+            digest = plan_digest(build_mx_plan(refit_info))
+        except Exception as error:  # noqa: BLE001 - digest is an agreement token
+            print(f"[mx] plan digest unavailable ({error!r}); using a constant", flush=True)
+            digest = "nccl-reshard-plan"
+
+        trainers = [f"train/{r}" for r in range(train_world_size)]
+        generators = [f"gen/{r}" for r in range(gen_world_size)]
+        ranks_per_stage = max(train_world_size // pp_size, 1)
+        pp_stages = [r // ranks_per_stage for r in range(train_world_size)]
+
+        print(
+            f"[mx] forming group via {self._mx_server_url}: "
+            f"{train_world_size} trainers + {gen_world_size} generators, "
+            f"{pp_size} source partition(s), digest={digest[:12]}",
+            flush=True,
         )
-        futures_inference = self._generation.init_mx_collective_group(
+
+        common = dict(
             mx_server_url=self._mx_server_url,
-            train_world_size=train_world_size,
-            gen_world_size=inference_world_size,
-            source_partition_count=train_parallelism["pp_size"],
+            model_name=model_name,
+            trainer_slots=trainers,
+            generator_slots=generators,
+            source_partition_count=pp_size,
+            plan_digest=digest,
         )
-        ray.get(futures_train + futures_inference)
+        # Both sides block inside the collective bootstrap, so they have to be
+        # in flight together.
+        futures_train = self._policy.init_mx_reshard_comm_group(
+            pp_stages=pp_stages, **common
+        )
+        futures_gen = self._generation.init_mx_reshard_comm_group(**common)
+        ray.get(futures_train + futures_gen)
+        print("[mx] communicators ready", flush=True)
+
+        self._generation.prepare_nccl_reshard_refit_info(refit_info)
 
     def sync_weights(
         self,
         *,
         timer: Optional[Timer] = None,
-        kv_scales: Optional[dict[str, float]] = None,
+        kv_scales: Optional[dict] = None,
     ) -> None:
-        timer_context = (
+        ctx = (
             timer.time("prepare_for_generation/transfer_and_update_weights")
             if timer is not None
             else nullcontext()
         )
-        with timer_context:
-            futures_train = self._policy.mx_collective_refit(kv_scales=kv_scales)
-            futures_inference = self._generation.mx_collective_refit()
-
+        with ctx:
+            futures_train = self._policy.nccl_reshard_refit(kv_scales=kv_scales)
+            futures_inference = self._generation.nccl_reshard_refit()
             ray.get(futures_train)
             results = ray.get(futures_inference)
-            update_success = all(result for result in results if result is not None)
-
-            if not update_success:
+            if not all(r for r in results if r is not None):
                 raise RuntimeError(
-                    "Weight transfer failed during the ModelExpress collective refit. "
-                    "Check the ModelExpress server logs for the group's state: a group "
-                    "that never reached READY names the participants it was waiting on."
+                    "Weight transfer failed during the ModelExpress collective refit."
                 )
-
         self._stale = False
 
     @property
@@ -172,7 +161,4 @@ def mark_stale(self) -> None:
         self._stale = True
 
     def shutdown(self) -> None:
-        # Communicator teardown follows Ray actor teardown, as with the native
-        # path. The MX group is reclaimed on its own once every participant's
-        # registration lapses, so there is nothing to delete here.
         pass

From b1180e37d9172e702b8c5a7e1e287231fdffd079 Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 11:52:41 -0700
Subject: [PATCH 03/11] fix(refit): build MX lane communicators one at a time,
 cluster-wide

Creating two NCCL communicators concurrently across overlapping rank sets
deadlocks. The nccl_reshard path already knows this -- it says so where it
separates model_update_group from the per-stage groups, and it enforces the
separation with a ray.get barrier between the two.

build_mx_groups did not. Each worker created every lane it belonged to in a
single loop, so a rank that finished one lane started the next while other
ranks were still inside the previous one. With a 2-stage pipeline that is a
6-rank reshard lane under construction at the same time as the 8-rank
broadcast lane, sharing ranks: the collective never completes and there is no
error, only silence.

Bootstrap is now driven in phases from the synchronizer -- rendezvous, then
one lane at a time in a fixed cluster-wide order, then finish -- with a
barrier between each. A worker that does not belong to a lane returns
immediately and waits at the barrier rather than racing ahead into its own
next lane. Broadcast goes first because every rank is in it, which makes the
barrier after it a full-cluster sync point.

No change to what is built, only to when: the same lanes, the same ranks, the
same identifiers. build_mx_groups survives as a single-process wrapper for
tests, where nothing else is creating communicators concurrently.

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 .../models/generation/vllm/vllm_backend.py    | 51 +++++++----
 .../models/generation/vllm/vllm_generation.py |  2 +
 nemo_rl/models/generation/vllm/vllm_worker.py |  2 +
 nemo_rl/models/policy/lm_policy.py            |  2 +
 .../policy/workers/base_policy_worker.py      | 56 +++++++-----
 .../weight_sync/mx_collective_bootstrap.py    | 90 +++++++++++++------
 .../mx_collective_weight_synchronizer.py      | 35 ++++++--
 7 files changed, 167 insertions(+), 71 deletions(-)

diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py
index e957950b96f..157ca67ba8c 100644
--- a/nemo_rl/models/generation/vllm/vllm_backend.py
+++ b/nemo_rl/models/generation/vllm/vllm_backend.py
@@ -293,29 +293,42 @@ def init_mx_reshard_comm_group(
         generator_slots: list,
         source_partition_count: int,
         plan_digest: str = "",
+        phase: object = "all",
     ) -> None:
-        """Bootstrap this gen worker's comm groups through ModelExpress."""
-        import os
+        """Bootstrap this gen worker's comm groups through ModelExpress.
 
-        from nemo_rl.weight_sync.mx_collective_bootstrap import build_mx_groups
+        Phased for the same reason the trainer side is: one lane at a time,
+        cluster-wide, with the driver barriering between them.
+        """
+        import os
 
-        local_rank = torch.distributed.get_rank()
-        index_in_role = rank_prefix + local_rank
-        reshard_groups, broadcast, _ = build_mx_groups(
-            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=f"gen-{index_in_role}-{os.getpid()}",
-            trainer_slots=trainer_slots,
-            generator_slots=generator_slots,
-            source_partition_count=source_partition_count,
-            plan_digest=plan_digest,
-            device=self.device,
+        from nemo_rl.weight_sync.mx_collective_bootstrap import (
+            mx_init_lane,
+            mx_rendezvous,
         )
-        self.pp_comm_groups = reshard_groups
-        self.model_update_group = broadcast
+
+        if phase == "rendezvous":
+            local_rank = torch.distributed.get_rank()
+            index_in_role = rank_prefix + local_rank
+            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=f"gen-{index_in_role}-{os.getpid()}",
+                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."""
diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py
index e40a7db3b88..599421933c2 100644
--- a/nemo_rl/models/generation/vllm/vllm_generation.py
+++ b/nemo_rl/models/generation/vllm/vllm_generation.py
@@ -1151,6 +1151,7 @@ def init_mx_reshard_comm_group(
         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:
@@ -1177,6 +1178,7 @@ def init_mx_reshard_comm_group(
                 "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
diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py
index 1a289e4ab6d..892bc1bf0e2 100644
--- a/nemo_rl/models/generation/vllm/vllm_worker.py
+++ b/nemo_rl/models/generation/vllm/vllm_worker.py
@@ -1240,6 +1240,7 @@ def init_mx_reshard_comm_group(
         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(
@@ -1252,6 +1253,7 @@ def init_mx_reshard_comm_group(
                 generator_slots,
                 source_partition_count,
                 plan_digest,
+                phase,
             ),
         )
 
diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py
index 53767a51424..03d1f5e1c29 100644
--- a/nemo_rl/models/policy/lm_policy.py
+++ b/nemo_rl/models/policy/lm_policy.py
@@ -1161,6 +1161,7 @@ def init_mx_reshard_comm_group(
         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(
@@ -1173,6 +1174,7 @@ def init_mx_reshard_comm_group(
                 "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
diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py
index a2825172f42..dbf4a82f0bc 100644
--- a/nemo_rl/models/policy/workers/base_policy_worker.py
+++ b/nemo_rl/models/policy/workers/base_policy_worker.py
@@ -92,35 +92,51 @@ def init_mx_reshard_comm_group(
         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.
         """
         import os
 
-        from nemo_rl.weight_sync.mx_collective_bootstrap import build_mx_groups
-
-        reshard_groups, broadcast, _ = build_mx_groups(
-            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=f"train-{self.rank}-{os.getpid()}",
-            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(),
+        from nemo_rl.weight_sync.mx_collective_bootstrap import (
+            mx_init_lane,
+            mx_rendezvous,
         )
-        # A trainer belongs to exactly one source partition, so it holds one
-        # reshard lane; generators hold all of them.
-        self.pp_comm_group = reshard_groups[my_pp_stage]
-        self.model_update_group = broadcast
-        self.my_pp_stage = my_pp_stage
+
+        if phase == "rendezvous":
+            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=f"train-{self.rank}-{os.getpid()}",
+                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,
diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
index b17e69ebbcb..209201b2674 100644
--- a/nemo_rl/weight_sync/mx_collective_bootstrap.py
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -66,7 +66,23 @@ def broadcast(self, tensor, src, stream=None):
         )
 
 
-def build_mx_groups(
+class MxBootstrapState:
+    """Rendezvous result plus the lane communicators built so far.
+
+    Held across phases because lane creation is driven from the driver, one
+    lane at a time with a barrier between them.
+    """
+
+    def __init__(self, membership, ids, device, worker_id):
+        self.membership = membership
+        self.ids = ids
+        self.device = device
+        self.worker_id = worker_id
+        self.reshard_groups = {}
+        self.broadcast_group = None
+
+
+def mx_rendezvous(
     *,
     mx_server_url: str,
     model_name: str,
@@ -81,12 +97,11 @@ def build_mx_groups(
     plan_digest: str = "",
     device=None,
     timeout_s: float = 900.0,
-):
-    """Join the MX group and return this worker's initialized lane groups.
+) -> MxBootstrapState:
+    """Join the MX group and fetch every lane's identifier.
 
-    Returns ``(reshard_group, broadcast_group, lane_ids)``. The reshard group is
-    this worker's own source partition on the trainer side; a generator joins
-    every reshard lane, so it takes the one it was asked for.
+    No communicator is created here. Creation is a separate phase per lane so
+    the driver can barrier between them.
     """
     import grpc
     from modelexpress_rl.collective import CollectiveRendezvous, Role
@@ -127,26 +142,51 @@ def build_mx_groups(
         group_id=membership.group_id, epoch=membership.epoch, timeout_s=timeout_s
     )
     ids = {lane.lane_id: bytes(lane.nccl_unique_id) for lane in group.lanes}
-
     if device is None:
         device = torch.cuda.current_device()
+    return MxBootstrapState(membership, ids, device, worker_id)
+
+
+def mx_init_lane(state: MxBootstrapState, lane_id: int) -> None:
+    """Create this worker's communicator for one lane, if it belongs to it.
+
+    One lane at a time, cluster-wide, is not an optimisation to be undone:
+    creating two different communicators concurrently across overlapping rank
+    sets deadlocks NCCL. A worker not in ``lane_id`` returns immediately and
+    waits at the driver's barrier rather than racing ahead into its next lane.
+    """
+    lane = next((l for l in state.membership.lanes if l.lane_id == lane_id), None)
+    if lane is None:
+        return
+    pg = MxProcessGroup(
+        unique_id=state.ids[lane.lane_id],
+        rank=lane.rank_in_lane,
+        world_size=lane.world_size,
+    )
+    pg.init_nccl_communicator(device=state.device)
+    if lane.kind == "BROADCAST":
+        state.broadcast_group = pg
+    else:
+        state.reshard_groups[lane.lane_id] = pg
 
-    reshard_groups = {}
-    broadcast_group = None
-    # Lane order matters: a generator joins every reshard lane, and the trainers
-    # in each lane are blocked inside their own bootstrap until it does. Creating
-    # them in ascending lane order is what makes the two sides unblock in the
-    # same sequence rather than deadlocking against each other.
-    for lane in sorted(membership.lanes, key=lambda l: l.lane_id):
-        pg = MxProcessGroup(
-            unique_id=ids[lane.lane_id],
-            rank=lane.rank_in_lane,
-            world_size=lane.world_size,
-        )
-        pg.init_nccl_communicator(device=device)
-        if lane.kind == "BROADCAST":
-            broadcast_group = pg
-        else:
-            reshard_groups[lane.lane_id] = pg
 
-    return reshard_groups, broadcast_group, membership
+def mx_lane_order(source_partition_count: int) -> list:
+    """Cluster-wide lane creation order.
+
+    Broadcast first because every rank is in it, so the barrier that follows is
+    a full-cluster sync point; then the reshard lanes in ascending order.
+    """
+    return [source_partition_count] + list(range(source_partition_count))
+
+
+def build_mx_groups(**kwargs):
+    """Single-process convenience wrapper: rendezvous then every lane in order.
+
+    Safe only when one process drives all ranks (tests). Real deployments must
+    use the phased API so the driver can barrier between lanes.
+    """
+    spc = kwargs["source_partition_count"]
+    state = mx_rendezvous(**kwargs)
+    for lane_id in mx_lane_order(spc):
+        mx_init_lane(state, lane_id)
+    return state.reshard_groups, state.broadcast_group, state.membership
diff --git a/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
index 8547f252d78..96ba9f39e3e 100644
--- a/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
@@ -30,6 +30,7 @@
 import ray
 
 from nemo_rl.utils.timer import Timer
+from nemo_rl.weight_sync.mx_collective_bootstrap import mx_lane_order
 from nemo_rl.weight_sync.interfaces import WeightSynchronizer
 
 MX_COLLECTIVE_TRANSPORT = "mx_nccl_reshard"
@@ -120,13 +121,33 @@ def init_communicator(self) -> None:
             source_partition_count=pp_size,
             plan_digest=digest,
         )
-        # Both sides block inside the collective bootstrap, so they have to be
-        # in flight together.
-        futures_train = self._policy.init_mx_reshard_comm_group(
-            pp_stages=pp_stages, **common
-        )
-        futures_gen = self._generation.init_mx_reshard_comm_group(**common)
-        ray.get(futures_train + futures_gen)
+
+        def both(phase):
+            """Run one phase on every worker and wait for all of them.
+
+            The wait is the point. Creating two different NCCL communicators
+            concurrently across overlapping rank sets deadlocks, so no worker
+            may start lane N+1 while another is still inside lane N. The
+            TCPStore path separates its all-ranks group from its per-stage
+            groups for exactly this reason.
+            """
+            futures_train = self._policy.init_mx_reshard_comm_group(
+                pp_stages=pp_stages, phase=phase, **common
+            )
+            futures_gen = self._generation.init_mx_reshard_comm_group(
+                phase=phase, **common
+            )
+            ray.get(futures_train + futures_gen)
+
+        # Both sides block inside the rendezvous, so they have to be in flight
+        # together; no communicator is created yet.
+        both("rendezvous")
+
+        # Broadcast lane first: every rank is in it, so the barrier after it is
+        # a full-cluster sync point. Then each reshard lane, one at a time.
+        for lane_id in mx_lane_order(pp_size):
+            both(lane_id)
+        both("finish")
         print("[mx] communicators ready", flush=True)
 
         self._generation.prepare_nccl_reshard_refit_info(refit_info)

From 73237b600f3e23db62b5c57912b3709f4927a8a0 Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 12:06:24 -0700
Subject: [PATCH 04/11] fix(refit): mint each lane's ncclUniqueId in the phase
 that uses it

ncclGetUniqueId opens the bootstrap listening socket as a side effect, so the
identifier is not inert data -- it names a socket that has to still be
accepting when the lane's peers dial in.

Minting every lane's identifier at rendezvous meant that socket had to survive
an actor-method return, a driver round-trip and a barrier before its lane was
built. When it did not, the peers got ECONNREFUSED from an address that looks
entirely correct, and NCCL reported it as "remote process exited or there was
a network error" -- which reads like a dead peer or a bad fabric, not a stale
listener.

Each lane now mints and publishes its identifier inside the phase that creates
the communicator, so minting and use sit microseconds apart in the same call,
the way the TCPStore path has always had it. Waiting is per-lane rather than
for whole-group READY, because with mint-at-use the later lanes are not
published yet and waiting for READY here would deadlock the ordering it is
meant to protect.

Client-side only: PublishGroupBootstrap and GetCollectiveGroup already carry
everything this needs, so the server and the proto are untouched.

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 .../weight_sync/mx_collective_bootstrap.py    | 84 ++++++++++++++-----
 1 file changed, 64 insertions(+), 20 deletions(-)

diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
index 209201b2674..1416dff32ec 100644
--- a/nemo_rl/weight_sync/mx_collective_bootstrap.py
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -80,6 +80,9 @@ def __init__(self, membership, ids, device, worker_id):
         self.worker_id = worker_id
         self.reshard_groups = {}
         self.broadcast_group = None
+        self.rz = None
+        self.channel = None
+        self.timeout_s = 900.0
 
 
 def mx_rendezvous(
@@ -124,27 +127,54 @@ def mx_rendezvous(
         source_partition=source_partition,
     )
 
-    # Rank 0 of a lane owes it an identifier. Publishing before waiting is what
-    # lets the group reach READY at all: readiness requires every lane
-    # bootstrapped for the current epoch.
-    if membership.is_bootstrap_leader:
-        for lane in membership.lanes:
-            if lane.rank_in_lane == 0:
-                rz.publish_bootstrap(
-                    group_id=membership.group_id,
-                    epoch=membership.epoch,
-                    lane_id=lane.lane_id,
-                    worker_id=worker_id,
-                    nccl_unique_id=new_unique_id(),
-                )
-
-    group = rz.await_ready(
-        group_id=membership.group_id, epoch=membership.epoch, timeout_s=timeout_s
-    )
-    ids = {lane.lane_id: bytes(lane.nccl_unique_id) for lane in group.lanes}
+    # No identifier is minted here. ncclGetUniqueId opens the bootstrap
+    # listening socket as a side effect, and that socket has to still be
+    # accepting when the lane's peers dial in. Minting at rendezvous means it
+    # must survive an actor-method return, a driver round-trip and a barrier;
+    # peers then get "connection refused" from an address that looks correct.
+    # Each lane mints its own identifier in the phase that uses it instead.
     if device is None:
         device = torch.cuda.current_device()
-    return MxBootstrapState(membership, ids, device, worker_id)
+    state = MxBootstrapState(membership, {}, device, worker_id)
+    state.rz = rz
+    state.channel = channel
+    state.timeout_s = timeout_s
+    return state
+
+
+def _await_lane_id(state: MxBootstrapState, lane_id: int) -> bytes:
+    """Block until MX has this lane's identifier at the admitted epoch.
+
+    Per-lane rather than whole-group: with mint-at-use the later lanes are not
+    published yet, so waiting for group READY here would deadlock the very
+    ordering it is meant to protect.
+    """
+    import time
+
+    from modelexpress_rl import refit_collective_pb2 as pb
+    from modelexpress_rl import refit_collective_pb2_grpc as pb_grpc
+
+    stub = pb_grpc.RefitCollectiveServiceStub(state.channel)
+    deadline = time.monotonic() + state.timeout_s
+    while True:
+        group = stub.GetCollectiveGroup(
+            pb.GetCollectiveGroupRequest(group_id=state.membership.group_id),
+            timeout=30.0,
+        )
+        if group.epoch != state.membership.epoch:
+            raise RuntimeError(
+                f"collective group epoch moved {state.membership.epoch} -> "
+                f"{group.epoch} while bootstrapping lane {lane_id}"
+            )
+        for lane in group.lanes:
+            if lane.lane_id == lane_id and len(lane.nccl_unique_id) == 128:
+                return bytes(lane.nccl_unique_id)
+        if time.monotonic() >= deadline:
+            raise TimeoutError(
+                f"lane {lane_id} identifier not published within "
+                f"{state.timeout_s:.0f}s"
+            )
+        time.sleep(0.2)
 
 
 def mx_init_lane(state: MxBootstrapState, lane_id: int) -> None:
@@ -158,8 +188,22 @@ def mx_init_lane(state: MxBootstrapState, lane_id: int) -> None:
     lane = next((l for l in state.membership.lanes if l.lane_id == lane_id), None)
     if lane is None:
         return
+    from modelexpress_rl.collective.comm import new_unique_id
+
+    if lane.rank_in_lane == 0:
+        # Mint and publish in the same breath as the init below, so the
+        # listening socket ncclGetUniqueId opens is still accepting when the
+        # peers dial it.
+        state.rz.publish_bootstrap(
+            group_id=state.membership.group_id,
+            epoch=state.membership.epoch,
+            lane_id=lane.lane_id,
+            worker_id=state.worker_id,
+            nccl_unique_id=new_unique_id(),
+        )
+    uid = _await_lane_id(state, lane.lane_id)
     pg = MxProcessGroup(
-        unique_id=state.ids[lane.lane_id],
+        unique_id=uid,
         rank=lane.rank_in_lane,
         world_size=lane.world_size,
     )

From 0824604aac3474e1ef8d6a2fd1e74e34395bd186 Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 12:32:09 -0700
Subject: [PATCH 05/11] fix(refit): fence MX communicator bootstrap

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 .../weight_sync/mx_collective_bootstrap.py    |  70 +++++-
 .../mx_collective_weight_synchronizer.py      |   8 +-
 .../test_mx_collective_bootstrap.py           | 202 ++++++++++++++++++
 3 files changed, 268 insertions(+), 12 deletions(-)
 create mode 100644 tests/unit/weight_sync/test_mx_collective_bootstrap.py

diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
index 1416dff32ec..a79d5464662 100644
--- a/nemo_rl/weight_sync/mx_collective_bootstrap.py
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -27,6 +27,9 @@
 they are the same code path.
 """
 
+import os
+from typing import Any
+
 import torch
 
 
@@ -38,11 +41,22 @@ class MxProcessGroup:
     would drag in the ``TCPStore`` this exists to replace.
     """
 
-    def __init__(self, *, unique_id: bytes, rank: int, world_size: int):
+    def __init__(
+        self,
+        *,
+        unique_id: bytes,
+        rank: int,
+        world_size: int,
+        root_unique_id: Any = None,
+    ):
         self.rank = rank
         self.world_size = world_size
         self.nccl_communicator = None
         self._unique_id = unique_id
+        # Rank zero keeps and consumes the exact object returned by
+        # ncclGetUniqueId, matching StatelessProcessGroup. Peers reconstruct
+        # it from the bytes MX brokered.
+        self._root_unique_id = root_unique_id
 
     def init_nccl_communicator(self, device):
         from nccl.core.communicator import Communicator
@@ -52,12 +66,29 @@ def init_nccl_communicator(self, device):
         # transport buffers have device-memory headroom.
         torch.cuda.empty_cache()
         with torch.cuda.device(device):
+            unique_id = self._root_unique_id
+            if unique_id is None:
+                unique_id = UniqueId.from_bytes(self._unique_id)
             self.nccl_communicator = Communicator.init(
                 nranks=self.world_size,
                 rank=self.rank,
-                unique_id=UniqueId.from_bytes(self._unique_id),
+                unique_id=unique_id,
             )
 
+            # Match the native StatelessProcessGroup bootstrap protocol. This
+            # proves the newly-created communicator can execute one collective
+            # before it is handed to the refit path.
+            stream = torch.cuda.current_stream()
+            data = (
+                torch.ones(1, device=device)
+                if self.rank == 0
+                else torch.zeros(1, device=device)
+            )
+            self.broadcast(data, 0, stream=stream)
+            stream.synchronize()
+            if not torch.allclose(data, torch.ones(1, device=device)):
+                raise RuntimeError("MX NCCL communicator bootstrap broadcast failed")
+
     def broadcast(self, tensor, src, stream=None):
         if stream is None:
             stream = torch.cuda.current_stream()
@@ -106,9 +137,16 @@ def mx_rendezvous(
     No communicator is created here. Creation is a separate phase per lane so
     the driver can barrier between them.
     """
+    if role not in ("TRAINER", "GENERATOR"):
+        raise ValueError(f"unsupported MX collective role {role!r}")
+    if os.environ.get("NCCL_COMM_ID"):
+        raise RuntimeError(
+            "NCCL_COMM_ID must be unset for the MX collective transport; "
+            "it overrides ncclGetUniqueId and bypasses MX's per-lane bootstrap"
+        )
+
     import grpc
     from modelexpress_rl.collective import CollectiveRendezvous, Role
-    from modelexpress_rl.collective.comm import new_unique_id
 
     role_enum = Role.TRAINER if role == "TRAINER" else Role.GENERATOR
     channel = grpc.insecure_channel(mx_server_url)
@@ -167,12 +205,15 @@ def _await_lane_id(state: MxBootstrapState, lane_id: int) -> bytes:
                 f"{group.epoch} while bootstrapping lane {lane_id}"
             )
         for lane in group.lanes:
-            if lane.lane_id == lane_id and len(lane.nccl_unique_id) == 128:
+            if (
+                lane.lane_id == lane_id
+                and lane.bootstrap_epoch == state.membership.epoch
+                and len(lane.nccl_unique_id) == 128
+            ):
                 return bytes(lane.nccl_unique_id)
         if time.monotonic() >= deadline:
             raise TimeoutError(
-                f"lane {lane_id} identifier not published within "
-                f"{state.timeout_s:.0f}s"
+                f"lane {lane_id} identifier not published within {state.timeout_s:.0f}s"
             )
         time.sleep(0.2)
 
@@ -188,24 +229,35 @@ def mx_init_lane(state: MxBootstrapState, lane_id: int) -> None:
     lane = next((l for l in state.membership.lanes if l.lane_id == lane_id), None)
     if lane is None:
         return
-    from modelexpress_rl.collective.comm import new_unique_id
-
+    root_unique_id = None
+    minted_uid = None
     if lane.rank_in_lane == 0:
         # Mint and publish in the same breath as the init below, so the
         # listening socket ncclGetUniqueId opens is still accepting when the
         # peers dial it.
+        from nccl.core.utils import get_unique_id
+
+        root_unique_id = get_unique_id()
+        minted_uid = bytes(root_unique_id.as_bytes)
         state.rz.publish_bootstrap(
             group_id=state.membership.group_id,
             epoch=state.membership.epoch,
             lane_id=lane.lane_id,
             worker_id=state.worker_id,
-            nccl_unique_id=new_unique_id(),
+            nccl_unique_id=minted_uid,
         )
     uid = _await_lane_id(state, lane.lane_id)
+    if minted_uid is not None and uid != minted_uid:
+        raise RuntimeError(
+            f"MX returned a different ncclUniqueId for lane {lane.lane_id} "
+            "than rank zero published"
+        )
+    state.ids[lane.lane_id] = uid
     pg = MxProcessGroup(
         unique_id=uid,
         rank=lane.rank_in_lane,
         world_size=lane.world_size,
+        root_unique_id=root_unique_id,
     )
     pg.init_nccl_communicator(device=state.device)
     if lane.kind == "BROADCAST":
diff --git a/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
index 96ba9f39e3e..dabe95f4ec9 100644
--- a/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py
@@ -97,9 +97,11 @@ def init_communicator(self) -> None:
             from modelexpress_rl.collective import plan_digest
 
             digest = plan_digest(build_mx_plan(refit_info))
-        except Exception as error:  # noqa: BLE001 - digest is an agreement token
-            print(f"[mx] plan digest unavailable ({error!r}); using a constant", flush=True)
-            digest = "nccl-reshard-plan"
+        except Exception as error:  # noqa: BLE001 - normalize an optional integration boundary
+            raise RuntimeError(
+                "failed to derive the MX collective reshard-plan digest; "
+                "refusing to form an unfenced collective"
+            ) from error
 
         trainers = [f"train/{r}" for r in range(train_world_size)]
         generators = [f"gen/{r}" for r in range(gen_world_size)]
diff --git a/tests/unit/weight_sync/test_mx_collective_bootstrap.py b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
new file mode 100644
index 00000000000..b16ad9d64ed
--- /dev/null
+++ b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
@@ -0,0 +1,202 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import sys
+from contextlib import nullcontext
+from types import ModuleType, SimpleNamespace
+
+import pytest
+
+from nemo_rl.weight_sync import mx_collective_bootstrap as bootstrap
+
+
+def rendezvous_kwargs(**overrides):
+    kwargs = {
+        "mx_server_url": "mx:8001",
+        "model_name": "model",
+        "role": "TRAINER",
+        "index_in_role": 0,
+        "slot_id": "train/0",
+        "worker_id": "worker-0",
+        "trainer_slots": ["train/0"],
+        "generator_slots": ["gen/0"],
+        "source_partition_count": 1,
+    }
+    kwargs.update(overrides)
+    return kwargs
+
+
+def bootstrap_state(*, rank=0, epoch=7):
+    lane = SimpleNamespace(lane_id=1, rank_in_lane=rank, world_size=2, kind="BROADCAST")
+    membership = SimpleNamespace(group_id="group", epoch=epoch, lanes=(lane,))
+    return bootstrap.MxBootstrapState(
+        membership, {}, device=0, worker_id=f"worker-{rank}"
+    )
+
+
+class FakeRendezvous:
+    def __init__(self):
+        self.published = []
+
+    def publish_bootstrap(self, **kwargs):
+        self.published.append(kwargs)
+
+
+class RootUniqueId:
+    def __init__(self, raw):
+        self.as_bytes = raw
+
+
+def install_unique_id_module(monkeypatch, unique_id):
+    utils = ModuleType("nccl.core.utils")
+    utils.get_unique_id = lambda: unique_id
+    core = ModuleType("nccl.core")
+    nccl = ModuleType("nccl")
+    monkeypatch.setitem(sys.modules, "nccl", nccl)
+    monkeypatch.setitem(sys.modules, "nccl.core", core)
+    monkeypatch.setitem(sys.modules, "nccl.core.utils", utils)
+
+
+def test_rendezvous_rejects_an_unknown_role_before_joining():
+    with pytest.raises(ValueError, match="unsupported MX collective role"):
+        bootstrap.mx_rendezvous(**rendezvous_kwargs(role="OBSERVER"))
+
+
+def test_rendezvous_rejects_nccl_comm_id_override(monkeypatch):
+    monkeypatch.setenv("NCCL_COMM_ID", "10.0.0.1:1234")
+    with pytest.raises(RuntimeError, match="must be unset"):
+        bootstrap.mx_rendezvous(**rendezvous_kwargs())
+
+
+def test_lane_root_rejects_a_uid_different_from_what_it_published(monkeypatch):
+    minted = b"a" * 128
+    install_unique_id_module(monkeypatch, RootUniqueId(minted))
+    state = bootstrap_state()
+    state.rz = FakeRendezvous()
+    monkeypatch.setattr(bootstrap, "_await_lane_id", lambda *_: b"b" * 128)
+
+    with pytest.raises(RuntimeError, match="different ncclUniqueId"):
+        bootstrap.mx_init_lane(state, 1)
+    assert state.rz.published[0]["nccl_unique_id"] == minted
+
+
+def test_lane_root_keeps_the_minted_uid_object_through_communicator_init(monkeypatch):
+    minted = b"a" * 128
+    root_unique_id = RootUniqueId(minted)
+    install_unique_id_module(monkeypatch, root_unique_id)
+    state = bootstrap_state()
+    state.rz = FakeRendezvous()
+    monkeypatch.setattr(bootstrap, "_await_lane_id", lambda *_: minted)
+    observed = {}
+
+    class FakeProcessGroup:
+        def __init__(self, **kwargs):
+            observed.update(kwargs)
+
+        def init_nccl_communicator(self, device):
+            observed["device"] = device
+
+    monkeypatch.setattr(bootstrap, "MxProcessGroup", FakeProcessGroup)
+    bootstrap.mx_init_lane(state, 1)
+
+    assert observed["root_unique_id"] is root_unique_id
+    assert observed["unique_id"] == minted
+    assert state.ids[1] == minted
+
+
+def test_lane_fetch_requires_a_bootstrap_stamp_from_the_current_epoch(monkeypatch):
+    from modelexpress_rl import refit_collective_pb2_grpc as pb_grpc
+
+    stale = SimpleNamespace(lane_id=1, bootstrap_epoch=6, nccl_unique_id=b"s" * 128)
+    current = SimpleNamespace(lane_id=1, bootstrap_epoch=7, nccl_unique_id=b"c" * 128)
+
+    class Stub:
+        def __init__(self):
+            self.responses = [
+                SimpleNamespace(epoch=7, lanes=[stale]),
+                SimpleNamespace(epoch=7, lanes=[current]),
+            ]
+
+        def GetCollectiveGroup(self, request, timeout):
+            return self.responses.pop(0)
+
+    stub = Stub()
+    monkeypatch.setattr(pb_grpc, "RefitCollectiveServiceStub", lambda channel: stub)
+    monkeypatch.setattr("time.sleep", lambda _: None)
+    state = bootstrap_state()
+    state.channel = object()
+    state.timeout_s = 1.0
+
+    assert bootstrap._await_lane_id(state, 1) == b"c" * 128
+
+
+def test_process_group_uses_root_uid_and_warms_up_the_communicator(monkeypatch):
+    root_unique_id = object()
+    events = []
+
+    class FakeTensor:
+        def __init__(self, value):
+            self.value = value
+
+    class FakeStream:
+        cuda_stream = 17
+
+        def synchronize(self):
+            events.append("sync")
+
+    class FakeCommunicator:
+        @staticmethod
+        def init(*, nranks, rank, unique_id):
+            events.append(("init", nranks, rank, unique_id))
+            return FakeCommunicator()
+
+        def broadcast(self, *, sendbuf, recvbuf, root, stream):
+            recvbuf.value = 1
+            events.append(("broadcast", root, stream))
+
+    class FakeUniqueId:
+        @staticmethod
+        def from_bytes(raw):
+            raise AssertionError("rank zero must use the original unique-id object")
+
+    communicator = ModuleType("nccl.core.communicator")
+    communicator.Communicator = FakeCommunicator
+    utils = ModuleType("nccl.core.utils")
+    utils.UniqueId = FakeUniqueId
+    monkeypatch.setitem(sys.modules, "nccl", ModuleType("nccl"))
+    monkeypatch.setitem(sys.modules, "nccl.core", ModuleType("nccl.core"))
+    monkeypatch.setitem(sys.modules, "nccl.core.communicator", communicator)
+    monkeypatch.setitem(sys.modules, "nccl.core.utils", utils)
+    monkeypatch.setattr(bootstrap.torch.cuda, "empty_cache", lambda: None)
+    monkeypatch.setattr(bootstrap.torch.cuda, "device", lambda device: nullcontext())
+    monkeypatch.setattr(bootstrap.torch.cuda, "current_stream", FakeStream)
+    monkeypatch.setattr(bootstrap.torch, "ones", lambda *args, **kwargs: FakeTensor(1))
+    monkeypatch.setattr(bootstrap.torch, "zeros", lambda *args, **kwargs: FakeTensor(0))
+    monkeypatch.setattr(
+        bootstrap.torch, "allclose", lambda left, right: left.value == right.value
+    )
+
+    group = bootstrap.MxProcessGroup(
+        unique_id=b"u" * 128,
+        rank=0,
+        world_size=2,
+        root_unique_id=root_unique_id,
+    )
+    group.init_nccl_communicator(device=0)
+
+    assert events == [
+        ("init", 2, 0, root_unique_id),
+        ("broadcast", 0, 17),
+        "sync",
+    ]

From e961ca94c402f32df159563b66a3284e215abc24 Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 12:33:43 -0700
Subject: [PATCH 06/11] fix(refit): make MX worker generations unique

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 nemo_rl/models/generation/vllm/vllm_backend.py      | 10 +++++++---
 nemo_rl/models/policy/workers/base_policy_worker.py | 10 +++++++---
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py
index 157ca67ba8c..9e86071ab1e 100644
--- a/nemo_rl/models/generation/vllm/vllm_backend.py
+++ b/nemo_rl/models/generation/vllm/vllm_backend.py
@@ -300,23 +300,27 @@ def init_mx_reshard_comm_group(
         Phased for the same reason the trainer side is: one lane at a time,
         cluster-wide, with the driver barriering between them.
         """
-        import os
-
         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=f"gen-{index_in_role}-{os.getpid()}",
+                worker_id=self._mx_worker_id,
                 trainer_slots=trainer_slots,
                 generator_slots=generator_slots,
                 source_partition_count=source_partition_count,
diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py
index dbf4a82f0bc..1667ee11dcc 100644
--- a/nemo_rl/models/policy/workers/base_policy_worker.py
+++ b/nemo_rl/models/policy/workers/base_policy_worker.py
@@ -105,21 +105,25 @@ def init_mx_reshard_comm_group(
         deadlocks NCCL -- the same reason the TCPStore path separates its
         all-ranks group from its per-stage groups.
         """
-        import os
-
         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=f"train-{self.rank}-{os.getpid()}",
+                worker_id=self._mx_worker_id,
                 trainer_slots=trainer_slots,
                 generator_slots=generator_slots,
                 source_partition_count=source_partition_count,

From ee196793aab91d057658b9a525aa7ee725489398 Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 13:03:41 -0700
Subject: [PATCH 07/11] fix(refit): bound MX communicator initialization

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 .../weight_sync/mx_collective_bootstrap.py    | 110 +++++++++++++++---
 .../test_mx_collective_bootstrap.py           |  72 +++++++++++-
 2 files changed, 164 insertions(+), 18 deletions(-)

diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
index a79d5464662..cdda21d86d3 100644
--- a/nemo_rl/weight_sync/mx_collective_bootstrap.py
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -27,7 +27,9 @@
 they are the same code path.
 """
 
+import math
 import os
+import time
 from typing import Any
 
 import torch
@@ -48,6 +50,7 @@ def __init__(
         rank: int,
         world_size: int,
         root_unique_id: Any = None,
+        timeout_s: float = 900.0,
     ):
         self.rank = rank
         self.world_size = world_size
@@ -57,9 +60,11 @@ def __init__(
         # ncclGetUniqueId, matching StatelessProcessGroup. Peers reconstruct
         # it from the bytes MX brokered.
         self._root_unique_id = root_unique_id
+        self._timeout_s = timeout_s
 
     def init_nccl_communicator(self, device):
-        from nccl.core.communicator import Communicator
+        from nccl.bindings import nccl as bindings
+        from nccl.core import communicator
         from nccl.core.utils import UniqueId
 
         # Mirror the native path: free cached blocks first so the communicator's
@@ -69,25 +74,71 @@ def init_nccl_communicator(self, device):
             unique_id = self._root_unique_id
             if unique_id is None:
                 unique_id = UniqueId.from_bytes(self._unique_id)
-            self.nccl_communicator = Communicator.init(
-                nranks=self.world_size,
-                rank=self.rank,
-                unique_id=unique_id,
-            )
+            config_type = getattr(communicator, "NCCLConfig", None)
+            if config_type is None:
+                raise RuntimeError(
+                    "MX communicator timeouts require nccl4py NCCLConfig support"
+                )
+            try:
+                self.nccl_communicator = communicator.Communicator.init(
+                    nranks=self.world_size,
+                    rank=self.rank,
+                    unique_id=unique_id,
+                    config=config_type(blocking=False),
+                )
+                deadline = time.monotonic() + self._timeout_s
+                while True:
+                    status = self.nccl_communicator.get_async_error()
+                    if int(status) == int(bindings.Result.Success):
+                        break
+                    if int(status) != int(bindings.Result.InProgress):
+                        raise RuntimeError(
+                            "MX NCCL communicator initialization failed with "
+                            f"{status!r}: {self.nccl_communicator.get_last_error()}"
+                        )
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        raise TimeoutError(
+                            "MX NCCL communicator initialization did not complete "
+                            f"within {self._timeout_s:.1f}s"
+                        )
+                    time.sleep(min(0.01, remaining))
+            except BaseException:
+                try:
+                    self.abort()
+                except Exception:
+                    pass
+                raise
 
             # Match the native StatelessProcessGroup bootstrap protocol. This
             # proves the newly-created communicator can execute one collective
             # before it is handed to the refit path.
-            stream = torch.cuda.current_stream()
-            data = (
-                torch.ones(1, device=device)
-                if self.rank == 0
-                else torch.zeros(1, device=device)
-            )
-            self.broadcast(data, 0, stream=stream)
-            stream.synchronize()
-            if not torch.allclose(data, torch.ones(1, device=device)):
-                raise RuntimeError("MX NCCL communicator bootstrap broadcast failed")
+            try:
+                stream = torch.cuda.current_stream()
+                data = (
+                    torch.ones(1, device=device)
+                    if self.rank == 0
+                    else torch.zeros(1, device=device)
+                )
+                self.broadcast(data, 0, stream=stream)
+                stream.synchronize()
+                if not torch.allclose(data, torch.ones(1, device=device)):
+                    raise RuntimeError(
+                        "MX NCCL communicator bootstrap broadcast failed"
+                    )
+            except BaseException:
+                try:
+                    self.abort()
+                except Exception:
+                    pass
+                raise
+
+    def abort(self) -> None:
+        """Abort a partially or fully initialized communicator exactly once."""
+        communicator = self.nccl_communicator
+        self.nccl_communicator = None
+        if communicator is not None:
+            communicator.abort()
 
     def broadcast(self, tensor, src, stream=None):
         if stream is None:
@@ -115,6 +166,19 @@ def __init__(self, membership, ids, device, worker_id):
         self.channel = None
         self.timeout_s = 900.0
 
+    def abort(self) -> None:
+        """Abort every lane already created for this group."""
+        groups = list(self.reshard_groups.values())
+        if self.broadcast_group is not None:
+            groups.append(self.broadcast_group)
+        self.reshard_groups.clear()
+        self.broadcast_group = None
+        for group in groups:
+            try:
+                group.abort()
+            except Exception:
+                pass
+
 
 def mx_rendezvous(
     *,
@@ -144,6 +208,13 @@ def mx_rendezvous(
             "NCCL_COMM_ID must be unset for the MX collective transport; "
             "it overrides ncclGetUniqueId and bypasses MX's per-lane bootstrap"
         )
+    if os.environ.get("NCCL_COMM_BLOCKING") not in (None, "", "0"):
+        raise RuntimeError(
+            "NCCL_COMM_BLOCKING must be unset or 0 for the MX collective "
+            "transport so communicator initialization remains bounded"
+        )
+    if not math.isfinite(timeout_s) or timeout_s <= 0:
+        raise ValueError(f"timeout_s must be finite and positive, got {timeout_s}")
 
     import grpc
     from modelexpress_rl.collective import CollectiveRendezvous, Role
@@ -258,8 +329,13 @@ def mx_init_lane(state: MxBootstrapState, lane_id: int) -> None:
         rank=lane.rank_in_lane,
         world_size=lane.world_size,
         root_unique_id=root_unique_id,
+        timeout_s=state.timeout_s,
     )
-    pg.init_nccl_communicator(device=state.device)
+    try:
+        pg.init_nccl_communicator(device=state.device)
+    except BaseException:
+        state.abort()
+        raise
     if lane.kind == "BROADCAST":
         state.broadcast_group = pg
     else:
diff --git a/tests/unit/weight_sync/test_mx_collective_bootstrap.py b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
index b16ad9d64ed..658d29ded20 100644
--- a/tests/unit/weight_sync/test_mx_collective_bootstrap.py
+++ b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
@@ -157,14 +157,25 @@ def synchronize(self):
 
     class FakeCommunicator:
         @staticmethod
-        def init(*, nranks, rank, unique_id):
+        def init(*, nranks, rank, unique_id, config):
+            assert config.blocking is False
             events.append(("init", nranks, rank, unique_id))
             return FakeCommunicator()
 
+        def get_async_error(self):
+            return 0
+
         def broadcast(self, *, sendbuf, recvbuf, root, stream):
             recvbuf.value = 1
             events.append(("broadcast", root, stream))
 
+        def abort(self):
+            events.append("abort")
+
+    class FakeConfig:
+        def __init__(self, *, blocking):
+            self.blocking = blocking
+
     class FakeUniqueId:
         @staticmethod
         def from_bytes(raw):
@@ -172,10 +183,15 @@ def from_bytes(raw):
 
     communicator = ModuleType("nccl.core.communicator")
     communicator.Communicator = FakeCommunicator
+    communicator.NCCLConfig = FakeConfig
     utils = ModuleType("nccl.core.utils")
     utils.UniqueId = FakeUniqueId
+    bindings = ModuleType("nccl.bindings.nccl")
+    bindings.Result = SimpleNamespace(Success=0, InProgress=7)
     monkeypatch.setitem(sys.modules, "nccl", ModuleType("nccl"))
     monkeypatch.setitem(sys.modules, "nccl.core", ModuleType("nccl.core"))
+    monkeypatch.setitem(sys.modules, "nccl.bindings", ModuleType("nccl.bindings"))
+    monkeypatch.setitem(sys.modules, "nccl.bindings.nccl", bindings)
     monkeypatch.setitem(sys.modules, "nccl.core.communicator", communicator)
     monkeypatch.setitem(sys.modules, "nccl.core.utils", utils)
     monkeypatch.setattr(bootstrap.torch.cuda, "empty_cache", lambda: None)
@@ -200,3 +216,57 @@ def from_bytes(raw):
         ("broadcast", 0, 17),
         "sync",
     ]
+
+
+def test_process_group_aborts_a_nonblocking_init_that_times_out(monkeypatch):
+    events = []
+
+    class FakeCommunicator:
+        @staticmethod
+        def init(**kwargs):
+            return FakeCommunicator()
+
+        def get_async_error(self):
+            return 7
+
+        def abort(self):
+            events.append("abort")
+
+    class FakeConfig:
+        def __init__(self, *, blocking):
+            self.blocking = blocking
+
+    class FakeUniqueId:
+        @staticmethod
+        def from_bytes(raw):
+            return raw
+
+    communicator = ModuleType("nccl.core.communicator")
+    communicator.Communicator = FakeCommunicator
+    communicator.NCCLConfig = FakeConfig
+    utils = ModuleType("nccl.core.utils")
+    utils.UniqueId = FakeUniqueId
+    bindings = ModuleType("nccl.bindings.nccl")
+    bindings.Result = SimpleNamespace(Success=0, InProgress=7)
+    monkeypatch.setitem(sys.modules, "nccl", ModuleType("nccl"))
+    monkeypatch.setitem(sys.modules, "nccl.core", ModuleType("nccl.core"))
+    monkeypatch.setitem(sys.modules, "nccl.bindings", ModuleType("nccl.bindings"))
+    monkeypatch.setitem(sys.modules, "nccl.bindings.nccl", bindings)
+    monkeypatch.setitem(sys.modules, "nccl.core.communicator", communicator)
+    monkeypatch.setitem(sys.modules, "nccl.core.utils", utils)
+    monkeypatch.setattr(bootstrap.torch.cuda, "empty_cache", lambda: None)
+    monkeypatch.setattr(bootstrap.torch.cuda, "device", lambda device: nullcontext())
+    ticks = iter((0.0, 0.0, 0.2))
+    monkeypatch.setattr(bootstrap.time, "monotonic", lambda: next(ticks))
+    monkeypatch.setattr(bootstrap.time, "sleep", lambda _: None)
+
+    group = bootstrap.MxProcessGroup(
+        unique_id=b"u" * 128,
+        rank=1,
+        world_size=2,
+        timeout_s=0.1,
+    )
+    with pytest.raises(TimeoutError, match="did not complete"):
+        group.init_nccl_communicator(device=0)
+
+    assert events == ["abort"]

From 3ab5aec9784712e2af58c8cdfb266d4312ed532d Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 13:19:10 -0700
Subject: [PATCH 08/11] fix(refit): expire failed MX bootstrap leases

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 nemo_rl/weight_sync/mx_collective_bootstrap.py      | 13 +++++++++++++
 .../weight_sync/test_mx_collective_bootstrap.py     | 11 +++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
index cdda21d86d3..3988869850a 100644
--- a/nemo_rl/weight_sync/mx_collective_bootstrap.py
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -178,6 +178,18 @@ def abort(self) -> None:
                 group.abort()
             except Exception:
                 pass
+        rendezvous = self.rz
+        self.rz = None
+        if rendezvous is not None:
+            close = getattr(rendezvous, "close", None)
+            if callable(close):
+                close()
+        channel = self.channel
+        self.channel = None
+        if channel is not None:
+            close = getattr(channel, "close", None)
+            if callable(close):
+                close()
 
 
 def mx_rendezvous(
@@ -319,6 +331,7 @@ def mx_init_lane(state: MxBootstrapState, lane_id: int) -> None:
         )
     uid = _await_lane_id(state, lane.lane_id)
     if minted_uid is not None and uid != minted_uid:
+        state.abort()
         raise RuntimeError(
             f"MX returned a different ncclUniqueId for lane {lane.lane_id} "
             "than rank zero published"
diff --git a/tests/unit/weight_sync/test_mx_collective_bootstrap.py b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
index 658d29ded20..4d012061115 100644
--- a/tests/unit/weight_sync/test_mx_collective_bootstrap.py
+++ b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
@@ -48,10 +48,14 @@ def bootstrap_state(*, rank=0, epoch=7):
 class FakeRendezvous:
     def __init__(self):
         self.published = []
+        self.closed = False
 
     def publish_bootstrap(self, **kwargs):
         self.published.append(kwargs)
 
+    def close(self):
+        self.closed = True
+
 
 class RootUniqueId:
     def __init__(self, raw):
@@ -83,12 +87,15 @@ def test_lane_root_rejects_a_uid_different_from_what_it_published(monkeypatch):
     minted = b"a" * 128
     install_unique_id_module(monkeypatch, RootUniqueId(minted))
     state = bootstrap_state()
-    state.rz = FakeRendezvous()
+    rendezvous = FakeRendezvous()
+    state.rz = rendezvous
     monkeypatch.setattr(bootstrap, "_await_lane_id", lambda *_: b"b" * 128)
 
     with pytest.raises(RuntimeError, match="different ncclUniqueId"):
         bootstrap.mx_init_lane(state, 1)
-    assert state.rz.published[0]["nccl_unique_id"] == minted
+    assert rendezvous.published[0]["nccl_unique_id"] == minted
+    assert rendezvous.closed
+    assert state.rz is None
 
 
 def test_lane_root_keeps_the_minted_uid_object_through_communicator_init(monkeypatch):

From 5f9f6ac748d56ea8c7fca3f0b37ae62af332d0fb Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 14:08:02 -0700
Subject: [PATCH 09/11] fix(refit): wait for nonblocking bootstrap collective

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 .../weight_sync/mx_collective_bootstrap.py    | 54 ++++++++----
 .../test_mx_collective_bootstrap.py           | 87 +++++++++++++++++++
 2 files changed, 124 insertions(+), 17 deletions(-)

diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
index 3988869850a..728b5b5af3d 100644
--- a/nemo_rl/weight_sync/mx_collective_bootstrap.py
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -62,6 +62,32 @@ def __init__(
         self._root_unique_id = root_unique_id
         self._timeout_s = timeout_s
 
+    def _wait_for_async_success(self, bindings, *, operation: str) -> None:
+        """Drive a nonblocking NCCL host operation to completion.
+
+        A nonblocking communicator can return ``ncclInProgress`` not only
+        from communicator creation, but also while the first collective lazily
+        connects its transports.  Until that state becomes ``ncclSuccess`` the
+        collective may not have been posted to the CUDA stream yet.
+        """
+        deadline = time.monotonic() + self._timeout_s
+        while True:
+            status = self.nccl_communicator.get_async_error()
+            if int(status) == int(bindings.Result.Success):
+                return
+            if int(status) != int(bindings.Result.InProgress):
+                raise RuntimeError(
+                    f"MX NCCL {operation} failed with {status!r}: "
+                    f"{self.nccl_communicator.get_last_error()}"
+                )
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise TimeoutError(
+                    f"MX NCCL {operation} did not complete within "
+                    f"{self._timeout_s:.1f}s"
+                )
+            time.sleep(min(0.01, remaining))
+
     def init_nccl_communicator(self, device):
         from nccl.bindings import nccl as bindings
         from nccl.core import communicator
@@ -86,23 +112,9 @@ def init_nccl_communicator(self, device):
                     unique_id=unique_id,
                     config=config_type(blocking=False),
                 )
-                deadline = time.monotonic() + self._timeout_s
-                while True:
-                    status = self.nccl_communicator.get_async_error()
-                    if int(status) == int(bindings.Result.Success):
-                        break
-                    if int(status) != int(bindings.Result.InProgress):
-                        raise RuntimeError(
-                            "MX NCCL communicator initialization failed with "
-                            f"{status!r}: {self.nccl_communicator.get_last_error()}"
-                        )
-                    remaining = deadline - time.monotonic()
-                    if remaining <= 0:
-                        raise TimeoutError(
-                            "MX NCCL communicator initialization did not complete "
-                            f"within {self._timeout_s:.1f}s"
-                        )
-                    time.sleep(min(0.01, remaining))
+                self._wait_for_async_success(
+                    bindings, operation="communicator initialization"
+                )
             except BaseException:
                 try:
                     self.abort()
@@ -121,6 +133,14 @@ def init_nccl_communicator(self, device):
                     else torch.zeros(1, device=device)
                 )
                 self.broadcast(data, 0, stream=stream)
+                # With a nonblocking communicator the first collective can
+                # still be establishing transport connections here. Polling
+                # is required before synchronizing the CUDA stream: otherwise
+                # the stream may have no NCCL work posted yet and peers can
+                # observe the untouched zero buffer, aborting a healthy lane.
+                self._wait_for_async_success(
+                    bindings, operation="communicator bootstrap broadcast"
+                )
                 stream.synchronize()
                 if not torch.allclose(data, torch.ones(1, device=device)):
                     raise RuntimeError(
diff --git a/tests/unit/weight_sync/test_mx_collective_bootstrap.py b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
index 4d012061115..27baeb1f863 100644
--- a/tests/unit/weight_sync/test_mx_collective_bootstrap.py
+++ b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
@@ -225,6 +225,93 @@ def from_bytes(raw):
     ]
 
 
+def test_process_group_waits_for_nonblocking_broadcast_before_reading(monkeypatch):
+    events = []
+    statuses = iter((0, 7, 0))
+
+    class FakeTensor:
+        def __init__(self, value):
+            self.value = value
+
+    class FakeStream:
+        cuda_stream = 17
+
+        def synchronize(self):
+            events.append("sync")
+
+    class FakeCommunicator:
+        def __init__(self):
+            self.recvbuf = None
+
+        @staticmethod
+        def init(**kwargs):
+            events.append("init")
+            return FakeCommunicator()
+
+        def get_async_error(self):
+            status = next(statuses)
+            events.append(("status", status))
+            if status == 0 and self.recvbuf is not None:
+                self.recvbuf.value = 1
+            return status
+
+        def broadcast(self, *, sendbuf, recvbuf, root, stream):
+            self.recvbuf = recvbuf
+            events.append("broadcast")
+
+        def abort(self):
+            events.append("abort")
+
+    class FakeConfig:
+        def __init__(self, *, blocking):
+            self.blocking = blocking
+
+    class FakeUniqueId:
+        @staticmethod
+        def from_bytes(raw):
+            return raw
+
+    communicator = ModuleType("nccl.core.communicator")
+    communicator.Communicator = FakeCommunicator
+    communicator.NCCLConfig = FakeConfig
+    utils = ModuleType("nccl.core.utils")
+    utils.UniqueId = FakeUniqueId
+    bindings = ModuleType("nccl.bindings.nccl")
+    bindings.Result = SimpleNamespace(Success=0, InProgress=7)
+    monkeypatch.setitem(sys.modules, "nccl", ModuleType("nccl"))
+    monkeypatch.setitem(sys.modules, "nccl.core", ModuleType("nccl.core"))
+    monkeypatch.setitem(sys.modules, "nccl.bindings", ModuleType("nccl.bindings"))
+    monkeypatch.setitem(sys.modules, "nccl.bindings.nccl", bindings)
+    monkeypatch.setitem(sys.modules, "nccl.core.communicator", communicator)
+    monkeypatch.setitem(sys.modules, "nccl.core.utils", utils)
+    monkeypatch.setattr(bootstrap.torch.cuda, "empty_cache", lambda: None)
+    monkeypatch.setattr(bootstrap.torch.cuda, "device", lambda device: nullcontext())
+    monkeypatch.setattr(bootstrap.torch.cuda, "current_stream", FakeStream)
+    monkeypatch.setattr(bootstrap.torch, "ones", lambda *args, **kwargs: FakeTensor(1))
+    monkeypatch.setattr(bootstrap.torch, "zeros", lambda *args, **kwargs: FakeTensor(0))
+    monkeypatch.setattr(
+        bootstrap.torch, "allclose", lambda left, right: left.value == right.value
+    )
+    monkeypatch.setattr(bootstrap.time, "sleep", lambda _: events.append("sleep"))
+
+    group = bootstrap.MxProcessGroup(
+        unique_id=b"u" * 128,
+        rank=1,
+        world_size=2,
+    )
+    group.init_nccl_communicator(device=0)
+
+    assert events == [
+        "init",
+        ("status", 0),
+        "broadcast",
+        ("status", 7),
+        "sleep",
+        ("status", 0),
+        "sync",
+    ]
+
+
 def test_process_group_aborts_a_nonblocking_init_that_times_out(monkeypatch):
     events = []
 

From 9a89fd8e2a709ec61b3a01b7a959341127c6622b Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 14:21:26 -0700
Subject: [PATCH 10/11] fix(refit): hand M2N a blocking child communicator

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 .../weight_sync/mx_collective_bootstrap.py    | 46 +++++++++++++++++--
 .../test_mx_collective_bootstrap.py           | 33 +++++++++++--
 2 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/nemo_rl/weight_sync/mx_collective_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py
index 728b5b5af3d..41a440eeac8 100644
--- a/nemo_rl/weight_sync/mx_collective_bootstrap.py
+++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py
@@ -55,6 +55,7 @@ def __init__(
         self.rank = rank
         self.world_size = world_size
         self.nccl_communicator = None
+        self._bootstrap_communicator = None
         self._unique_id = unique_id
         # Rank zero keeps and consumes the exact object returned by
         # ncclGetUniqueId, matching StatelessProcessGroup. Peers reconstruct
@@ -62,7 +63,9 @@ def __init__(
         self._root_unique_id = root_unique_id
         self._timeout_s = timeout_s
 
-    def _wait_for_async_success(self, bindings, *, operation: str) -> None:
+    def _wait_for_async_success(
+        self, bindings, *, operation: str, communicator=None
+    ) -> None:
         """Drive a nonblocking NCCL host operation to completion.
 
         A nonblocking communicator can return ``ncclInProgress`` not only
@@ -70,15 +73,16 @@ def _wait_for_async_success(self, bindings, *, operation: str) -> None:
         connects its transports.  Until that state becomes ``ncclSuccess`` the
         collective may not have been posted to the CUDA stream yet.
         """
+        communicator = communicator or self.nccl_communicator
         deadline = time.monotonic() + self._timeout_s
         while True:
-            status = self.nccl_communicator.get_async_error()
+            status = communicator.get_async_error()
             if int(status) == int(bindings.Result.Success):
                 return
             if int(status) != int(bindings.Result.InProgress):
                 raise RuntimeError(
                     f"MX NCCL {operation} failed with {status!r}: "
-                    f"{self.nccl_communicator.get_last_error()}"
+                    f"{communicator.get_last_error()}"
                 )
             remaining = deadline - time.monotonic()
             if remaining <= 0:
@@ -115,6 +119,30 @@ def init_nccl_communicator(self, device):
                 self._wait_for_async_success(
                     bindings, operation="communicator initialization"
                 )
+
+                # M2N currently requires ordinary blocking NCCL semantics and
+                # treats ncclInProgress from its internal collectives as an
+                # error. Keep the original communicator as a bounded bootstrap
+                # parent, then collectively split it into an independent
+                # blocking child with identical ranks for the data path.
+                bootstrap_communicator = self.nccl_communicator
+                blocking_communicator = bootstrap_communicator.split(
+                    color=0,
+                    key=self.rank,
+                    config=config_type(blocking=True),
+                )
+                self._bootstrap_communicator = bootstrap_communicator
+                self.nccl_communicator = blocking_communicator
+                self._wait_for_async_success(
+                    bindings,
+                    operation="blocking communicator split",
+                    communicator=bootstrap_communicator,
+                )
+                self._wait_for_async_success(
+                    bindings,
+                    operation="blocking communicator activation",
+                    communicator=blocking_communicator,
+                )
             except BaseException:
                 try:
                     self.abort()
@@ -155,9 +183,17 @@ def init_nccl_communicator(self, device):
 
     def abort(self) -> None:
         """Abort a partially or fully initialized communicator exactly once."""
-        communicator = self.nccl_communicator
+        communicators = (
+            self.nccl_communicator,
+            self._bootstrap_communicator,
+        )
         self.nccl_communicator = None
-        if communicator is not None:
+        self._bootstrap_communicator = None
+        seen = set()
+        for communicator in communicators:
+            if communicator is None or id(communicator) in seen:
+                continue
+            seen.add(id(communicator))
             communicator.abort()
 
     def broadcast(self, tensor, src, stream=None):
diff --git a/tests/unit/weight_sync/test_mx_collective_bootstrap.py b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
index 27baeb1f863..240aa721dab 100644
--- a/tests/unit/weight_sync/test_mx_collective_bootstrap.py
+++ b/tests/unit/weight_sync/test_mx_collective_bootstrap.py
@@ -163,21 +163,30 @@ def synchronize(self):
             events.append("sync")
 
     class FakeCommunicator:
+        def __init__(self, kind):
+            self.kind = kind
+
         @staticmethod
         def init(*, nranks, rank, unique_id, config):
             assert config.blocking is False
             events.append(("init", nranks, rank, unique_id))
-            return FakeCommunicator()
+            return FakeCommunicator("bootstrap")
+
+        def split(self, *, color, key, config):
+            assert config.blocking is True
+            events.append(("split", color, key))
+            return FakeCommunicator("blocking")
 
         def get_async_error(self):
             return 0
 
         def broadcast(self, *, sendbuf, recvbuf, root, stream):
+            assert self.kind == "blocking"
             recvbuf.value = 1
             events.append(("broadcast", root, stream))
 
         def abort(self):
-            events.append("abort")
+            events.append(f"abort-{self.kind}")
 
     class FakeConfig:
         def __init__(self, *, blocking):
@@ -217,17 +226,21 @@ def from_bytes(raw):
         root_unique_id=root_unique_id,
     )
     group.init_nccl_communicator(device=0)
+    group.abort()
 
     assert events == [
         ("init", 2, 0, root_unique_id),
+        ("split", 0, 0),
         ("broadcast", 0, 17),
         "sync",
+        "abort-blocking",
+        "abort-bootstrap",
     ]
 
 
 def test_process_group_waits_for_nonblocking_broadcast_before_reading(monkeypatch):
     events = []
-    statuses = iter((0, 7, 0))
+    statuses = iter((0, 7, 0, 0, 7, 0))
 
     class FakeTensor:
         def __init__(self, value):
@@ -240,13 +253,18 @@ def synchronize(self):
             events.append("sync")
 
     class FakeCommunicator:
-        def __init__(self):
+        def __init__(self, kind):
+            self.kind = kind
             self.recvbuf = None
 
         @staticmethod
         def init(**kwargs):
             events.append("init")
-            return FakeCommunicator()
+            return FakeCommunicator("bootstrap")
+
+        def split(self, **kwargs):
+            events.append("split")
+            return FakeCommunicator("blocking")
 
         def get_async_error(self):
             status = next(statuses)
@@ -304,6 +322,11 @@ def from_bytes(raw):
     assert events == [
         "init",
         ("status", 0),
+        "split",
+        ("status", 7),
+        "sleep",
+        ("status", 0),
+        ("status", 0),
         "broadcast",
         ("status", 7),
         "sleep",

From 7c2cc385eb212e0e1d019ede25349d629d61003f Mon Sep 17 00:00:00 2001
From: Yixin Huang <yixinh@nvidia.com>
Date: Thu, 20 Aug 2026 14:39:56 -0700
Subject: [PATCH 11/11] fix(grpo): shut down Ray owners before exit

Signed-off-by: Yixin Huang <yixinh@nvidia.com>
---
 examples/run_grpo.py                | 51 ++++++++++++++++++++++++++---
 tests/unit/test_run_grpo_cleanup.py | 51 +++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 4 deletions(-)
 create mode 100644 tests/unit/test_run_grpo_cleanup.py

diff --git a/examples/run_grpo.py b/examples/run_grpo.py
index 51b1f08698e..19a23ea86fc 100644
--- a/examples/run_grpo.py
+++ b/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__":
diff --git a/tests/unit/test_run_grpo_cleanup.py b/tests/unit/test_run_grpo_cleanup.py
new file mode 100644
index 00000000000..66f1eaa60e7
--- /dev/null
+++ b/tests/unit/test_run_grpo_cleanup.py
@@ -0,0 +1,51 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 examples import run_grpo
+
+
+def test_shutdown_runtime_drains_unique_owners_before_ray(monkeypatch):
+    events = []
+
+    class Resource:
+        def __init__(self, name):
+            self.name = name
+
+        def shutdown(self):
+            events.append(self.name)
+
+    generation = Resource("generation")
+    teacher = Resource("teacher")
+    policy = Resource("policy")
+    train_cluster = Resource("train_cluster")
+    inference_cluster = Resource("inference_cluster")
+
+    monkeypatch.setattr(run_grpo.ray, "is_initialized", lambda: True)
+    monkeypatch.setattr(run_grpo.ray, "shutdown", lambda: events.append("ray_shutdown"))
+
+    run_grpo._shutdown_runtime(
+        policy,
+        generation,
+        (train_cluster, inference_cluster, train_cluster),
+        {"teacher": teacher, "teacher_alias": teacher},
+    )
+
+    assert events == [
+        "generation",
+        "teacher",
+        "policy",
+        "train_cluster",
+        "inference_cluster",
+        "ray_shutdown",
+    ]