diff --git a/docs/design-docs/nccl-reshard-refit.md b/docs/design-docs/nccl-reshard-refit.md index b208533572..973db27992 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 .yaml \ + policy.generation.colocated.enabled=false \ + policy.generation.refit_transport=mx_nccl_reshard \ + policy.generation.mx_server_url=: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/examples/run_grpo.py b/examples/run_grpo.py index 51b1f08698..19a23ea86f 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/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index d8d5f26763..0e9b2e47a8 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 f6786d1e8e..a8fc147a3e 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 e4e9891105..9e86071ab1 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -284,6 +284,56 @@ def init_nccl_reshard_comm_group( group.init_nccl_communicator(device=self.device) self.pp_comm_groups[stage] = group + def init_mx_reshard_comm_group( + self, + rank_prefix: int, + mx_server_url: str, + model_name: str, + trainer_slots: list, + generator_slots: list, + source_partition_count: int, + plan_digest: str = "", + phase: object = "all", + ) -> None: + """Bootstrap this gen worker's comm groups through ModelExpress. + + Phased for the same reason the trainer side is: one lane at a time, + cluster-wide, with the driver barriering between them. + """ + from nemo_rl.weight_sync.mx_collective_bootstrap import ( + mx_init_lane, + mx_rendezvous, + ) + + if phase == "rendezvous": + import uuid + + local_rank = torch.distributed.get_rank() + index_in_role = rank_prefix + local_rank + if not hasattr(self, "_mx_worker_id"): + self._mx_worker_id = ( # pyrefly: ignore[implicitly-defined-attribute] + f"gen-{index_in_role}-{uuid.uuid4().hex}" + ) + self._mx_state = mx_rendezvous( + mx_server_url=mx_server_url, + model_name=model_name, + role="GENERATOR", + index_in_role=index_in_role, + slot_id=f"gen/{index_in_role}", + worker_id=self._mx_worker_id, + trainer_slots=trainer_slots, + generator_slots=generator_slots, + source_partition_count=source_partition_count, + plan_digest=plan_digest, + device=self.device, + ) + return + if phase == "finish": + self.pp_comm_groups = self._mx_state.reshard_groups + self.model_update_group = self._mx_state.broadcast_group + return + mx_init_lane(self._mx_state, int(phase)) + def report_device_id(self) -> str: """Retrieve the UUID of the current CUDA device.""" from nemo_rl.utils.nvml import get_device_uuid diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index a8cb78bd13..599421933c 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1143,6 +1143,47 @@ def init_nccl_reshard_comm_group( # co-works with lm_policy; wait for all futures to complete outside return futures + def init_mx_reshard_comm_group( + self, + mx_server_url: str, + model_name: str, + trainer_slots: list, + generator_slots: list, + source_partition_count: int, + plan_digest: str = "", + phase: object = "all", + ) -> list[ray.ObjectRef]: + """Initialize the ModelExpress-brokered comm groups on all gen workers.""" + if not self.worker_group or not self.worker_group.workers: + raise RuntimeError("Worker group is not initialized") + + method_name = ( + "init_mx_reshard_comm_group_async" + if self.cfg["vllm_cfg"]["async_engine"] + else "init_mx_reshard_comm_group" + ) + + total_workers = len(self.worker_group.workers) + workers_per_group = total_workers // self.dp_size + rank_prefix_list = list(range(0, total_workers, workers_per_group)) + + futures = self.worker_group.run_all_workers_multiple_data( + method_name, + rank_prefix=rank_prefix_list, + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + common_kwargs={ + "mx_server_url": mx_server_url, + "model_name": model_name, + "trainer_slots": trainer_slots, + "generator_slots": generator_slots, + "source_partition_count": source_partition_count, + "plan_digest": plan_digest, + "phase": phase, + }, + ) + # co-works with lm_policy; wait for all futures to complete outside + return futures + def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: """Forward per-layer param metadata to vLLM workers for nccl_reshard refit.""" method_name = ( diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index ed3f045274..892bc1bf0e 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -1231,6 +1231,32 @@ def init_nccl_reshard_comm_group( ), ) + def init_mx_reshard_comm_group( + self, + rank_prefix: int, + mx_server_url: str, + model_name: str, + trainer_slots: list, + generator_slots: list, + source_partition_count: int, + plan_digest: str = "", + phase: object = "all", + ) -> None: + """Forward ModelExpress comm-group init to vLLM backend workers.""" + self.llm.collective_rpc( + "init_mx_reshard_comm_group", + args=( + rank_prefix, + mx_server_url, + model_name, + trainer_slots, + generator_slots, + source_partition_count, + plan_digest, + phase, + ), + ) + def prepare_nccl_reshard_refit_info(self, refit_info: dict) -> None: """Forward refit info to vLLM backend workers.""" self.llm.collective_rpc("prepare_nccl_reshard_refit_info", args=(refit_info,)) diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 81438afa3b..03d1f5e1c2 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -1152,6 +1152,34 @@ def init_nccl_reshard_comm_group( # co-works with vllm; wait for all futures to complete outside return futures + def init_mx_reshard_comm_group( + self, + mx_server_url: str, + model_name: str, + trainer_slots: list, + generator_slots: list, + source_partition_count: int, + pp_stages: list, + plan_digest: str = "", + phase: object = "all", + ) -> list[ray.ObjectRef]: + """Initialize the ModelExpress-brokered comm groups on all train workers.""" + futures = self.worker_group.run_all_workers_multiple_data( + "init_mx_reshard_comm_group", + my_pp_stage=pp_stages, + common_kwargs={ + "mx_server_url": mx_server_url, + "model_name": model_name, + "trainer_slots": trainer_slots, + "generator_slots": generator_slots, + "source_partition_count": source_partition_count, + "plan_digest": plan_digest, + "phase": phase, + }, + ) + # co-works with vllm; wait for all futures to complete outside + return futures + def prepare_nccl_reshard_refit_info( self, train_parallelism, diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py index e371e4f5ca..1667ee11dc 100644 --- a/nemo_rl/models/policy/workers/base_policy_worker.py +++ b/nemo_rl/models/policy/workers/base_policy_worker.py @@ -83,6 +83,65 @@ def init_nccl_reshard_comm_group( self.pp_comm_group.init_nccl_communicator(device=device) self.my_pp_stage = my_pp_stage + def init_mx_reshard_comm_group( + self, + mx_server_url: str, + model_name: str, + trainer_slots: list, + generator_slots: list, + source_partition_count: int, + my_pp_stage: int = 0, + plan_digest: str = "", + phase: object = "all", + ) -> None: + """Bootstrap this train worker's comm groups through ModelExpress. + + Same groups the TCPStore path produces, same downstream refit loop; only + the ncclUniqueId's provenance differs. + + Driven in phases: ``"rendezvous"``, then one call per lane id, then + ``"finish"``. The driver barriers between them because creating two + different communicators concurrently across overlapping rank sets + deadlocks NCCL -- the same reason the TCPStore path separates its + all-ranks group from its per-stage groups. + """ + from nemo_rl.weight_sync.mx_collective_bootstrap import ( + mx_init_lane, + mx_rendezvous, + ) + + if phase == "rendezvous": + import uuid + + if not hasattr(self, "_mx_worker_id"): + self._mx_worker_id = ( # pyrefly: ignore[implicitly-defined-attribute] + f"train-{self.rank}-{uuid.uuid4().hex}" + ) + self._mx_state = mx_rendezvous( + mx_server_url=mx_server_url, + model_name=model_name, + role="TRAINER", + index_in_role=self.rank, + slot_id=f"train/{self.rank}", + worker_id=self._mx_worker_id, + trainer_slots=trainer_slots, + generator_slots=generator_slots, + source_partition_count=source_partition_count, + source_partition=my_pp_stage, + plan_digest=plan_digest, + device=torch.cuda.current_device(), + ) + return + if phase == "finish": + state = self._mx_state + # A trainer belongs to exactly one source partition, so it holds one + # reshard lane; generators hold all of them. + self.pp_comm_group = state.reshard_groups[my_pp_stage] + self.model_update_group = state.broadcast_group + self.my_pp_stage = my_pp_stage + return + mx_init_lane(self._mx_state, int(phase)) + def prepare_nccl_reshard_refit_info( self, train_parallelism: dict[str, int], diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py index c48ed19f5d..8312132a9e 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_bootstrap.py b/nemo_rl/weight_sync/mx_collective_bootstrap.py new file mode 100644 index 0000000000..41a440eeac --- /dev/null +++ b/nemo_rl/weight_sync/mx_collective_bootstrap.py @@ -0,0 +1,433 @@ +# 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 math +import os +import time +from typing import Any + +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, + root_unique_id: Any = None, + timeout_s: float = 900.0, + ): + 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 + # it from the bytes MX brokered. + self._root_unique_id = root_unique_id + self._timeout_s = timeout_s + + 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 + 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. + """ + communicator = communicator or self.nccl_communicator + deadline = time.monotonic() + self._timeout_s + while True: + 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"{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 + 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): + unique_id = self._root_unique_id + if unique_id is None: + unique_id = UniqueId.from_bytes(self._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), + ) + 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() + 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. + 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) + # 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( + "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.""" + communicators = ( + self.nccl_communicator, + self._bootstrap_communicator, + ) + self.nccl_communicator = 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): + if stream is None: + stream = torch.cuda.current_stream() + self.nccl_communicator.broadcast( + sendbuf=tensor, recvbuf=tensor, root=src, stream=int(stream.cuda_stream) + ) + + +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 + self.rz = None + 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 + 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( + *, + 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, +) -> MxBootstrapState: + """Join the MX group and fetch every lane's identifier. + + 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" + ) + 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 + + 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, + ) + + # 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() + 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 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 {state.timeout_s:.0f}s" + ) + time.sleep(0.2) + + +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 + 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=minted_uid, + ) + 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" + ) + 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, + timeout_s=state.timeout_s, + ) + try: + pg.init_nccl_communicator(device=state.device) + except BaseException: + state.abort() + raise + if lane.kind == "BROADCAST": + state.broadcast_group = pg + else: + state.reshard_groups[lane.lane_id] = pg + + +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_plan.py b/nemo_rl/weight_sync/mx_collective_plan.py new file mode 100644 index 0000000000..a56510f60b --- /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 0000000000..dabe95f4ec --- /dev/null +++ b/nemo_rl/weight_sync/mx_collective_weight_synchronizer.py @@ -0,0 +1,187 @@ +# 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. + +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 +from typing import Any, Optional + +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" + + +class MxCollectiveWeightSynchronizer(WeightSynchronizer): + """nccl_reshard refit whose communicator bootstrap is brokered by ModelExpress.""" + + 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 + 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: + cfg = self._policy.cfg["megatron_cfg"] + return { + "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: + cfg = self._policy.cfg["generation"].get("vllm_cfg", {}) + return { + "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: + train_parallelism = self._train_parallelism() + gen_parallelism = self._gen_parallelism() + train_world_size = self._train_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, gen_world_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 - 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)] + 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, + ) + + common = dict( + mx_server_url=self._mx_server_url, + model_name=model_name, + trainer_slots=trainers, + generator_slots=generators, + source_partition_count=pp_size, + plan_digest=digest, + ) + + 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) + + def sync_weights( + self, + *, + timer: Optional[Timer] = None, + kv_scales: Optional[dict] = None, + ) -> None: + ctx = ( + timer.time("prepare_for_generation/transfer_and_update_weights") + if timer is not None + else nullcontext() + ) + 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) + if not all(r for r in results if r is not None): + raise RuntimeError( + "Weight transfer failed during the ModelExpress collective refit." + ) + self._stale = False + + @property + def is_stale(self) -> bool: + return self._stale + + def mark_stale(self) -> None: + self._stale = True + + def shutdown(self) -> None: + pass diff --git a/tests/unit/test_run_grpo_cleanup.py b/tests/unit/test_run_grpo_cleanup.py new file mode 100644 index 0000000000..66f1eaa60e --- /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", + ] 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 0000000000..240aa721da --- /dev/null +++ b/tests/unit/weight_sync/test_mx_collective_bootstrap.py @@ -0,0 +1,389 @@ +# 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 = [] + self.closed = False + + def publish_bootstrap(self, **kwargs): + self.published.append(kwargs) + + def close(self): + self.closed = True + + +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() + 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 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): + 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: + 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("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(f"abort-{self.kind}") + + class FakeConfig: + def __init__(self, *, blocking): + self.blocking = blocking + + 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 + 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 + ) + + 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) + 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, 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, kind): + self.kind = kind + self.recvbuf = None + + @staticmethod + def init(**kwargs): + events.append("init") + return FakeCommunicator("bootstrap") + + def split(self, **kwargs): + events.append("split") + return FakeCommunicator("blocking") + + 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), + "split", + ("status", 7), + "sleep", + ("status", 0), + ("status", 0), + "broadcast", + ("status", 7), + "sleep", + ("status", 0), + "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"] 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 0000000000..162060ba82 --- /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)