From 6b6a6c5c98e4a23c77bb07d344c515f48c157193 Mon Sep 17 00:00:00 2001 From: pkousha <43781676+pkousha@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:41:17 -0700 Subject: [PATCH] [RL] Support per-parameter source layouts for NCCL M2N (#60934) Signed-off-by: pkousha Co-authored-by: OpenAI Codex Co-authored-by: aoshen02 --- docs/training/weight_transfer/m2n.md | 31 +- examples/rl/rlhf_m2n.py | 35 +- tests/distributed/test_weight_transfer_m2n.py | 1223 ++++++++++++++--- .../distributed/weight_transfer/m2n_common.py | 766 +++++++++-- .../distributed/weight_transfer/m2n_engine.py | 432 ++++-- .../distributed/weight_transfer/m2n_source.py | 124 +- .../weight_transfer/m2n_trainer.py | 619 ++++++--- .../weight_transfer/nccl_common.py | 55 +- 8 files changed, 2590 insertions(+), 695 deletions(-) diff --git a/docs/training/weight_transfer/m2n.md b/docs/training/weight_transfer/m2n.md index 61510978c8..a0e4e8089b 100644 --- a/docs/training/weight_transfer/m2n.md +++ b/docs/training/weight_transfer/m2n.md @@ -10,18 +10,14 @@ The `nccl_m2n` weight transfer engine uses [NCCL M2N](https://github.com/NVIDIA/ ## Requirements -This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. It needs NCCL 2.30.5 or newer, and `libnccl_m2n.so` is linked against a specific NCCL build, so vLLM must load the same library: - -```bash -export VLLM_NCCL_SO_PATH=/path/to/the/same/libnccl.so -``` +This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. Install both against a compatible NCCL 2.30.5 or newer runtime. The backend creates its data-plane communicator through `nccl4py` and passes that communicator directly to `nccl.m2n`. The backend imports the runtime lazily, so vLLM is unaffected unless you select it. ## How It Works -1. The trainer ranks and all inference workers join **one** NCCL communicator via `StatelessProcessGroup`. Trainer ranks occupy `[0, T)` and workers `[T, T + N)` — m2n meshes are contiguous rank intervals, which is exactly what vLLM's existing `rank_offset` convention produces. -2. At the init handshake the trainer ships the full transfer plan: both meshes — one per side, shared by every parameter — then each parameter's name, dtype, shape and placement. The workers resolve the matching **destination placements** and publish them back over the shared group; the meshes themselves are declared, not negotiated. +1. A `StatelessProcessGroup` exchanges metadata and the NCCL unique ID. The trainer ranks and all inference workers then create **one** `nccl.core.Communicator`: trainer ranks occupy `[0, T)` and workers `[T, T + N)`. The communicator is passed directly to each M2N collective. +2. At the init handshake the trainer ships the full transfer plan: the session-wide destination mesh, then each parameter's name, dtype, shape, source mesh, and source placements. The workers resolve the matching **destination placements** and publish them over the data-plane communicator; the meshes themselves are declared, not negotiated. 3. Each round, both sides issue one `reshard` per parameter in the same order. Every rank in the communicator participates in every reshard, so — as with the NCCL engine — the worker must be inside `update_weights` while the trainer sends. The trainer engine handles that concurrency itself. ### Destination resolution @@ -83,19 +79,28 @@ engine.send_weights() `send_weights()` drives the whole round — `start_weight_update`, `update_weights`, `finish_weight_update`. Every rank in the communicator must be inside the same reshard, so `update_weights` has to be in flight while the trainer sends; the engine runs it on a helper thread so the trainer loop does not have to open-code that. -Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass `M2NWeightSource` and supply layouts explicitly. +Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass +`M2NWeightSource` and supply layouts explicitly. A source may yield any +nccl4py-compatible CUDA buffer: a torch tensor, a CUDA Array Interface object, +or a raw device pointer. vLLM keeps buffer owners alive through stream +completion, using allocator stream tracking when the buffer provides it. +Sources must provide buffers on the communicator's current device; nccl4py +resolves their pointers but does not verify device ownership. Raw-pointer +sources are additionally responsible for pointer lifetime and matching their +declared dtype and local shape. See [`examples/rl/rlhf_m2n.py`](../../../examples/rl/rlhf_m2n.py) for a runnable FSDP → TP example. ## Limitations -- Tensor rank 1..3; 4-D parameters are rejected. -- dtypes: int8/uint8, fp8 e4m3/e5m2, fp16, bf16, int32/uint32, fp32, int64/uint64, fp64. No fp4. -- The trainer's device mesh must cover a contiguous rank interval starting at 0, and be 1-D or 2-D. +- vLLM validates tensor-rank and dtype support while agreeing on the complete + plan; the installed M2N runtime validates them again for each reshard. +- Each parameter's source mesh must cover the same contiguous trainer-rank interval starting at 0, and be 1-D or 2-D. - `Partial` placements are not supported. - Models with known model-level checkpoint preprocessing use the full-tensor fallback for the affected module. - A reshard cannot be issued from inside a CUDA graph capture (weight updates run outside capture, so this only matters for custom callers). -- The reshard plan is capped per *shard*, not per mesh: at most 16 source shards may feed one destination shard, and one source shard may feed at most 64 destination shards (`MAX_SOURCES` / `MAX_TARGETS` in the m2n build; raising either needs a rebuild). -- A directly-resharded parameter spreads over `tp_size` destination shards, so each one is fed by fewer source shards. A parameter on the fallback path is replicated, i.e. a single destination shard fed by every source shard, so a trainer that **sharded** its parameters is still limited to 16 ranks whenever any parameter falls back — which fused parameters always do. A replicated trainer is unaffected. +- vLLM checks every parameter's shard overlap against conservative 16-source + and 64-target capacities before committing the plan. The installed M2N + runtime remains authoritative and validates each reshard again. - Each communicator holds a cached staging pool (~256 MiB by default: 4 channels x 64 MiB). diff --git a/examples/rl/rlhf_m2n.py b/examples/rl/rlhf_m2n.py index e02229ed09..6fc64463e1 100644 --- a/examples/rl/rlhf_m2n.py +++ b/examples/rl/rlhf_m2n.py @@ -17,9 +17,8 @@ Every FSDP rank builds an ``M2NTrainerWeightTransferEngine`` and calls ``send_weights()``; all ranks run every reshard, and only rank 0 drives the inference side through its ``RayVLLMWeightSyncClient``. -Requires the ``nccl-extensions`` package (NCCL 2.30.5+) and a -``VLLM_NCCL_SO_PATH`` pointing at the same ``libnccl.so`` that -``libnccl_m2n.so`` was linked against. +Requires the ``nccl-extensions`` package (NCCL 2.30.5+). The backend creates +the data-plane communicator with ``nccl4py`` and passes it directly to M2N. This example was written for 4xH100. """ @@ -28,6 +27,7 @@ from __future__ import annotations import os +import nccl.core as nccl import pybase64 as base64 import ray import torch @@ -38,7 +38,6 @@ from transformers import AutoModelForCausalLM from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary from vllm.distributed.weight_transfer import ( RayVLLMWeightSyncClient, WeightTransferTrainerFactory, @@ -94,7 +93,15 @@ class FSDPTrainWorker: def ready(self): return True - def setup_engine(self, llm_handle, nccl_unique_id_b64, world_size, num_workers): + def setup_engine( + self, + llm_handle, + metadata_address, + metadata_port, + nccl_unique_id_b64, + world_size, + num_workers, + ): """Build the trainer engine on every FSDP rank. `DTensorModuleSource` reads each parameter's FSDP device mesh and @@ -103,6 +110,8 @@ class FSDPTrainWorker: """ self.engine = WeightTransferTrainerFactory.trainer_init( init_info=M2NTrainerInitInfo( + master_address=metadata_address, + master_port=metadata_port, nccl_unique_id_b64=nccl_unique_id_b64, world_size=world_size, num_trainer_ranks=FSDP_WORLD_SIZE, @@ -166,16 +175,22 @@ def main(): # Trainer ranks [0, FSDP_WORLD_SIZE) and inference ranks after them share # one communicator, so the two meshes are contiguous rank intervals. - nccl = NCCLLibrary() - nccl_unique_id_b64 = base64.b64encode( - bytes(nccl.ncclGetUniqueId().internal) - ).decode() + nccl_unique_id_b64 = base64.b64encode(nccl.get_unique_id().as_bytes).decode() world_size = FSDP_WORLD_SIZE + num_workers + metadata_address = get_ip() + metadata_port = get_open_port() print("[transfer] Initializing nccl_m2n weight transfer (all FSDP ranks)...") ray.get( [ - w.setup_engine.remote(llm, nccl_unique_id_b64, world_size, num_workers) + w.setup_engine.remote( + llm, + metadata_address, + metadata_port, + nccl_unique_id_b64, + world_size, + num_workers, + ) for w in fsdp_workers ] ) diff --git a/tests/distributed/test_weight_transfer_m2n.py b/tests/distributed/test_weight_transfer_m2n.py index 588ffeeb26..991852c034 100644 --- a/tests/distributed/test_weight_transfer_m2n.py +++ b/tests/distributed/test_weight_transfer_m2n.py @@ -2,36 +2,44 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for the NCCL M2N weight transfer backend. -These cover how a transfer is *described* — layout encoding and the wire-type -validation that keeps a bad plan from reaching a collective, since a mismatch -there is a hang rather than an error. The transfer itself needs the m2n runtime -and multiple GPUs, so it is exercised separately. +These cover how a transfer is *described* — layout encoding, whole-plan +validation, and the vLLM-owned source contract. The transfer itself needs the +M2N runtime and multiple GPUs, so it is exercised separately. """ import logging +from contextlib import contextmanager from functools import wraps -from types import MethodType -from unittest.mock import Mock +from types import MethodType, SimpleNamespace +from unittest.mock import Mock, call import pybase64 as base64 import pytest import ray import torch +from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer import ( WeightTransferEngineFactory, WeightTransferTrainerFactory, ) from vllm.distributed.weight_transfer.m2n_common import ( + M2N_WIRE_SCHEMA_VERSION, REPLICATE, REPLICATED, + M2NLayout, M2NMesh, M2NParamMeta, - check_placements, - check_transferable, - publish_destination_placements, + M2NWireParam, + check_complete_plan_agreement, + check_data_plane_agreement, + check_plan_limits, + check_runtime_ready_agreement, + check_source_plan_agreement, + create_m2n_communicator, + expected_local_shape, resolve_layout, - validate_layout, + source_plan_digest, ) from vllm.distributed.weight_transfer.m2n_engine import ( M2NWeightTransferEngine, @@ -43,6 +51,9 @@ from vllm.distributed.weight_transfer.m2n_layout import ( resolve_parameter_destinations, ) from vllm.distributed.weight_transfer.m2n_source import ( + M2NManifestEntry, + M2NWeightSource, + ManifestM2NWeightSource, mesh_from_tensor, placements_from_tensor, ) @@ -50,14 +61,16 @@ from vllm.distributed.weight_transfer.m2n_trainer import ( M2NTrainerInitInfo, M2NTrainerWeightTransferEngine, ) +from vllm.distributed.weight_transfer.nccl_common import ( + stateless_init_metadata_group, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.platforms import current_platform - -VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() +from vllm.utils.network_utils import get_open_port @pytest.mark.parametrize( @@ -67,18 +80,114 @@ VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() (M2NTrainerWeightTransferEngine, "group"), ], ) -def test_shutdown_destroys_communicator_once(engine_type, group_attribute): +def test_shutdown_releases_resources_once(monkeypatch, engine_type, group_attribute): engine = object.__new__(engine_type) - engine._handle = None - engine._executor = None + handle = Mock() + engine._handle = handle group = Mock() setattr(engine, group_attribute, group) + synchronize = Mock() + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) + + if engine_type is M2NWeightTransferEngine: + engine._parameter_destinations = [Mock()] + else: + engine._dst_placements = [REPLICATED] + engine._metadata_group = Mock() engine.shutdown() engine.shutdown() - group.destroy.assert_called_once_with() + synchronize.assert_called_once_with() + handle.destroy.assert_called_once_with() + group.abort.assert_called_once_with() + group.destroy.assert_not_called() + assert engine._handle is None assert getattr(engine, group_attribute) is None + if engine_type is M2NWeightTransferEngine: + assert engine._parameter_destinations == [] + else: + assert engine._dst_placements == [] + assert engine._metadata_group is None + + +@pytest.mark.parametrize( + "engine_type,group_attribute", + [ + (M2NWeightTransferEngine, "model_update_group"), + (M2NTrainerWeightTransferEngine, "group"), + ], +) +@pytest.mark.parametrize("failure", ["synchronize", "handle"]) +def test_shutdown_attempts_all_cleanup_after_failure( + monkeypatch, engine_type, group_attribute, failure +): + engine = object.__new__(engine_type) + handle = Mock() + group = Mock() + engine._handle = handle + setattr(engine, group_attribute, group) + if engine_type is M2NTrainerWeightTransferEngine: + engine._metadata_group = Mock() + + synchronize = Mock() + if failure == "synchronize": + synchronize.side_effect = RuntimeError("synchronize failed") + else: + handle.destroy.side_effect = RuntimeError("handle destroy failed") + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) + + with pytest.raises(RuntimeError, match=failure): + engine.shutdown() + + synchronize.assert_called_once_with() + handle.destroy.assert_called_once_with() + group.abort.assert_called_once_with() + assert engine._handle is None + assert getattr(engine, group_attribute) is None + if engine_type is M2NTrainerWeightTransferEngine: + assert engine._metadata_group is None + + engine.shutdown() + synchronize.assert_called_once_with() + handle.destroy.assert_called_once_with() + group.abort.assert_called_once_with() + + +class _BufferLike: + """Minimal nccl4py buffer protocol without torch-only methods.""" + + def __init__(self, dtype=torch.float32, shape=(4,)): + self.dtype = dtype + self.shape = shape + + def data_ptr(self): + return 1234 + + +class _AllocatorAwareBuffer(_BufferLike): + def __init__(self): + super().__init__() + self.record_stream = Mock() + + +class _FakeDistTensor: + def __init__(self, buffer, local_shape=None, dtype=None, **kwargs): + del kwargs + self.local_shape = ( + tuple(local_shape) if local_shape is not None else tuple(buffer.shape) + ) + self.dtype = dtype if dtype is not None else buffer.dtype + + +def _fake_m2n(): + return SimpleNamespace( + DistTensor=_FakeDistTensor, + Mesh=lambda dims, start_rank=0: (tuple(dims), start_rank), + Replicate=lambda: REPLICATE, + Shard=lambda dim: dim, + reshard=Mock(), + ) class TestLayout: @@ -87,10 +196,10 @@ class TestLayout: dims are deliberately coprime with the mesh size: a layout that actually split the tensor would reject them.""" mesh = M2NMesh((2, 2), start_rank=1) - indivisible_shape = (7, 13) # neither dim divisible by any mesh axis resolved_mesh, placements = resolve_layout(mesh, REPLICATED) - validate_layout(resolved_mesh, placements, indivisible_shape, "destination") + assert resolved_mesh.dims == (4, 1) + assert placements == (REPLICATE, 0) def test_replicated_keeps_the_same_ranks(self): """Replication re-factors the mesh to get a size-1 axis for its no-op @@ -108,23 +217,27 @@ class TestLayout: assert resolved_mesh == mesh assert placements == (REPLICATE, 0) - def test_two_shard_axes_rejected(self): - """One axis has to replicate; a 2-D mesh that shards both is not - something a single reshard can express.""" - with pytest.raises(ValueError, match="shards both"): - check_placements((0, 1)) + @pytest.mark.parametrize( + ("placements", "match"), + [((REPLICATE, 1), "tensor dim 1"), ((REPLICATE, 0), "divide evenly")], + ) + def test_layout_validation_always_returns_shape_or_raises(self, placements, match): + meta = M2NParamMeta( + "w", + torch.float32, + (5,), + M2NLayout(M2NMesh((1, 2), 0), placements), + ) - def test_placement_below_replicate_rejected(self): - with pytest.raises(ValueError, match=r"parameter 'w'.*-2"): - check_placements((-2, REPLICATE), "parameter 'w' source placements") + with pytest.raises(ValueError, match=match): + expected_local_shape(meta) - def test_shard_dim_must_exist(self): - with pytest.raises(ValueError, match="rank 2"): - validate_layout(M2NMesh((1, 2), 0), (REPLICATE, 2), (8, 16), "source") + def test_cross_sharded_fanout_over_capacity_is_rejected(self): + src = (M2NMesh((1, 2), 0), (REPLICATE, 0)) + dst = (M2NMesh((1, 65), 2), (REPLICATE, 1)) - def test_shard_must_divide_evenly(self): - with pytest.raises(ValueError, match="does not divide evenly"): - validate_layout(M2NMesh((1, 3), 0), (REPLICATE, 0), (8, 16), "source") + with pytest.raises(ValueError, match=r"parameter 'w'.*MAX_TARGETS=64"): + check_plan_limits(src, dst, "w", (2, 65)) class TestSourceLayout: @@ -134,76 +247,398 @@ class TestSourceLayout: assert placements_from_tensor(torch.zeros(4)) is REPLICATED assert mesh_from_tensor(torch.zeros(4), 4) == M2NMesh((4, 1), 0) + def test_manifest_preserves_order_and_materializes_lazily(self): + calls: list[int] = [] -class TestTransferable: - def test_unsupported_dtype_names_the_parameter(self): - with pytest.raises(ValueError, match="'w'"): - check_transferable("w", torch.complex64, (4,)) + def provider(index, value): + def get_tensor(): + calls.append(index) + return value - def test_rank_four_rejected(self): - with pytest.raises(ValueError, match="rank 4"): - check_transferable("w", torch.bfloat16, (2, 2, 2, 2)) + return get_tensor + + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "second", + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + provider(0, torch.ones(4)), + ), + M2NManifestEntry( + "first", + torch.float32, + (2,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + provider(1, torch.full((2,), 2.0)), + ), + ] + ) + + assert [meta.name for meta in source.metadata()] == ["second", "first"] + assert calls == [] + assert [name for name, _ in source] == ["second", "first"] + assert calls == [0, 1] + + def test_trainer_boundary_rejects_duplicate_manifest_names(self): + entry = M2NManifestEntry( + "w", + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + lambda: torch.zeros(4), + ) + source = ManifestM2NWeightSource([entry, entry]) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match="duplicate parameter 'w'"): + engine._prepare_source_plan(source, num_trainer_ranks=1) + + @pytest.mark.parametrize( + ("dtype", "shape"), + [(torch.complex64, (4,)), (torch.float32, (2, 2, 2, 2))], + ) + def test_transport_capabilities_rejected_during_initialization(self, dtype, shape): + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "w", + dtype, + shape, + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + lambda: _BufferLike(dtype, shape), + ) + ] + ) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match=r"parameter 'w'.*(dtype|rank)"): + engine._prepare_source_plan(source, num_trainer_ranks=1) + + @pytest.mark.parametrize("placements", [(REPLICATE, 1), (REPLICATE, 0)]) + def test_layout_geometry_rejected_during_initialization(self, placements): + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "w", + torch.float32, + (5,), + M2NLayout(M2NMesh((1, 2), 0), placements), + lambda: _BufferLike(shape=(5,)), + ) + ] + ) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match=r"parameter 'w'.*source layout"): + engine._prepare_source_plan(source, num_trainer_ranks=2) + + def test_raw_pointer_invalid_layout_rejected_during_initialization(self): + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "pointer", + torch.float32, + (5,), + M2NLayout(M2NMesh((1, 2), 0), (REPLICATE, 0)), + lambda: 1234, + ) + ] + ) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match=r"parameter 'pointer'.*divide evenly"): + engine._prepare_source_plan(source, num_trainer_ranks=2) + + +class TestPreflightAgreement: + def test_source_error_is_reported_before_data_plane(self): + group = Mock(world_size=2) + group.all_gather_obj.return_value = [ + { + "phase": "source", + "digest": None, + "declared_digest": None, + "error": "ValueError: bad local manifest", + }, + { + "phase": "source", + "digest": None, + "declared_digest": None, + "error": "ValueError: bad local manifest", + }, + ] + + with pytest.raises(RuntimeError, match="bad local manifest"): + check_source_plan_agreement( + group, [], local_error="ValueError: bad local manifest" + ) + + def test_data_plane_identity_must_match(self): + group = Mock(world_size=2) + + def gather(envelope): + peer = dict(envelope) + peer["mode"] = "uid" + peer["uid_digest"] = "different" + return [envelope, peer] + + group.all_gather_obj.side_effect = gather + with pytest.raises(RuntimeError, match="data-plane identity disagrees"): + check_data_plane_agreement(group, None) + + def test_runtime_error_is_shared_before_nccl4py(self): + group = Mock(world_size=2) + group.all_gather_obj.return_value = [ + {"phase": "runtime_ready", "error": None}, + {"phase": "runtime_ready", "error": "ImportError: no nccl_m2n"}, + ] + + with pytest.raises(RuntimeError, match="no nccl_m2n"): + check_runtime_ready_agreement(group, None) + + def test_complete_plan_error_is_shared_before_first_reshard(self): + group = Mock(world_size=2) + group.all_gather_obj.return_value = [ + {"phase": "complete_plan", "digest": "same", "error": None}, + { + "phase": "complete_plan", + "digest": None, + "error": "ValueError: bad destination", + }, + ] + + with pytest.raises(RuntimeError, match="bad destination"): + check_complete_plan_agreement( + group, + "source", + M2NMesh((1, 1), 1), + [REPLICATED], + ) + + def test_complete_plan_digest_must_match(self): + group = Mock(world_size=2) + + def gather(envelope): + peer = dict(envelope) + peer["digest"] = "different" + return [envelope, peer] + + group.all_gather_obj.side_effect = gather + with pytest.raises(RuntimeError, match="complete-plan digest disagrees"): + check_complete_plan_agreement( + group, + "source", + M2NMesh((1, 1), 1), + [REPLICATED], + ) class TestWireTypes: + @staticmethod + def _param( + name="w", + dtype_name="bfloat16", + shape=(16, 16), + src_mesh_dims=(1, 1), + src_placements=None, + ): + return M2NWireParam( + name=name, + dtype_name=dtype_name, + shape=shape, + src_mesh_dims=src_mesh_dims, + src_placements=src_placements, + ) + def _init_info(self, **overrides): + param = self._param() fields = dict( + schema_version=M2N_WIRE_SCHEMA_VERSION, master_address="127.0.0.1", master_port=1234, rank_offset=1, world_size=3, - src_mesh_dims=[1, 1], dst_mesh_dims=[2, 1], - names=["w"], - dtype_names=["bfloat16"], - shapes=[[16, 16]], - src_placements=[None], + source_digest=source_plan_digest([param]), + params=[param.to_dict()], ) fields.update(overrides) + if "params" in overrides and "source_digest" not in overrides: + params = [M2NWireParam.from_dict(value) for value in fields["params"]] + fields["source_digest"] = source_plan_digest(params) return M2NWeightTransferInitInfo(**fields) def test_accepts_a_consistent_plan(self): - assert self._init_info().names == ["w"] + assert [param.name for param in self._init_info().parse_wire_params()] == ["w"] - def test_accepts_pre_shared_nccl_unique_id(self): + def test_m2n_communicator_uses_nccl4py_uid_under_unpinned_env(self, monkeypatch): + nccl = Mock() + group = Mock(rank=1, world_size=3) + unique_id = object() + communicator = object() + nccl.UniqueId.from_bytes.return_value = unique_id + active = False + + @contextmanager + def unpinned(): + nonlocal active + active = True + try: + yield + finally: + active = False + + def init(**kwargs): + assert active + assert kwargs == {"nranks": 3, "rank": 1, "unique_id": unique_id} + return communicator + + nccl.Communicator.init.side_effect = init + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_common.unpinned_nccl_env", + unpinned, + ) + + result = create_m2n_communicator(nccl, group, b"shared-id") + + assert result is communicator + nccl.UniqueId.from_bytes.assert_called_once_with(b"shared-id") + nccl.Communicator.init.assert_called_once_with( + nranks=3, rank=1, unique_id=unique_id + ) + group.broadcast_obj.assert_not_called() + + @pytest.mark.parametrize( + "mutation", + [ + lambda value: value.pop("shape"), + lambda value: value.update(extra="field"), + lambda value: value.update(src_mesh_dims=[1]), + lambda value: value.update(src_placements=[0, 1]), + ], + ) + def test_malformed_wire_parameter_rejected(self, mutation): + value = self._param().to_dict() + mutation(value) + with pytest.raises((TypeError, ValueError)): + M2NWireParam.from_dict(value) + + def test_source_digest_authenticates_ordered_plan(self): + params = [self._param("a"), self._param("b")] info = self._init_info( - master_address=None, - master_port=None, - nccl_unique_id_b64=VALID_UID_B64, + params=[param.to_dict() for param in params], + source_digest=source_plan_digest(list(reversed(params))), ) - assert info.nccl_unique_id_bytes == b"\x00" * 128 - assert VALID_UID_B64 not in repr(info) + with pytest.raises(ValueError, match="source digest"): + info.parse_wire_params() - def test_destination_plan_uses_uid_communicator_without_group(self): - comm = Mock(rank=0, device=torch.device("cpu"), group=None) - - def receive_plan(tensor, src): - assert src == 1 - tensor.copy_(torch.tensor([[-1, 0], [-2, -2]], dtype=torch.int8)) - - comm.broadcast.side_effect = receive_plan - - placements = publish_destination_placements(comm, 1, None, 2) - - assert placements == [(REPLICATE, 0), REPLICATED] - - def test_rejects_both_rendezvous_modes(self): - with pytest.raises(ValueError, match="not both"): - self._init_info(nccl_unique_id_b64=VALID_UID_B64) - - def test_trainer_accepts_pre_shared_nccl_unique_id(self): - info = M2NTrainerInitInfo( - nccl_unique_id_b64=VALID_UID_B64, - world_size=3, - num_trainer_ranks=1, - rank=0, + @pytest.fixture + def worker_init_env(self, monkeypatch): + metadata_group = Mock(rank=1, world_size=3) + metadata_group.all_gather_obj.side_effect = lambda envelope: [envelope] * 3 + init_metadata = Mock(return_value=metadata_group) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.worker_init_metadata_group", + init_metadata, + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.import_m2n", Mock() + ) + handle = Mock() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.prepare_m2n_local_runtime", + Mock(return_value=(Mock(), 0, handle)), + ) + communicator = Mock() + make_communicator = Mock(return_value=communicator) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.create_m2n_communicator", + make_communicator, + ) + publish = Mock(return_value=[]) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "publish_destination_placements", + publish, + ) + vllm_config = SimpleNamespace( + parallel_config=Mock(pipeline_parallel_size=1), + model_config=SimpleNamespace(quantization=None), + ) + engine = M2NWeightTransferEngine( + WeightTransferConfig(backend="nccl_m2n"), + vllm_config, + torch.device("cuda:0"), + torch.nn.Module(), + ) + return SimpleNamespace( + engine=engine, + metadata_group=metadata_group, + init_metadata=init_metadata, + handle=handle, + communicator=communicator, + make_communicator=make_communicator, + publish=publish, ) - assert info.nccl_unique_id_bytes == b"\x00" * 128 - assert VALID_UID_B64 not in repr(info) - def test_ragged_plan_rejected(self): - with pytest.raises(ValueError, match="`shapes`"): - self._init_info(shapes=[]) + def test_trainer_disagreement_is_rejected_before_nccl4py(self, worker_init_env): + group = worker_init_env.metadata_group + + def gather(envelope): + peer = dict(envelope) + peer["digest"] = "different" + return [envelope, peer, envelope] + + group.all_gather_obj.side_effect = gather + with pytest.raises(RuntimeError, match="source digest disagrees"): + worker_init_env.engine.init_transfer_engine(self._init_info()) + + worker_init_env.make_communicator.assert_not_called() + + def test_worker_rejects_active_reinit_but_allows_init_after_shutdown( + self, monkeypatch, worker_init_env + ): + engine = worker_init_env.engine + old_handle = Mock() + old_group = Mock() + engine._handle = old_handle + engine.model_update_group = old_group + + with pytest.raises(RuntimeError, match="already initialized"): + engine.init_transfer_engine(self._init_info(params=[])) + worker_init_env.init_metadata.assert_not_called() + + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + engine.shutdown() + + engine.init_transfer_engine(self._init_info(params=[])) + + assert engine._handle is worker_init_env.handle + assert engine.model_update_group is worker_init_env.communicator + + def test_partial_worker_initialization_releases_transport_resources( + self, worker_init_env + ): + worker_init_env.publish.side_effect = RuntimeError("late init failure") + + with pytest.raises(RuntimeError, match="late init failure"): + worker_init_env.engine.init_transfer_engine(self._init_info(params=[])) + + worker_init_env.communicator.abort.assert_called_once_with() + worker_init_env.communicator.destroy.assert_not_called() + worker_init_env.handle.destroy.assert_called_once_with() + assert worker_init_env.engine.model_update_group is None + assert worker_init_env.engine._handle is None def test_destination_mesh_must_cover_the_workers(self): """The trainer declares the inference mesh, so one that does not cover @@ -217,69 +652,140 @@ class TestWireTypes: self._init_info(rank_offset=3, world_size=3) @pytest.mark.parametrize("dtype_name", ["not_a_dtype", "Tensor"]) - def test_invalid_dtype_name_names_parameter(self, monkeypatch, dtype_name): - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.import_m2n", - lambda: object(), - ) + def test_invalid_dtype_name_names_parameter(self, dtype_name): engine = object.__new__(M2NWeightTransferEngine) + param = self._param(dtype_name=dtype_name) with pytest.raises(ValueError, match=r"parameter 'w'.*dtype"): - engine.init_transfer_engine(self._init_info(dtype_names=[dtype_name])) + engine._prepare_source_plan((param,), num_trainer_ranks=1) def _init_32_to_4_plan(self, monkeypatch, destination): - m2n = Mock() - m2n.Handle.create.return_value = object() - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.import_m2n", lambda: m2n - ) monkeypatch.setattr( "vllm.distributed.weight_transfer.m2n_engine." "resolve_parameter_destinations", lambda *args, **kwargs: [destination], ) - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.worker_init_process_group", - lambda *args, **kwargs: object(), - ) - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine." - "publish_destination_placements", - lambda *args: args[2], - ) engine = object.__new__(M2NWeightTransferEngine) engine.parallel_config = Mock(pipeline_parallel_size=1) engine.model_config = Mock(quantization=None) engine.model = torch.nn.Module() - engine.init_transfer_engine( + param = self._param( + shape=(32, 8), + src_mesh_dims=(1, 32), + src_placements=(REPLICATE, 0), + ) + engine._prepare_source_plan((param,), num_trainer_ranks=32) + engine._prepare_destination_plan( self._init_info( rank_offset=32, world_size=36, - src_mesh_dims=[1, 32], dst_mesh_dims=[1, 4], - shapes=[[32, 8]], - src_placements=[[REPLICATE, 0]], + params=[param.to_dict()], ) ) - def test_limits_use_resolved_sharded_destination(self, monkeypatch): + def test_sharded_destination_within_capacity_initializes(self, monkeypatch): destination = M2NDestination("w", (REPLICATE, 0), torch.empty(8, 8)) self._init_32_to_4_plan(monkeypatch, destination) - def test_limits_still_reject_replicated_destination(self, monkeypatch): + def test_cross_sharded_destination_over_capacity_rejected_at_init( + self, monkeypatch + ): + destination = M2NDestination("w", (REPLICATE, 1), torch.empty(32, 2)) + + with pytest.raises(ValueError, match=r"parameter 'w'.*MAX_SOURCES=16"): + self._init_32_to_4_plan(monkeypatch, destination) + + def test_replicated_destination_over_capacity_rejected_at_init(self, monkeypatch): destination = M2NDestination("w", REPLICATED, None) - with pytest.raises(ValueError, match=r"32 source shards.*MAX_SOURCES=16"): + with pytest.raises(ValueError, match=r"parameter 'w'.*MAX_SOURCES=16"): self._init_32_to_4_plan(monkeypatch, destination) + def test_invalid_destination_placement_rejected_at_init(self, monkeypatch): + destination = M2NDestination("w", (REPLICATE, 2), torch.empty(8, 8)) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "resolve_parameter_destinations", + lambda *args, **kwargs: [destination], + ) + engine = object.__new__(M2NWeightTransferEngine) + engine.parallel_config = Mock(pipeline_parallel_size=1) + engine.model_config = Mock(quantization=None) + engine.model = torch.nn.Module() + param = self._param(shape=(32, 8)) + engine._prepare_source_plan((param,), num_trainer_ranks=1) + + with pytest.raises( + ValueError, match=r"parameter 'w'.*destination layout.*tensor dim 2" + ): + engine._prepare_destination_plan( + self._init_info( + rank_offset=1, + world_size=5, + dst_mesh_dims=[1, 4], + params=[param.to_dict()], + ) + ) + + def test_later_invalid_parameter_rejects_plan_before_any_reshard(self, monkeypatch): + destinations = [ + M2NDestination("valid", (REPLICATE, 0), torch.empty(8, 8)), + M2NDestination("invalid", (REPLICATE, 1), torch.empty(32, 2)), + ] + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "resolve_parameter_destinations", + lambda *args, **kwargs: destinations, + ) + engine = object.__new__(M2NWeightTransferEngine) + engine.parallel_config = Mock(pipeline_parallel_size=1) + engine.model_config = Mock(quantization=None) + engine.model = torch.nn.Module() + params = ( + self._param( + name="valid", + shape=(32, 8), + src_mesh_dims=(8, 4), + src_placements=(REPLICATE, 0), + ), + self._param( + name="invalid", + shape=(32, 8), + src_mesh_dims=(1, 32), + src_placements=(REPLICATE, 0), + ), + ) + engine._prepare_source_plan(params, num_trainer_ranks=32) + engine._reshard = Mock() + + with pytest.raises(ValueError, match=r"parameter 'invalid'.*MAX_SOURCES=16"): + engine._prepare_destination_plan( + self._init_info( + rank_offset=32, + world_size=36, + dst_mesh_dims=[1, 4], + params=[param.to_dict() for param in params], + ) + ) + + engine._reshard.assert_not_called() + def test_update_preflights_all_names_before_reshard(self): engine = object.__new__(M2NWeightTransferEngine) engine._handle = object() engine.model_update_group = object() engine._index = {"valid": 0} - engine._metas = [M2NParamMeta("valid", torch.float32, (4,), REPLICATED)] + engine._metas = [ + M2NParamMeta( + "valid", + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + ) + ] engine._reshard = Mock() with pytest.raises(ValueError, match=r"parameter 'unknown'"): @@ -310,7 +816,12 @@ class TestWireTypes: engine.model = model engine._index = {"pair.weight": 0, "direct": 1, "pair.scale": 2} engine._metas = [ - M2NParamMeta(name, torch.float32, (1,), REPLICATED) + M2NParamMeta( + name, + torch.float32, + (1,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + ) for name in engine._index ] engine._parameter_destinations = [ @@ -326,9 +837,6 @@ class TestWireTypes: engine._reshard = reshard stream = Mock() - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.comm_ptr", lambda group: 0 - ) monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) engine.receive_weights( @@ -385,6 +893,284 @@ class TestWireTypes: assert info.is_sender +def test_metadata_group_receives_reserved_socket(monkeypatch): + reservation = Mock() + create = Mock(return_value=Mock()) + monkeypatch.setattr("vllm.distributed.utils.StatelessProcessGroup.create", create) + + stateless_init_metadata_group( + "127.0.0.1", + 1234, + rank=0, + world_size=2, + listen_socket=reservation, + ) + + create.assert_called_once_with( + host="127.0.0.1", + port=1234, + rank=0, + world_size=2, + listen_socket=reservation, + ) + + +class _StaticM2NSource(M2NWeightSource): + def __init__(self, metadata, values): + self._metadata = metadata + self._values = values + + def metadata(self): + return list(self._metadata) + + def __iter__(self): + return iter(self._values) + + +class TestTrainerSourceContract: + @staticmethod + def _meta(name="weight"): + return M2NParamMeta( + name, + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + ) + + @staticmethod + def _engine(source): + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + engine._m2n = _fake_m2n() + engine._handle = object() + engine._metas = source.metadata() + engine._dst_mesh = M2NMesh((1, 1), 1) + engine._dst_placements = [REPLICATED] * len(engine._metas) + engine.group = Mock(comm=1) + return engine + + def test_partial_trainer_initialization_releases_transport_resources( + self, monkeypatch + ): + metadata_group = Mock(rank=0, world_size=2) + metadata_group.all_gather_obj.side_effect = lambda envelope: [envelope] * 2 + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer." + "stateless_init_metadata_group", + Mock(return_value=metadata_group), + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer.import_m2n", Mock() + ) + handle = Mock() + nccl = Mock() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer.prepare_m2n_local_runtime", + Mock(return_value=(nccl, 0, handle)), + ) + communicator = Mock() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer.create_m2n_communicator", + Mock(return_value=communicator), + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer." + "publish_destination_placements", + Mock(return_value=[]), + ) + source = _StaticM2NSource([], []) + client = Mock() + client.init_weight_transfer_engine.side_effect = RuntimeError( + "worker init failed" + ) + engine = M2NTrainerWeightTransferEngine( + client=client, source=source, is_sender=True + ) + init_info = M2NTrainerInitInfo( + rank=0, + master_address="127.0.0.1", + master_port=1234, + world_size=2, + ) + + with pytest.raises(RuntimeError, match="worker init failed"): + engine._initialize(init_info, source) + + communicator.abort.assert_called_once_with() + communicator.destroy.assert_not_called() + handle.destroy.assert_called_once_with() + assert engine.group is None + assert engine._handle is None + + def test_non_sender_observes_rank_zero_finish_failure(self): + outcomes: dict[int, dict[str, object]] = {} + sender_group = Mock() + sender_group.send_obj.side_effect = lambda outcome, rank: outcomes.__setitem__( + rank, outcome + ) + sender_group.recv_obj.return_value = {"phase": "round_outcome_consumed"} + + source = _StaticM2NSource([], []) + sender = self._engine(source) + sender.is_sender = True + sender._metadata_group = sender_group + sender._num_trainer_ranks = 2 + sender._send = Mock() + sender.client = Mock() + sender.client.finish_weight_update.side_effect = RuntimeError("finish failed") + + with pytest.raises(RuntimeError, match="finish failed"): + sender.send_weights() + + sender._send.assert_called_once_with() + sender_group.recv_obj.assert_called_once_with(1) + receiver_group = Mock() + receiver_group.recv_obj.side_effect = lambda _source_rank: outcomes.pop(1) + receiver = self._engine(source) + receiver._metadata_group = receiver_group + receiver._num_trainer_ranks = 2 + receiver._send = Mock() + + with pytest.raises( + RuntimeError, + match=( + "worker update failed on trainer rank zero: RuntimeError: finish failed" + ), + ): + receiver.send_weights() + + receiver._send.assert_called_once_with() + receiver_group.recv_obj.assert_called_once_with(0) + receiver_group.send_obj.assert_called_once_with( + {"phase": "round_outcome_consumed"}, 0 + ) + assert not outcomes + + def test_sender_rejects_invalid_round_outcome_acknowledgement(self): + source = _StaticM2NSource([], []) + sender = self._engine(source) + sender.is_sender = True + sender._metadata_group = Mock() + sender._metadata_group.recv_obj.return_value = {"phase": "wrong"} + sender._num_trainer_ranks = 2 + sender._send = Mock() + sender.client = Mock() + + with pytest.raises(RuntimeError, match="invalid round-outcome acknowledgement"): + sender.send_weights() + + def test_non_sender_consumes_outcome_after_local_failure(self): + source = _StaticM2NSource([], []) + receiver = self._engine(source) + receiver._send = Mock(side_effect=[ValueError("local send failed"), None]) + receiver._metadata_group = Mock() + receiver._metadata_group.recv_obj.side_effect = [ + {"ok": False, "error": "RuntimeError: first round failed"}, + {"ok": True, "error": None}, + ] + + with pytest.raises(ValueError, match="local send failed"): + receiver.send_weights() + receiver.send_weights() + + assert receiver._send.call_count == 2 + assert receiver._metadata_group.recv_obj.call_count == 2 + assert receiver._metadata_group.send_obj.call_args_list == [ + call({"phase": "round_outcome_consumed"}, 0), + call({"phase": "round_outcome_consumed"}, 0), + ] + + @pytest.mark.parametrize( + ("case", "match"), + [ + ("missing", "first missing: 'weight'"), + ("extra", "first extra"), + ("reordered", "yielded 'other'.*declared 'weight'"), + ], + ) + def test_source_cardinality_and_order(self, monkeypatch, case, match): + buffer = _BufferLike() + values = { + "missing": [], + "extra": [("weight", buffer), ("extra", buffer)], + "reordered": [("other", buffer)], + }[case] + source = _StaticM2NSource([self._meta()], values) + engine = self._engine(source) + monkeypatch.setattr(torch.cuda, "current_stream", Mock()) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + with pytest.raises(RuntimeError, match=match): + engine._send() + + @pytest.mark.parametrize( + ("case", "match"), [("dtype", "dtype"), ("shape", "local shape")] + ) + def test_materialized_buffer_matches_declared_metadata( + self, monkeypatch, case, match + ): + buffer = _BufferLike() + if case == "dtype": + buffer.dtype = torch.float16 + else: + buffer.shape = (2,) + + source = _StaticM2NSource([self._meta()], [("weight", buffer)]) + engine = self._engine(source) + monkeypatch.setattr(torch.cuda, "current_stream", Mock()) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + with pytest.raises(ValueError, match=match): + engine._send() + + engine._m2n.reshard.assert_not_called() + + def test_m2n_compatible_buffer_needs_no_torch_methods(self, monkeypatch): + buffer = _BufferLike() + source = _StaticM2NSource([self._meta()], [("weight", buffer)]) + engine = self._engine(source) + stream = Mock() + monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + engine._send() + + assert engine._m2n.reshard.call_args.args[0] is buffer + stream.synchronize.assert_called_once_with() + + def test_allocator_aware_buffer_records_send_stream(self, monkeypatch): + buffer = _AllocatorAwareBuffer() + source = _StaticM2NSource([self._meta()], [("weight", buffer)]) + engine = self._engine(source) + stream = Mock() + monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + engine._send() + + buffer.record_stream.assert_called_once_with(stream) + + def test_raw_pointer_uses_declared_shape_and_dtype(self, monkeypatch): + pointer = 1234 + source = _StaticM2NSource([self._meta()], [("weight", pointer)]) + engine = self._engine(source) + monkeypatch.setattr(torch.cuda, "current_stream", Mock()) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + engine._send() + + call = engine._m2n.reshard.call_args + assert call.args[0] == pointer + assert call.kwargs["src_local_shape"] == (4,) + assert call.kwargs["src_dtype"] == torch.float32 + + class TestRegistration: def test_both_registries_expose_the_backend(self): """Both worker and trainer factories register nccl_m2n.""" @@ -433,8 +1219,15 @@ def _assigned_device() -> "torch.device": @ray.remote(num_gpus=1) -def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: - """Send one parameter through the real trainer engine.""" +def _m2n_trainer_send( + master_address: str, + master_port: int, + world_size: int, + nccl_unique_id_b64: str | None = None, + trainer_rank: int = 0, + mixed_layout: bool = False, +) -> bool: + """Send either the basic or mixed-layout manifest through the real engine.""" device = _assigned_device() from vllm.distributed.weight_transfer import WeightTransferTrainerFactory @@ -463,16 +1256,47 @@ def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: ).reshape(SHAPE) ) + num_trainer_ranks = 2 if mixed_layout else 1 + if mixed_layout: + row_full = torch.arange(16, dtype=torch.float32, device=device).reshape(4, 4) + column_full = row_full + 100 + row_local = row_full.chunk(2, dim=0)[trainer_rank].contiguous() + column_local = column_full.chunk(2, dim=1)[trainer_rank].contiguous() + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "row_sharded", + torch.float32, + (4, 4), + M2NLayout(M2NMesh((1, 2), 0), (REPLICATE, 0)), + lambda: row_local, + ), + M2NManifestEntry( + "column_sharded", + torch.float32, + (4, 4), + M2NLayout(M2NMesh((2, 1), 0), (1, REPLICATE)), + lambda: column_local, + ), + ] + ) + dst_mesh_dims = (2, 1) + else: + source = DTensorModuleSource(Tiny(), num_trainer_ranks=1) + dst_mesh_dims = (1, world_size - 1) + engine = WeightTransferTrainerFactory.trainer_init( init_info=M2NTrainerInitInfo( + master_address=master_address, + master_port=master_port, nccl_unique_id_b64=nccl_unique_id_b64, world_size=world_size, - num_trainer_ranks=1, - dst_mesh_dims=(1, world_size - 1), - rank=0, + num_trainer_ranks=num_trainer_ranks, + dst_mesh_dims=dst_mesh_dims, + rank=trainer_rank, ), client=NoopClient(), - source=DTensorModuleSource(Tiny(), num_trainer_ranks=1), + source=source, ) engine.send_weights() torch.accelerator.synchronize() @@ -482,12 +1306,14 @@ def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: @ray.remote(num_gpus=1) def _m2n_worker_receive( - nccl_unique_id_b64: str, + master_address: str, + master_port: int, world_size: int, worker_rank: int = 0, + nccl_unique_id_b64: str | None = None, + mixed_layout: bool = False, ) -> dict: - """Receive that parameter through the real worker engine.""" - import contextlib + """Receive either the basic or mixed-layout manifest through the engine.""" from unittest.mock import MagicMock device = _assigned_device() @@ -503,8 +1329,9 @@ def _m2n_worker_receive( class Recorder(torch.nn.Module): def __init__(self): super().__init__() - self.received = [] - if world_size > 2: + self.received: list[tuple[str, torch.Tensor]] = [] + self.received_by_name: dict[str, torch.Tensor] = {} + if not mixed_layout and world_size > 2: self.weight = torch.nn.Parameter( torch.zeros( SHAPE[0] // (world_size - 1), @@ -521,14 +1348,19 @@ def _m2n_worker_receive( def load_weights(self, weights): for name, tensor in weights: - self.received.append((name, tensor.clone())) + if mixed_layout: + self.received_by_name[name] = tensor.clone() + else: + self.received.append((name, tensor.clone())) + num_trainer_ranks = 2 if mixed_layout else 1 + num_workers = world_size - num_trainer_ranks parallel_config = MagicMock(spec=ParallelConfig) parallel_config.rank = worker_rank - parallel_config.world_size = world_size - 1 + parallel_config.world_size = num_workers parallel_config.data_parallel_rank = 0 parallel_config.data_parallel_index = 0 - parallel_config.tensor_parallel_size = world_size - 1 + parallel_config.tensor_parallel_size = num_workers parallel_config.pipeline_parallel_size = 1 vllm_config = MagicMock() vllm_config.parallel_config = parallel_config @@ -539,37 +1371,52 @@ def _m2n_worker_receive( engine = M2NWeightTransferEngine( WeightTransferConfig(backend="nccl_m2n"), vllm_config, device, recorder ) - # Transport-only: receive_weights enters set_current_vllm_config, and - # vllm_config here is a mock. - import vllm.config as _vllm_config_mod - - _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext() - + if mixed_layout: + params = [ + M2NWireParam("row_sharded", "float32", (4, 4), (1, 2), (REPLICATE, 0)), + M2NWireParam("column_sharded", "float32", (4, 4), (2, 1), (1, REPLICATE)), + ] + names = [param.name for param in params] + dst_mesh_dims = [2, 1] + else: + params = [M2NWireParam("weight", DTYPE, tuple(SHAPE), (1, 1), REPLICATED)] + names = ["weight"] + dst_mesh_dims = [1, num_workers] engine.init_transfer_engine( M2NWeightTransferInitInfo( + schema_version=M2N_WIRE_SCHEMA_VERSION, + master_address=master_address, + master_port=master_port, nccl_unique_id_b64=nccl_unique_id_b64, - rank_offset=1, # trainer occupies rank 0 + rank_offset=num_trainer_ranks, world_size=world_size, - src_mesh_dims=[1, 1], - dst_mesh_dims=[1, world_size - 1], - names=["weight"], - dtype_names=[DTYPE], - shapes=[SHAPE], - src_placements=[None], # replicated on the single trainer rank + dst_mesh_dims=dst_mesh_dims, + source_digest=source_plan_digest(params), + params=[param.to_dict() for param in params], ) ) - engine.receive_weights(M2NWeightTransferUpdateInfo(names=["weight"])) + engine.start_weight_update() + engine.receive_weights(M2NWeightTransferUpdateInfo(names=names)) + engine.finish_weight_update() torch.accelerator.synchronize() + if mixed_layout: + result = { + name: tensor.cpu().tolist() + for name, tensor in recorder.received_by_name.items() + } + engine.shutdown() + return result + full = torch.arange( SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device ).reshape(SHAPE) - direct = world_size > 2 + direct = num_workers > 1 name: str | None got: torch.Tensor | None if direct: name, got = "weight", recorder.weight - expected = full.chunk(world_size - 1, dim=0)[worker_rank] + expected = full.chunk(num_workers, dim=0)[worker_rank] else: name, got = recorder.received[0] if recorder.received else (None, None) expected = full @@ -588,7 +1435,8 @@ def _m2n_worker_receive( torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs: one trainer rank and one inference worker.", ) -def test_m2n_weight_transfer_between_processes(): +@pytest.mark.parametrize("data_plane", ["tcp", "uid"]) +def test_m2n_weight_transfer_between_processes(data_plane): """A parameter survives a real reshard from a trainer process to a worker. This is the only test here that moves bytes: it builds both engines, joins @@ -598,17 +1446,30 @@ def test_m2n_weight_transfer_between_processes(): pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") _init_ray() - from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary - - nccl = NCCLLibrary() - nccl_unique_id_b64 = base64.b64encode( - bytes(nccl.ncclGetUniqueId().internal) - ).decode() + master_address = "127.0.0.1" + master_port = get_open_port() world_size = 2 # 1 trainer + 1 inference worker + nccl_unique_id_b64 = None + if data_plane == "uid": + import nccl.core as nccl - worker = _m2n_worker_receive.remote(nccl_unique_id_b64, world_size) - trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) - trainer_ok, result = ray.get([trainer, worker]) + nccl_unique_id_b64 = base64.b64encode( + bytes(nccl.get_unique_id().as_bytes) + ).decode() + + worker = _m2n_worker_receive.remote( + master_address, + master_port, + world_size, + nccl_unique_id_b64=nccl_unique_id_b64, + ) + trainer = _m2n_trainer_send.remote( + master_address, + master_port, + world_size, + nccl_unique_id_b64, + ) + trainer_ok, result = ray.get([trainer, worker], timeout=300) assert trainer_ok, "trainer engine did not complete" assert result["count"] == 1, f"expected one parameter, got {result['count']}" @@ -626,20 +1487,16 @@ def test_m2n_weight_transfer_to_tp2_shards(): pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") _init_ray() - from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary - - nccl = NCCLLibrary() - nccl_unique_id_b64 = base64.b64encode( - bytes(nccl.ncclGetUniqueId().internal) - ).decode() + master_address = "127.0.0.1" + master_port = get_open_port() world_size = 3 workers = [ - _m2n_worker_receive.remote(nccl_unique_id_b64, world_size, rank) + _m2n_worker_receive.remote(master_address, master_port, world_size, rank) for rank in range(2) ] - trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) - trainer_ok, *results = ray.get([trainer, *workers]) + trainer = _m2n_trainer_send.remote(master_address, master_port, world_size) + trainer_ok, *results = ray.get([trainer, *workers], timeout=300) assert trainer_ok, "trainer engine did not complete" for rank, result in enumerate(results): @@ -650,6 +1507,52 @@ def test_m2n_weight_transfer_to_tp2_shards(): assert result["exact"], f"worker {rank} received the wrong shard" +@pytest.mark.skipif( + torch.accelerator.device_count() < 4, + reason="Need 4 GPUs: two trainer ranks and two inference workers.", +) +def test_m2n_transfer_with_per_parameter_source_meshes(): + """One collective sequence accepts heterogeneous source factorizations.""" + pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") + _init_ray() + + master_address = "127.0.0.1" + master_port = get_open_port() + trainers = [ + _m2n_trainer_send.remote( + master_address, + master_port, + 4, + trainer_rank=rank, + mixed_layout=True, + ) + for rank in range(2) + ] + workers = [ + _m2n_worker_receive.remote( + master_address, + master_port, + 4, + worker_rank=rank, + mixed_layout=True, + ) + for rank in range(2) + ] + results = ray.get(trainers + workers, timeout=300) + trainer_results, worker_results = results[:2], results[2:] + + assert trainer_results == [True, True] + expected_row = torch.arange(16, dtype=torch.float32).reshape(4, 4).tolist() + expected_column = ( + torch.arange(16, dtype=torch.float32).reshape(4, 4) + 100 + ).tolist() + for result in worker_results: + assert result == { + "row_sharded": expected_row, + "column_sharded": expected_column, + } + + class _Model(torch.nn.Module): """Stand-in for a TP=2 shard of a tiny model. diff --git a/vllm/distributed/weight_transfer/m2n_common.py b/vllm/distributed/weight_transfer/m2n_common.py index a2c6b73844..13f7521408 100644 --- a/vllm/distributed/weight_transfer/m2n_common.py +++ b/vllm/distributed/weight_transfer/m2n_common.py @@ -9,16 +9,19 @@ runtime import, layout descriptors, and the conversion to `nccl.m2n` types — so the engine module stays about the transfer itself. """ -from collections.abc import Sequence +import hashlib +import json +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast import torch from vllm.distributed.weight_transfer.base import ParamMeta +from vllm.utils.nccl import unpinned_nccl_env if TYPE_CHECKING: - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup # Placement code for a mesh axis that replicates. Any other (non-negative) code # is the tensor dim that axis shards. Ints rather than `nccl.m2n` objects so @@ -35,15 +38,15 @@ MESH_NDIMS = 2 DESTINATION_REPLICA_AXIS = 0 DESTINATION_SHARD_AXIS = 1 -# m2n bounds how many source shards may feed one destination shard, and how many -# destination shards one source shard may feed, with compile-time arrays in -# `reshard_limits.h`. The bindings do not expose them, so they are mirrored -# here; a build with larger arrays makes these conservative, never wrong. +# Conservative capacities accepted by the supported nccl_m2n algorithms. They +# are not exposed by its Python API yet; remove these mirrors once nccl_m2n +# provides communication-free plan validation. MAX_SOURCE_SHARDS = 16 MAX_DEST_SHARDS = 64 -# Wire dtypes `ncclReshard` accepts. Notably excludes fp4 and any packed -# sub-byte type, so quantized checkpoints are out of scope for now. +# Increment when the JSON init payload changes incompatibly. +M2N_WIRE_SCHEMA_VERSION = 2 + SUPPORTED_DTYPES: frozenset[torch.dtype] = frozenset( dtype for dtype in ( @@ -70,12 +73,62 @@ _IMPORT_HINT = ( "The nccl_m2n weight transfer backend requires the `nccl-extensions` " "package (and its `nccl4py` dependency), which is not installed by vLLM. " "See https://github.com/NVIDIA/nccl-extensions for build " - "and install instructions. It needs NCCL 2.30.5 or newer, and " - "VLLM_NCCL_SO_PATH must point at the same libnccl.so that libnccl_m2n.so " - "was linked against." + "and install instructions. It needs NCCL 2.30.5 or newer." ) +def prepare_m2n_local_runtime(m2n: Any, max_cta: int | None) -> tuple[Any, int, Any]: + """Prepare rank-local resources before any rank enters NCCL init.""" + import nccl.core as nccl + + device = torch.accelerator.current_device_index() + if not isinstance(device, int) or isinstance(device, bool) or device < 0: + raise RuntimeError(f"invalid current accelerator device index: {device!r}") + handle = m2n.Handle.create(m2n.Config(max_cta=max_cta)) + return nccl, device, handle + + +def create_m2n_communicator( + nccl: Any, + group: "StatelessProcessGroup", + unique_id_bytes: bytes | None, +) -> Any: + """Create the M2N data-plane communicator through NCCL4Py.""" + if unique_id_bytes is None: + envelope: dict[str, bytes | str | None] | None = None + if group.rank == 0: + try: + envelope = { + "unique_id": bytes(nccl.get_unique_id().as_bytes), + "error": None, + } + except Exception as exc: + envelope = { + "unique_id": None, + "error": f"{type(exc).__name__}: {exc}", + } + envelope = group.broadcast_obj(envelope, src=0) + if not isinstance(envelope, Mapping) or set(envelope) != {"unique_id", "error"}: + raise RuntimeError("invalid NCCL4Py unique-id envelope") + error = envelope["error"] + received_unique_id = envelope["unique_id"] + if error is not None: + if not isinstance(error, str) or not error: + raise RuntimeError("invalid NCCL4Py unique-id error") + raise RuntimeError(f"NCCL4Py unique-id generation failed: {error}") + if not isinstance(received_unique_id, bytes): + raise RuntimeError("NCCL4Py unique-id broadcast did not return bytes") + unique_id_bytes = received_unique_id + + unique_id = nccl.UniqueId.from_bytes(unique_id_bytes) + with unpinned_nccl_env(): + return nccl.Communicator.init( + nranks=group.world_size, + rank=group.rank, + unique_id=unique_id, + ) + + def import_m2n() -> Any: """Import `nccl.m2n` lazily, with an actionable error when it is missing. @@ -97,17 +150,25 @@ class M2NMesh: `[start_rank, start_rank + dims[0] * dims[1])`. There is no 1-D mesh; a single-axis topology is spelled with a second axis of size 1. - One mesh describes every tensor on its side, so it is exchanged once at the - init handshake rather than per parameter. + A model may use a different factorization for each tensor. For example, + dense weights can use a DP x TP mesh while expert weights use EDP x EP. """ dims: tuple[int, int] start_rank: int def __post_init__(self) -> None: - if len(self.dims) != MESH_NDIMS or any(d <= 0 for d in self.dims): + object.__setattr__(self, "dims", tuple(self.dims)) + if len(self.dims) != MESH_NDIMS or any( + not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 + for dim in self.dims + ): raise ValueError(f"mesh dims must be {MESH_NDIMS} positive ints") - if self.start_rank < 0: + if ( + not isinstance(self.start_rank, int) + or isinstance(self.start_rank, bool) + or self.start_rank < 0 + ): raise ValueError( f"mesh start_rank must be non-negative, got {self.start_rank}" ) @@ -124,16 +185,26 @@ Placements = tuple[int, int] REPLICATED: Placements | None = None -def check_placements(placements: Placements, context: str = "placements") -> None: - """Reject placement pairs m2n cannot express. +@dataclass(frozen=True) +class M2NLayout: + """A tensor's mesh and placement on one side of a transfer.""" - The header requires exactly one SHARD axis and one REPLICATE axis. - `{REPLICATE, REPLICATE}` hits a degenerate prepare branch, which is why - full replication is carried as `REPLICATED` and resolved separately; - sharding both axes is not expressible at all. - """ + mesh: M2NMesh + placements: Placements | None + + def __post_init__(self) -> None: + if not isinstance(self.mesh, M2NMesh): + raise TypeError("M2N layout mesh must be an M2NMesh") + if self.placements is not None: + object.__setattr__(self, "placements", tuple(self.placements)) + + +def check_placements(placements: Placements, context: str = "placements") -> None: + """Reject placement pairs nccl_m2n cannot express.""" if len(placements) != MESH_NDIMS: raise ValueError(f"{context} must have {MESH_NDIMS} entries") + if any(not isinstance(code, int) or isinstance(code, bool) for code in placements): + raise TypeError(f"{context} must contain integer placement codes") invalid = [code for code in placements if code < REPLICATE] if invalid: raise ValueError( @@ -144,8 +215,7 @@ def check_placements(placements: Placements, context: str = "placements") -> Non if num_sharded == 0: raise ValueError( f"{context}: a fully replicated tensor is carried as REPLICATED, " - "not as two " - "REPLICATE axes" + "not as two REPLICATE axes" ) if num_sharded == MESH_NDIMS: raise ValueError( @@ -175,87 +245,46 @@ def resolve_layout( return mesh, placements -def shard_count(mesh: M2NMesh, placements: Placements) -> int: - """How many pieces this layout splits the tensor into.""" - return next( - (mesh.dims[axis] for axis, code in enumerate(placements) if code != REPLICATE), - 1, +def validate_layout( + mesh: M2NMesh, + placements: Placements, + shape: Sequence[int], + context: str, +) -> tuple[int, ...]: + """Validate one resolved layout and return its required local shape.""" + local_shape = list(shape) + for axis, tensor_dim in enumerate(placements): + if tensor_dim == REPLICATE: + continue + if tensor_dim >= len(shape): + raise ValueError( + f"{context} shards tensor dim {tensor_dim}, but parameter has " + f"rank {len(shape)}" + ) + factor = mesh.dims[axis] + if shape[tensor_dim] % factor: + raise ValueError( + f"{context} shards dim {tensor_dim} (size {shape[tensor_dim]}) " + f"over {factor} ranks, which does not divide evenly" + ) + local_shape[tensor_dim] //= factor + return tuple(local_shape) + + +def expected_local_shape(meta: "M2NParamMeta") -> tuple[int, ...]: + """Validate a parameter's source layout and return its local shape.""" + mesh, placements = resolve_layout( + meta.source_layout.mesh, + meta.source_layout.placements, + f"parameter '{meta.name}' source placements", + ) + return validate_layout( + mesh, placements, meta.shape, f"parameter '{meta.name}' source layout" ) -def check_plan_limits( - src: tuple[M2NMesh, Placements], - dst: tuple[M2NMesh, Placements], - name: str, -) -> None: - """Reject plans that exceed m2n's static per-shard fan-in / fan-out arrays. - - Only the unambiguous cases are checked: when one side is a single shard it - is fed by (or feeds) every shard on the other side, so the count is exactly - the other side's shard count. In the general sharded-to-sharded case the - overlap depends on the library's chunking, and reproducing that arithmetic - here would duplicate internals that can change under us. - - m2n re-checks authoritatively and, because the plan is derived from the - shared descriptors, fails identically on every rank -- so this is about - reporting at init with a message that names the parameter, not about - avoiding a hang. - """ - src_shards = shard_count(*src) - dst_shards = shard_count(*dst) - if dst_shards == 1 and src_shards > MAX_SOURCE_SHARDS: - raise ValueError( - f"parameter '{name}' would feed {src_shards} source shards into one " - f"replicated destination, over m2n's MAX_SOURCES={MAX_SOURCE_SHARDS}. " - "Shard the destination, replicate the source, or rebuild m2n with " - "larger arrays." - ) - if src_shards == 1 and dst_shards > MAX_DEST_SHARDS: - raise ValueError( - f"parameter '{name}' would fan one source shard out to {dst_shards} " - f"destination shards, over m2n's MAX_TARGETS={MAX_DEST_SHARDS}. " - "Rebuild m2n with larger arrays to raise the bound." - ) - - -def validate_layout( - mesh: M2NMesh, placements: Placements, shape: Sequence[int], side: str -) -> None: - """Check a resolved layout can describe `shape` before any collective runs. - - Called on both sides at init so a bad plan surfaces as an error from the - init RPC rather than as a hang inside the first reshard. - """ - for axis, code in enumerate(placements): - if code == REPLICATE: - continue - if code >= len(shape): - raise ValueError( - f"{side} layout shards tensor dim {code}, but the tensor " - f"has rank {len(shape)}" - ) - factor = mesh.dims[axis] - if shape[code] % factor: - raise ValueError( - f"{side} layout shards dim {code} (size {shape[code]}) over " - f"{factor} ranks, which does not divide evenly" - ) - - -@dataclass(frozen=True) -class M2NParamMeta(ParamMeta): - """`ParamMeta` extended with how the trainer places this tensor. - - The base class carries only name / dtype / full shape, which is not enough - to plan a reshard. `placements` is relative to its side's `M2NMesh`, or - `REPLICATED` when every rank holds the whole tensor. - """ - - placements: Placements | None - - def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> None: - """Reject tensors m2n cannot move at all, with the parameter named.""" + """Reject parameters nccl_m2n cannot transfer.""" if dtype not in SUPPORTED_DTYPES: raise ValueError( f"parameter '{name}' has dtype {dtype}, which nccl_m2n does not " @@ -268,6 +297,514 @@ def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> N ) +def _shard_spec(mesh: M2NMesh, placements: Placements) -> tuple[int, int]: + return next( + (code, mesh.dims[axis]) + for axis, code in enumerate(placements) + if code != REPLICATE + ) + + +def _max_partition_overlaps( + extent: int, + partitions: int, + other_partitions: int, +) -> int: + """Return the maximum other partitions intersected by any one partition.""" + if extent == 0: + return 0 + width = extent // partitions + other_width = extent // other_partitions + max_overlaps = 0 + for index in range(partitions): + start = index * width + end = start + width + first_other = start // other_width + last_other = (end - 1) // other_width + max_overlaps = max(max_overlaps, last_other - first_other + 1) + return max_overlaps + + +def check_plan_limits( + src: tuple[M2NMesh, Placements], + dst: tuple[M2NMesh, Placements], + name: str, + shape: Sequence[int], +) -> None: + """Reject plans whose per-shard fan-in or fan-out exceeds M2N capacity.""" + src_dim, src_shards = _shard_spec(*src) + dst_dim, dst_shards = _shard_spec(*dst) + if any(extent == 0 for extent in shape): + max_sources = max_targets = 0 + elif src_dim != dst_dim: + max_sources = src_shards + max_targets = dst_shards + else: + extent = shape[src_dim] + max_sources = _max_partition_overlaps(extent, dst_shards, src_shards) + max_targets = _max_partition_overlaps(extent, src_shards, dst_shards) + + if max_sources > MAX_SOURCE_SHARDS: + raise ValueError( + f"parameter '{name}' would feed up to {max_sources} source shards " + f"into one destination shard, over nccl_m2n MAX_SOURCES=" + f"{MAX_SOURCE_SHARDS}" + ) + if max_targets > MAX_DEST_SHARDS: + raise ValueError( + f"parameter '{name}' would fan one source shard out to up to " + f"{max_targets} destination shards, over nccl_m2n " + f"MAX_TARGETS={MAX_DEST_SHARDS}" + ) + + +def validate_transfer_plan( + metas: Sequence["M2NParamMeta"], + dst_mesh: M2NMesh, + dst_placements: Sequence[Placements | None], +) -> None: + """Validate the complete ordered source-to-destination transfer plan.""" + if len(dst_placements) != len(metas): + raise ValueError( + f"destination plan has {len(dst_placements)} parameters, expected " + f"{len(metas)}" + ) + for meta, placement in zip(metas, dst_placements): + check_transferable(meta.name, meta.dtype, meta.shape) + src = resolve_layout( + meta.source_layout.mesh, + meta.source_layout.placements, + f"parameter '{meta.name}' source placements", + ) + validate_layout(*src, meta.shape, f"parameter '{meta.name}' source layout") + dst = resolve_layout( + dst_mesh, + placement, + f"parameter '{meta.name}' destination placements", + ) + validate_layout(*dst, meta.shape, f"parameter '{meta.name}' destination layout") + check_plan_limits(src, dst, meta.name, meta.shape) + + +@dataclass(frozen=True) +class M2NParamMeta(ParamMeta): + """`ParamMeta` extended with how the trainer places this tensor. + + The base class carries only name / dtype / full shape, which is not enough + to plan a reshard. `source_layout` carries both the per-parameter mesh and + its placement on that mesh. + """ + + source_layout: M2NLayout + + def __post_init__(self) -> None: + if not isinstance(self.source_layout, M2NLayout): + raise TypeError("source_layout must be an M2NLayout") + + +@dataclass(frozen=True) +class M2NWireParam: + """JSON-safe source metadata for one parameter in the init handshake.""" + + name: str + dtype_name: str + shape: tuple[int, ...] + src_mesh_dims: tuple[int, int] + src_placements: Placements | None + + def __post_init__(self) -> None: + shape = tuple(self.shape) + dims = tuple(self.src_mesh_dims) + placements = None if self.src_placements is None else tuple(self.src_placements) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "src_mesh_dims", dims) + object.__setattr__(self, "src_placements", placements) + if not isinstance(self.name, str) or not self.name: + raise ValueError("M2N wire parameter name must not be empty") + if not isinstance(self.dtype_name, str) or not self.dtype_name: + raise ValueError(f"parameter '{self.name}' has an empty wire dtype name") + if not shape or any( + not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 + for dim in shape + ): + raise ValueError( + f"parameter '{self.name}' shape must contain positive integers" + ) + M2NMesh(cast(tuple[int, int], dims), 0) + if placements is not None: + check_placements( + cast(Placements, placements), + f"parameter '{self.name}' source placements", + ) + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "M2NWireParam": + """Parse one strict wire record from its JSON-decoded mapping.""" + fields = { + "name", + "dtype_name", + "shape", + "src_mesh_dims", + "src_placements", + } + missing = fields - value.keys() + extra = value.keys() - fields + if missing or extra: + details = [] + if missing: + details.append(f"missing {sorted(missing)}") + if extra: + details.append(f"unexpected {sorted(extra)}") + raise ValueError(f"invalid M2N wire parameter: {', '.join(details)}") + return cls(**value) + + def to_dict(self) -> dict[str, Any]: + """Serialize this record using only JSON-compatible values.""" + return { + "name": self.name, + "dtype_name": self.dtype_name, + "shape": list(self.shape), + "src_mesh_dims": list(self.src_mesh_dims), + "src_placements": ( + None if self.src_placements is None else list(self.src_placements) + ), + } + + +def source_plan_digest(params: Sequence[M2NWireParam]) -> str: + """Return a rank-independent digest of the ordered source wire plan.""" + payload = json.dumps( + { + "schema_version": M2N_WIRE_SCHEMA_VERSION, + "params": [param.to_dict() for param in params], + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def complete_plan_digest( + source_digest: str, + dst_mesh: M2NMesh, + dst_placements: Sequence[Placements | None], +) -> str: + """Return a digest of the agreed source and destination plan.""" + if not isinstance(source_digest, str) or not source_digest: + raise ValueError("source digest must be a non-empty string") + payload = json.dumps( + { + "schema_version": M2N_WIRE_SCHEMA_VERSION, + "source_digest": source_digest, + "destination_mesh": { + "dims": list(dst_mesh.dims), + "start_rank": dst_mesh.start_rank, + }, + "destination_placements": [ + None if placement is None else list(placement) + for placement in dst_placements + ], + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def check_source_plan_agreement( + group: "StatelessProcessGroup", + params: Sequence[M2NWireParam], + expected_digest: str | None = None, + local_error: str | None = None, +) -> str: + """Verify every participant entered with the same source plan.""" + prepared_error = local_error + digest: str | None = None + declared: str | None = expected_digest + if prepared_error is None: + try: + digest = source_plan_digest(params) + declared = digest if expected_digest is None else expected_digest + if not isinstance(declared, str) or not declared: + raise ValueError("declared source digest must be a non-empty string") + except Exception as exc: + prepared_error = f"{type(exc).__name__}: {exc}" + elif not isinstance(prepared_error, str) or not prepared_error: + prepared_error = "TypeError: local source error must be a non-empty string" + + gathered = group.all_gather_obj( + { + "phase": "source", + "digest": digest, + "declared_digest": declared, + "error": prepared_error, + } + ) + errors: list[str] = [] + expected_fields = {"phase", "digest", "declared_digest", "error"} + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + agreed: list[tuple[int, str, str]] = [] + for rank, item in enumerate(gathered): + if not isinstance(item, Mapping) or set(item) != expected_fields: + errors.append(f"rank {rank}: invalid source preflight envelope") + continue + if item["phase"] != "source": + errors.append(f"rank {rank}: invalid source preflight phase") + continue + error = item["error"] + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid source preflight error") + else: + errors.append(f"rank {rank}: {error}") + continue + computed = item["digest"] + claimed = item["declared_digest"] + if ( + not isinstance(computed, str) + or not computed + or not isinstance(claimed, str) + or not claimed + ): + errors.append(f"rank {rank}: invalid source digest") + continue + agreed.append((rank, computed, claimed)) + if errors: + raise RuntimeError( + "nccl_m2n source preflight rejected before NCCL init: " + "; ".join(errors) + ) + + canonical = agreed[0][1] + for rank, computed, claimed in agreed: + if computed != claimed: + errors.append( + f"rank {rank}: computed digest {computed} does not match " + f"declared digest {claimed}" + ) + if computed != canonical: + errors.append(f"rank {rank}: source digest disagrees with rank 0") + if errors: + raise RuntimeError( + "nccl_m2n source preflight rejected before NCCL init: " + "; ".join(errors) + ) + return canonical + + +def check_complete_plan_agreement( + group: "StatelessProcessGroup", + source_digest: str, + dst_mesh: M2NMesh, + dst_placements: Sequence[Placements | None], + local_error: str | None = None, +) -> str: + """Require every rank to accept the same complete transfer plan.""" + prepared_error = local_error + digest: str | None = None + if prepared_error is None: + try: + digest = complete_plan_digest(source_digest, dst_mesh, dst_placements) + except Exception as exc: + prepared_error = f"{type(exc).__name__}: {exc}" + elif not isinstance(prepared_error, str) or not prepared_error: + prepared_error = "TypeError: local plan error must be a non-empty string" + + gathered = group.all_gather_obj( + {"phase": "complete_plan", "digest": digest, "error": prepared_error} + ) + errors: list[str] = [] + valid: list[tuple[int, str]] = [] + expected_fields = {"phase", "digest", "error"} + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + for rank, item in enumerate(gathered): + if ( + not isinstance(item, Mapping) + or set(item) != expected_fields + or item.get("phase") != "complete_plan" + ): + errors.append(f"rank {rank}: invalid complete-plan envelope") + continue + error = item["error"] + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid complete-plan error") + else: + errors.append(f"rank {rank}: {error}") + continue + rank_digest = item["digest"] + if not isinstance(rank_digest, str) or not rank_digest: + errors.append(f"rank {rank}: invalid complete-plan digest") + continue + valid.append((rank, rank_digest)) + if not errors: + canonical = valid[0][1] + for rank, rank_digest in valid[1:]: + if rank_digest != canonical: + errors.append( + f"rank {rank}: complete-plan digest disagrees with rank 0" + ) + if errors: + raise RuntimeError( + "nccl_m2n complete-plan preflight rejected before first reshard: " + + "; ".join(errors) + ) + return valid[0][1] + + +def check_runtime_ready_agreement( + group: "StatelessProcessGroup", local_error: str | None +) -> None: + """Require every rank to prepare local resources before NCCL init.""" + prepared_error = local_error + if prepared_error is not None and ( + not isinstance(prepared_error, str) or not prepared_error + ): + prepared_error = "TypeError: local runtime error must be a non-empty string" + gathered = group.all_gather_obj({"phase": "runtime_ready", "error": prepared_error}) + errors: list[str] = [] + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + for rank, item in enumerate(gathered): + if ( + not isinstance(item, Mapping) + or set(item) != {"phase", "error"} + or item.get("phase") != "runtime_ready" + ): + errors.append(f"rank {rank}: invalid runtime readiness envelope") + continue + error = item.get("error") + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid runtime readiness error") + else: + errors.append(f"rank {rank}: {error}") + if errors: + raise RuntimeError( + "nccl_m2n local runtime preflight failed before NCCL init: " + + "; ".join(errors) + ) + + +def check_data_plane_agreement( + group: "StatelessProcessGroup", + unique_id_bytes: bytes | None, + local_error: str | None = None, +) -> None: + """Agree on the NCCL rendezvous identity before NCCL init.""" + prepared_error = local_error + uid_digest: str | None = None + if prepared_error is None: + try: + if unique_id_bytes is not None: + if not isinstance(unique_id_bytes, bytes): + raise TypeError("NCCL unique id must decode to bytes") + uid_digest = hashlib.sha256(unique_id_bytes).hexdigest() + except Exception as exc: + prepared_error = f"{type(exc).__name__}: {exc}" + elif not isinstance(prepared_error, str) or not prepared_error: + prepared_error = "TypeError: local data-plane error must be non-empty" + + gathered = group.all_gather_obj( + { + "phase": "data_plane", + "mode": "uid" if unique_id_bytes is not None else "tcp", + "uid_digest": uid_digest, + "error": prepared_error, + } + ) + expected_fields = { + "phase", + "mode", + "uid_digest", + "error", + } + errors: list[str] = [] + valid: list[tuple[int, str, str | None]] = [] + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + for rank, item in enumerate(gathered): + if ( + not isinstance(item, Mapping) + or set(item) != expected_fields + or item.get("phase") != "data_plane" + ): + errors.append(f"rank {rank}: invalid data-plane preflight envelope") + continue + error = item["error"] + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid data-plane preflight error") + else: + errors.append(f"rank {rank}: {error}") + continue + mode = item["mode"] + uid = item["uid_digest"] + if mode not in {"tcp", "uid"} or (mode == "uid") != isinstance(uid, str): + errors.append(f"rank {rank}: invalid data-plane identity") + continue + valid.append((rank, mode, uid)) + if not errors: + _, mode, uid = valid[0] + for rank, other_mode, other_uid in valid[1:]: + if (other_mode, other_uid) != (mode, uid): + errors.append(f"rank {rank}: NCCL data-plane identity disagrees") + if errors: + raise RuntimeError( + "nccl_m2n data-plane preflight rejected before NCCL init: " + + "; ".join(errors) + ) + + +def validate_local_buffer(m2n: Any, meta: M2NParamMeta, buffer: Any) -> None: + """Check source dtype/shape while delegating buffer validity to M2N.""" + # A raw pointer carries no inspectable metadata. The source owns its + # correctness; _send supplies the declared dtype and local shape to M2N. + if isinstance(buffer, int): + return + + mesh, placements = resolve_layout( + meta.source_layout.mesh, meta.source_layout.placements + ) + actual = m2n.DistTensor( + buffer, + mesh=to_mesh(m2n, mesh), + placements=to_placements(m2n, placements), + ) + declared = m2n.DistTensor( + None, + local_shape=actual.local_shape, + dtype=meta.dtype, + mesh=to_mesh(m2n, mesh), + placements=to_placements(m2n, placements), + ) + if actual.dtype != declared.dtype: + raise ValueError( + f"parameter '{meta.name}' source returned dtype {actual.dtype}, " + f"but its metadata declared {declared.dtype}" + ) + + expected = expected_local_shape(meta) + if actual.local_shape != expected: + raise ValueError( + f"parameter '{meta.name}' source returned local shape " + f"{actual.local_shape}, but layout {meta.source_layout} implies " + f"{expected} from global shape {meta.shape}" + ) + + def to_mesh(m2n: Any, mesh: M2NMesh) -> Any: return m2n.Mesh(mesh.dims, start_rank=mesh.start_rank) @@ -279,7 +816,9 @@ def to_placements(m2n: Any, placements: Placements) -> list[Any]: def publish_destination_placements( - comm: "PyNcclCommunicator", + comm: Any, + device: torch.device, + stream: torch.cuda.Stream, first_worker_rank: int, placements: "Sequence[Placements | None] | None", num_parameters: int, @@ -294,15 +833,15 @@ def publish_destination_placements( rather than deadlocking later. Only placements travel: the destination *mesh* is still derived from the - rank split, identically on both sides. Use the NCCL communicator directly - because unique-id rendezvous deliberately has no bootstrap process group. + rank split, identically on both sides. The plan is published on the data + plane before any M2N reshard collective. """ replicated_sentinel = REPLICATE - 1 encoded = torch.full( (num_parameters, MESH_NDIMS), replicated_sentinel, dtype=torch.int8, - device=comm.device, + device=device, ) if comm.rank == first_worker_rank: if placements is None or len(placements) != num_parameters: @@ -314,25 +853,12 @@ def publish_destination_placements( for placement in placements ] if rows: - encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=comm.device)) + encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=device)) - comm.broadcast(encoded, src=first_worker_rank) + comm.broadcast(encoded, encoded, root=first_worker_rank, stream=stream) return [ REPLICATED if all(code == replicated_sentinel for code in row) else cast(Placements, tuple(row)) for row in encoded.tolist() ] - - -def comm_ptr(comm: "PyNcclCommunicator") -> int: - """Raw `ncclComm_t` behind vLLM's `PyNcclCommunicator`. - - m2n links its own NCCL, so the handle only means anything if vLLM loaded - the same `libnccl.so` — set `VLLM_NCCL_SO_PATH` accordingly. - """ - handle = comm.comm - ptr = getattr(handle, "value", handle) - if not ptr: - raise RuntimeError("PyNcclCommunicator has no live NCCL communicator") - return int(ptr) diff --git a/vllm/distributed/weight_transfer/m2n_engine.py b/vllm/distributed/weight_transfer/m2n_engine.py index e480e60eaf..66dbb6ba24 100644 --- a/vllm/distributed/weight_transfer/m2n_engine.py +++ b/vllm/distributed/weight_transfer/m2n_engine.py @@ -20,7 +20,7 @@ inside `receive_weights` while the trainer is sending. Driving that from the trainer is the trainer engine's job. """ -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast @@ -34,20 +34,28 @@ from vllm.distributed.weight_transfer.base import ( ) from vllm.distributed.weight_transfer.m2n_common import ( DESTINATION_SHARD_AXIS, + M2N_WIRE_SCHEMA_VERSION, MESH_NDIMS, - REPLICATED, + M2NLayout, M2NMesh, M2NParamMeta, + M2NWireParam, Placements, - check_plan_limits, + check_complete_plan_agreement, + check_data_plane_agreement, + check_runtime_ready_agreement, + check_source_plan_agreement, check_transferable, - comm_ptr, + create_m2n_communicator, + expected_local_shape, import_m2n, + prepare_m2n_local_runtime, publish_destination_placements, resolve_layout, + source_plan_digest, to_mesh, to_placements, - validate_layout, + validate_transfer_plan, ) from vllm.distributed.weight_transfer.m2n_layout import ( M2NDestination, @@ -56,12 +64,14 @@ from vllm.distributed.weight_transfer.m2n_layout import ( from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, decode_nccl_unique_id, - worker_init_process_group, + worker_init_metadata_group, ) +from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + +logger = init_logger(__name__) __all__ = [ "M2NWeightTransferInitInfo", @@ -84,59 +94,99 @@ class M2NWeightTransferInitInfo(WeightTransferInitInfo): plain JSON so the HTTP control plane carries it unchanged. """ - master_address: str | None = None - master_port: int | None = None + schema_version: int + master_address: str + master_port: int rank_offset: int """First worker rank, i.e. the number of trainer ranks.""" world_size: int """Trainer ranks + all inference workers.""" - src_mesh_dims: list[int] - """The trainer's mesh, shared by every parameter (`start_rank` is 0: the - trainer occupies the front of the communicator).""" dst_mesh_dims: list[int] """The inference mesh, starting at `rank_offset`. Declared rather than derived so both sides describe the destination identically.""" - names: list[str] - dtype_names: list[str] - shapes: list[list[int]] - src_placements: list[list[int] | None] - """Per parameter, relative to `src_mesh_dims`; `None` means replicated.""" + source_digest: str + """SHA-256 of the ordered, versioned source parameter plan.""" + params: list[dict[str, Any]] + """Ordered per-parameter source metadata and layouts.""" nccl_unique_id_b64: str | None = field(default=None, repr=False) max_cta: int | None = None def __post_init__(self) -> None: - _ = self.nccl_unique_id_bytes - num_params = len(self.names) - for label, values in ( - ("dtype_names", self.dtype_names), - ("shapes", self.shapes), - ("src_placements", self.src_placements), + if not isinstance(self.master_address, str) or not self.master_address: + raise ValueError("`master_address` must name the metadata rendezvous") + if ( + not isinstance(self.master_port, int) + or isinstance(self.master_port, bool) + or not 0 < self.master_port < 65536 + ): + raise ValueError("`master_port` must be a valid metadata port") + if ( + not isinstance(self.rank_offset, int) + or isinstance(self.rank_offset, bool) + or not isinstance(self.world_size, int) + or isinstance(self.world_size, bool) + or self.rank_offset < 1 + or self.rank_offset >= self.world_size ): - if len(values) != num_params: - raise ValueError( - f"`{label}` should be of the same size as `names`: " - f"got {len(values)} and {num_params}" - ) - if self.rank_offset < 1 or self.rank_offset >= self.world_size: raise ValueError( f"`rank_offset` ({self.rank_offset}) must leave at least one " f"trainer rank and one worker in world_size {self.world_size}" ) num_workers = self.world_size - self.rank_offset dst = tuple(self.dst_mesh_dims) - if len(dst) != MESH_NDIMS or dst[0] * dst[1] != num_workers: + try: + dst_mesh = M2NMesh(cast(tuple[int, int], dst), self.rank_offset) + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid `dst_mesh_dims` {dst}: {exc}") from exc + if dst_mesh.size != num_workers: raise ValueError( f"`dst_mesh_dims` {dst} must be {MESH_NDIMS} dims covering the " f"{num_workers} inference workers" ) + def parse_wire_params(self) -> tuple[M2NWireParam, ...]: + """Parse and authenticate the ordered source plan.""" + if self.schema_version != M2N_WIRE_SCHEMA_VERSION: + raise ValueError( + "unsupported nccl_m2n init schema version " + f"{self.schema_version}; expected {M2N_WIRE_SCHEMA_VERSION}" + ) + if not isinstance(self.params, list): + raise TypeError("`params` must be a list of wire parameter mappings") + if not isinstance(self.source_digest, str) or not self.source_digest: + raise ValueError("`source_digest` must be a non-empty string") + wire_params: list[M2NWireParam] = [] + names: set[str] = set() + for index, value in enumerate(self.params): + if not isinstance(value, Mapping): + raise TypeError( + f"`params[{index}]` must be a mapping, got {type(value).__name__}" + ) + param = M2NWireParam.from_dict(value) + if param.name in names: + raise ValueError( + f"nccl_m2n init params contain duplicate parameter '{param.name}'" + ) + names.add(param.name) + wire_params.append(param) + parsed = tuple(wire_params) + computed_digest = source_plan_digest(parsed) + if self.source_digest != computed_digest: + raise ValueError( + "nccl_m2n init source digest does not match params: " + f"computed {computed_digest}, declared {self.source_digest}" + ) + return parsed + @property def nccl_unique_id_bytes(self) -> bytes | None: + if self.nccl_unique_id_b64 is None: + return None return decode_nccl_unique_id( - master_address=self.master_address, - master_port=self.master_port, + master_address=None, + master_port=None, nccl_unique_id_b64=self.nccl_unique_id_b64, - ctx="M2NWeightTransferInitInfo", + ctx="M2NWeightTransferInitInfo data plane", ) @@ -182,33 +232,172 @@ class M2NWeightTransferEngine( model: torch.nn.Module, ) -> None: super().__init__(config, vllm_config, device, model) - self.model_update_group: PyNcclCommunicator | None = None + self.model_update_group: Any = None self._m2n: Any = None self._handle: Any = None self._metas: list[M2NParamMeta] = [] - self._src_mesh: M2NMesh | None = None self._dst_mesh: M2NMesh | None = None self._parameter_destinations: list[M2NDestination] = [] self._index: dict[str, int] = {} self._uses_load_weights = False def init_transfer_engine(self, init_info: M2NWeightTransferInitInfo) -> None: - """Join the trainer's communicator and rebuild the transfer plan. - - Every precondition (dtype, tensor rank, divisibility) is checked here so - a bad plan fails the init RPC instead of hanging the first collective. - """ - self._m2n = import_m2n() - self._metas = [] - - self._src_mesh = M2NMesh( - cast(tuple[int, int], tuple(init_info.src_mesh_dims)), 0 - ) - if self._src_mesh.size != init_info.rank_offset: - raise ValueError( - f"source mesh covers {self._src_mesh.size} ranks, but there are " - f"{init_info.rank_offset} trainer ranks" + """Agree on source and local readiness before constructing NCCL.""" + if self._handle is not None or self.model_update_group is not None: + raise RuntimeError( + "nccl_m2n weight transfer is already initialized; call " + "shutdown_weight_transfer_engine() before initializing again" ) + rendezvous_info = NCCLWeightTransferInitInfo( + master_address=init_info.master_address, + master_port=init_info.master_port, + rank_offset=init_info.rank_offset, + world_size=init_info.world_size, + ) + metadata_group = worker_init_metadata_group( + rendezvous_info, self.parallel_config + ) + + wire_params: tuple[M2NWireParam, ...] = () + source_error: str | None = None + try: + self._m2n = import_m2n() + wire_params = init_info.parse_wire_params() + self._prepare_source_plan(wire_params, init_info.rank_offset) + except Exception as exc: + source_error = f"{type(exc).__name__}: {exc}" + self._metas = [] + check_source_plan_agreement( + metadata_group, + wire_params, + expected_digest=init_info.source_digest, + local_error=source_error, + ) + + destination_error: str | None = None + try: + self._prepare_destination_plan(init_info) + except Exception as exc: + destination_error = f"{type(exc).__name__}: {exc}" + self._parameter_destinations = [] + + unique_id_bytes: bytes | None = None + data_plane_error: str | None = None + try: + unique_id_bytes = init_info.nccl_unique_id_bytes + except Exception as exc: + data_plane_error = f"{type(exc).__name__}: {exc}" + check_data_plane_agreement( + metadata_group, + unique_id_bytes, + data_plane_error, + ) + + handle = None + group = None + try: + runtime_error = destination_error + nccl = None + device = None + if runtime_error is None: + try: + nccl, device, handle = prepare_m2n_local_runtime( + self._m2n, init_info.max_cta + ) + except Exception as exc: + runtime_error = f"{type(exc).__name__}: {exc}" + check_runtime_ready_agreement(metadata_group, runtime_error) + assert nccl is not None and device is not None and handle is not None + group = create_m2n_communicator(nccl, metadata_group, unique_id_bytes) + + # Keep the merged destination-placement publication unchanged: source + # agreement is independent of worker-side destination selection. + mine = [ + destination.placements for destination in self._parameter_destinations + ] + agreed = publish_destination_placements( + group, + self.device, + torch.cuda.current_stream(), + init_info.rank_offset, + mine, + len(self._metas), + ) + complete_plan_error: str | None = None + try: + if agreed != mine: + mismatched = next( + meta.name + for meta, ours, theirs in zip(self._metas, mine, agreed) + if ours != theirs + ) + raise RuntimeError( + "inference workers resolved different nccl_m2n " + f"destinations (first mismatch: '{mismatched}'); all " + "workers must run the same model and parallel config" + ) + assert self._dst_mesh is not None + validate_transfer_plan(self._metas, self._dst_mesh, agreed) + except Exception as exc: + complete_plan_error = f"{type(exc).__name__}: {exc}" + assert self._dst_mesh is not None + check_complete_plan_agreement( + metadata_group, + init_info.source_digest, + self._dst_mesh, + agreed, + complete_plan_error, + ) + except BaseException: + if group is not None: + try: + # Transport cleanup only; reload validity belongs to the + # model-update lifecycle, not this backend. + group.abort() + except BaseException: + logger.exception( + "failed to abort partially initialized nccl_m2n communicator" + ) + if handle is not None: + try: + handle.destroy() + except BaseException: + logger.exception( + "failed to destroy partially initialized nccl_m2n handle" + ) + raise + + self._handle = handle + self.model_update_group = group + + def _prepare_source_plan( + self, + wire_params: tuple[M2NWireParam, ...], + num_trainer_ranks: int, + ) -> None: + """Validate worker-visible source metadata without constructing NCCL.""" + self._metas = [] + for param in wire_params: + dtype = getattr(torch, param.dtype_name, None) + if not isinstance(dtype, torch.dtype): + raise ValueError( + f"parameter '{param.name}' has invalid dtype name " + f"'{param.dtype_name}'; expected the name of a torch.dtype" + ) + check_transferable(param.name, dtype, param.shape) + layout = M2NLayout(M2NMesh(param.src_mesh_dims, 0), param.src_placements) + if layout.mesh.size != num_trainer_ranks: + raise ValueError( + f"parameter '{param.name}' source mesh covers " + f"{layout.mesh.size} ranks, but there are " + f"{num_trainer_ranks} trainer ranks" + ) + meta = M2NParamMeta(param.name, dtype, param.shape, layout) + expected_local_shape(meta) + self._metas.append(meta) + + def _prepare_destination_plan(self, init_info: M2NWeightTransferInitInfo) -> None: + """Resolve the merged-#51520 in-place/full-fallback destination plan.""" # The inference topology, as declared by the trainer. This is the mesh # a sharded destination is placed over; a replicated one is described # by the alternate descriptor `resolve_layout` derives from it, which @@ -217,33 +406,6 @@ class M2NWeightTransferEngine( cast(tuple[int, int], tuple(init_info.dst_mesh_dims)), init_info.rank_offset, ) - - for name, name_dtype, shape, placements in zip( - init_info.names, - init_info.dtype_names, - init_info.shapes, - init_info.src_placements, - ): - dtype = getattr(torch, name_dtype, None) - if not isinstance(dtype, torch.dtype): - raise ValueError( - f"parameter '{name}' has invalid dtype name " - f"'{name_dtype}'; expected the name of a torch.dtype" - ) - check_transferable(name, dtype, shape) - codes = ( - REPLICATED - if placements is None - else cast(Placements, tuple(placements)) - ) - src_mesh, src_placements = resolve_layout( - self._src_mesh, codes, f"parameter '{name}' source placements" - ) - validate_layout(src_mesh, src_placements, shape, "source") - # Destination checks are deferred until loader-aware resolution; - # assuming replication here can reject a valid sharded plan. - self._metas.append(M2NParamMeta(name, dtype, tuple(shape), codes)) - parallel_config = self.parallel_config shard_axis_size = self._dst_mesh.dims[DESTINATION_SHARD_AXIS] quantization = getattr(self.model_config, "quantization", None) @@ -266,27 +428,11 @@ class M2NWeightTransferEngine( shard_axis_size=shard_axis_size, allow_direct=allow_direct, ) - # Check the resolved layout rather than a provisional replicated one. - # For example, 32 source shards feeding 4 destination shards may need - # only 8 sources per destination, while a replicated plan needs all 32. - for meta, destination in zip(self._metas, self._parameter_destinations): - src_mesh, src_placements = resolve_layout( - self._src_mesh, - meta.placements, - f"parameter '{meta.name}' source placements", - ) - dst_mesh, dst_placements = resolve_layout( - self._dst_mesh, - destination.placements, - f"parameter '{meta.name}' destination placements", - ) - validate_layout(dst_mesh, dst_placements, meta.shape, "destination") - check_plan_limits( - (src_mesh, src_placements), - (dst_mesh, dst_placements), - meta.name, - ) - + validate_transfer_plan( + self._metas, + self._dst_mesh, + [destination.placements for destination in self._parameter_destinations], + ) # Update requests carry names only, so cache their plan indices. Track # whether any fallback entry requires the `load_weights` lifecycle. self._index = {meta.name: i for i, meta in enumerate(self._metas)} @@ -294,38 +440,6 @@ class M2NWeightTransferEngine( not destination.direct for destination in self._parameter_destinations ) - self.model_update_group = worker_init_process_group( - NCCLWeightTransferInitInfo( - master_address=init_info.master_address, - master_port=init_info.master_port, - nccl_unique_id_b64=init_info.nccl_unique_id_b64, - rank_offset=init_info.rank_offset, - world_size=init_info.world_size, - ), - parallel_config, - ) - # Destinations are per-parameter and depend on the inference model, so - # the trainer cannot derive them. The first worker publishes the plan - # and every other rank checks its own against it — a silent - # disagreement between workers would mean mismatched collectives. - mine = [destination.placements for destination in self._parameter_destinations] - agreed = publish_destination_placements( - self.model_update_group, init_info.rank_offset, mine, len(self._metas) - ) - if agreed != mine: - mismatched = next( - meta.name for meta, a, b in zip(self._metas, mine, agreed) if a != b - ) - raise RuntimeError( - "inference workers resolved different nccl_m2n destinations " - f"(first mismatch: '{mismatched}'); all workers must run the " - "same model and parallel config" - ) - - self._handle = self._m2n.Handle.create( - self._m2n.Config(max_cta=init_info.max_cta) - ) - def start_weight_update(self) -> None: """Set up layerwise reloading, but only if some parameter needs it. @@ -335,7 +449,9 @@ class M2NWeightTransferEngine( """ if not self._uses_load_weights: return - from vllm.model_executor.model_loader.reload import initialize_layerwise_reload + from vllm.model_executor.model_loader.reload import ( + initialize_layerwise_reload, + ) initialize_layerwise_reload(self.model) @@ -354,6 +470,10 @@ class M2NWeightTransferEngine( "nccl_m2n weight transfer not initialized. " "Call init_transfer_engine() first." ) + if not isinstance(update_info.names, list) or any( + not isinstance(name, str) for name in update_info.names + ): + raise TypeError("nccl_m2n update names must be a list of strings") requested = [] for name in update_info.names: @@ -369,7 +489,6 @@ class M2NWeightTransferEngine( disable_mtp_completeness_check, ) - comm = comm_ptr(self.model_update_group) stream = torch.cuda.current_stream() # Reshard lazily so one `load_weights` invocation sees the complete @@ -385,7 +504,13 @@ class M2NWeightTransferEngine( meta.shape, dtype=meta.dtype, device=self.device ) - self._reshard(comm, stream, meta, destination.placements, buffer) + self._reshard( + self.model_update_group, + stream, + meta, + destination.placements, + buffer, + ) if destination.tensor is None: # `load_weights` reads on the host stream, so the transfer @@ -410,7 +535,7 @@ class M2NWeightTransferEngine( def _reshard( self, - comm: int, + comm: Any, stream: torch.cuda.Stream, src_meta: M2NParamMeta, dst_placements: Placements | None, @@ -421,11 +546,10 @@ class M2NWeightTransferEngine( `dst_placements` describes the buffer's logical placement over the worker mesh; `dst_buffer` is this rank's physical storage. """ - assert self._src_mesh is not None assert self._dst_mesh is not None m2n = self._m2n src_mesh, resolved_src_placements = resolve_layout( - self._src_mesh, src_meta.placements + src_meta.source_layout.mesh, src_meta.source_layout.placements ) dst_mesh, resolved_dst_placements = resolve_layout( self._dst_mesh, dst_placements @@ -444,12 +568,32 @@ class M2NWeightTransferEngine( def shutdown(self) -> None: """Finish pending GPU work and release M2N communication state.""" - if self._handle is not None: - # M2N does not synchronize caller streams on finalize. - torch.accelerator.synchronize() - self._handle.destroy() - self._handle = None - if self.model_update_group is not None: - self.model_update_group.destroy() - self.model_update_group = None + handle, self._handle = self._handle, None + group, self.model_update_group = self.model_update_group, None + self._m2n = None + self._metas = [] + self._dst_mesh = None self._parameter_destinations = [] + self._index = {} + self._uses_load_weights = False + + first_error: BaseException | None = None + + def attempt(action: Any, message: str) -> None: + nonlocal first_error + try: + action() + except BaseException as exc: + if first_error is None: + first_error = exc + logger.exception(message) + + if handle is not None or group is not None: + # M2N does not synchronize caller streams on finalize. + attempt(torch.accelerator.synchronize, "failed to synchronize nccl_m2n") + if handle is not None: + attempt(handle.destroy, "failed to destroy nccl_m2n handle") + if group is not None: + attempt(group.abort, "failed to abort nccl_m2n communicator") + if first_error is not None: + raise first_error diff --git a/vllm/distributed/weight_transfer/m2n_source.py b/vllm/distributed/weight_transfer/m2n_source.py index dd08d24b56..d9c9e9194a 100644 --- a/vllm/distributed/weight_transfer/m2n_source.py +++ b/vllm/distributed/weight_transfer/m2n_source.py @@ -5,21 +5,22 @@ m2n plans a transfer from both sides' layouts, so the trainer has to say how it holds each parameter — which the base `WeightSource` / `ParamMeta` pair does not express. This module supplies the m2n flavor of both, and the DTensor-backed -implementation that covers the common trainer. - -A source declares its mesh once (`mesh()`) and a placement per parameter, since -one topology describes every tensor on the side. +implementation that covers the common trainer. Each parameter owns its source +mesh and placements because dense and expert weights can use different +factorizations over the same trainer ranks. """ -from collections.abc import Iterator +from collections.abc import Callable, Iterator, Sequence +from dataclasses import dataclass from typing import Any import torch -from vllm.distributed.weight_transfer.base import ParamMeta, WeightSource +from vllm.distributed.weight_transfer.base import WeightSource from vllm.distributed.weight_transfer.m2n_common import ( REPLICATE, REPLICATED, + M2NLayout, M2NMesh, M2NParamMeta, Placements, @@ -27,6 +28,8 @@ from vllm.distributed.weight_transfer.m2n_common import ( __all__ = [ "M2NWeightSource", + "M2NManifestEntry", + "ManifestM2NWeightSource", "DTensorModuleSource", "mesh_from_tensor", "placements_from_tensor", @@ -40,17 +43,76 @@ class M2NWeightSource(WeightSource): materialized full tensor: gathering is exactly the cost m2n removes. """ - def mesh(self) -> M2NMesh: - """The trainer's rank topology, shared by every parameter.""" + def metadata(self) -> list[M2NParamMeta]: # type: ignore[override] + """Name, dtype, full shape, and source layout for each parameter.""" raise NotImplementedError - def metadata(self) -> list[ParamMeta]: - """Name, dtype, full shape, and trainer placement for each parameter.""" + def __iter__(self) -> Iterator[tuple[str, Any]]: # type: ignore[override] + """Yield `(name, local buffer)` pairs in metadata order. + + Providers must enqueue materialization on the caller's current CUDA + stream, or return a buffer whose writes are already visible to it. + Buffers may be torch tensors, CUDA Array Interface objects, or raw + device pointers accepted by nccl4py. Providers own correct + communicator-device placement. Raw-pointer providers additionally own + allocation lifetime, dtype, and shape correctness. + """ raise NotImplementedError - def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: - """Yield `(name, local shard)` pairs in the same order as `metadata()`.""" - raise NotImplementedError + +@dataclass(frozen=True) +class M2NManifestEntry: + """One immutable source descriptor with a lazy rank-local tensor view.""" + + key: str + dtype: torch.dtype + global_shape: tuple[int, ...] + source_layout: M2NLayout + local_tensor: Callable[[], Any] + + def __post_init__(self) -> None: + object.__setattr__(self, "global_shape", tuple(self.global_shape)) + if not self.key: + raise ValueError("M2N manifest key must not be empty") + if not isinstance(self.dtype, torch.dtype): + raise TypeError( + f"manifest entry '{self.key}' has invalid dtype {self.dtype!r}" + ) + if not isinstance(self.source_layout, M2NLayout): + raise TypeError( + f"manifest entry '{self.key}' source_layout must be an M2NLayout" + ) + if not callable(self.local_tensor): + raise TypeError( + f"manifest entry '{self.key}' local_tensor must be callable" + ) + + @property + def metadata(self) -> M2NParamMeta: + """Return the serializable portion consumed during initialization.""" + return M2NParamMeta( + self.key, + self.dtype, + self.global_shape, + self.source_layout, + ) + + +class ManifestM2NWeightSource(M2NWeightSource): + """An ordered M2N source backed by lazy, rank-local manifest entries.""" + + def __init__(self, entries: Sequence[M2NManifestEntry]) -> None: + self._entries = tuple(entries) + self._metadata = tuple(entry.metadata for entry in self._entries) + + def metadata(self) -> list[M2NParamMeta]: # type: ignore[override] + """Return a fresh list while preserving the immutable manifest order.""" + return list(self._metadata) + + def __iter__(self) -> Iterator[tuple[str, Any]]: # type: ignore[override] + """Materialize one rank-local buffer at a time.""" + for entry in self._entries: + yield entry.key, entry.local_tensor() def _placement_code(placement: Any) -> int: @@ -69,7 +131,8 @@ def _placement_code(placement: Any) -> int: def mesh_from_tensor(tensor: torch.Tensor, num_trainer_ranks: int) -> M2NMesh: """The mesh a parameter lives on, as an `M2NMesh`. - A plain tensor is identical on every trainer rank, so it spans all of them. + Tensors without DeviceMesh metadata are treated as replicated. Implicitly + sharded tensors require a custom source with an explicit layout. """ device_mesh = getattr(tensor, "device_mesh", None) if device_mesh is None: @@ -122,29 +185,18 @@ class DTensorModuleSource(M2NWeightSource): self._module = module self._num_trainer_ranks = num_trainer_ranks - def mesh(self) -> M2NMesh: - """The mesh every sharded parameter agrees on. - - Replicated parameters span the trainer without a device mesh of their - own, so they never decide this; a model whose *sharded* parameters - disagree cannot be described by one side mesh and is rejected. - """ - meshes = { - mesh_from_tensor(p, self._num_trainer_ranks) - for _, p in self._module.named_parameters() - if placements_from_tensor(p) is not REPLICATED - } - if len(meshes) > 1: - raise ValueError( - "nccl_m2n needs one mesh per side, but this module's sharded " - f"parameters span several: {sorted(m.dims for m in meshes)}" - ) - return meshes.pop() if meshes else M2NMesh((self._num_trainer_ranks, 1), 0) - - def metadata(self) -> list[ParamMeta]: - """Read global shape/dtype and placements without gathering shards.""" + def metadata(self) -> list[M2NParamMeta]: # type: ignore[override] + """Read each global shape and source layout without gathering shards.""" return [ - M2NParamMeta(name, p.dtype, tuple(p.shape), placements_from_tensor(p)) + M2NParamMeta( + name, + p.dtype, + tuple(p.shape), + M2NLayout( + mesh_from_tensor(p, self._num_trainer_ranks), + placements_from_tensor(p), + ), + ) for name, p in self._module.named_parameters() ] diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py index 8b9b73f9f9..19dd973c73 100644 --- a/vllm/distributed/weight_transfer/m2n_trainer.py +++ b/vllm/distributed/weight_transfer/m2n_trainer.py @@ -13,6 +13,8 @@ engine -- where the `TrainerWeightTransferEngine` contract puts it -- rather tha in every caller. """ +import socket +from collections.abc import Mapping from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Self, cast @@ -26,39 +28,43 @@ from vllm.distributed.weight_transfer.base import ( WeightSource, ) from vllm.distributed.weight_transfer.m2n_common import ( + M2N_WIRE_SCHEMA_VERSION, M2NMesh, M2NParamMeta, + M2NWireParam, Placements, + check_complete_plan_agreement, + check_data_plane_agreement, + check_runtime_ready_agreement, + check_source_plan_agreement, check_transferable, - comm_ptr, + create_m2n_communicator, + expected_local_shape, import_m2n, + prepare_m2n_local_runtime, publish_destination_placements, resolve_layout, + source_plan_digest, to_mesh, to_placements, - validate_layout, + validate_local_buffer, + validate_transfer_plan, ) from vllm.distributed.weight_transfer.m2n_source import M2NWeightSource from vllm.distributed.weight_transfer.nccl_common import ( decode_nccl_unique_id, - trainer_init, - uid_init_process_group, + stateless_init_metadata_group, ) from vllm.logger import init_logger if TYPE_CHECKING: - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup logger = init_logger(__name__) __all__ = ["M2NTrainerInitInfo", "M2NTrainerWeightTransferEngine"] -def _dtype_name(dtype: torch.dtype) -> str: - """Short name for a dtype, e.g. `bfloat16`, for the worker-side init dict.""" - return str(dtype).split(".")[-1] - - @dataclass(kw_only=True) class M2NTrainerInitInfo(TrainerInitInfo): """Trainer-side init info for nccl_m2n. @@ -70,8 +76,8 @@ class M2NTrainerInitInfo(TrainerInitInfo): backend: ClassVar[str] = "nccl_m2n" - master_address: str | None = None - master_port: int | None = None + master_address: str + master_port: int world_size: int """Trainer ranks + all inference workers.""" num_trainer_ranks: int = 1 @@ -79,20 +85,34 @@ class M2NTrainerInitInfo(TrainerInitInfo): """How the inference ranks are laid out: axis 0 replicates and axis 1 shards. The trainer declares it so both sides describe the destination identically; defaults to a flat `(num_workers, 1)`, which is all a - replicated destination needs.""" - max_cta: int | None = None + replicated destination needs. + + Destination meshes remain common to every parameter; per-parameter + destination meshes are intentionally future work. + """ nccl_unique_id_b64: str | None = field(default=None, repr=False) + """Optional UID for the NCCL data plane; TCP remains the metadata plane.""" + max_cta: int | None = None + listen_socket: socket.socket | None = field(default=None, repr=False, compare=False) + """Optional rank-zero TCPStore reservation; never sent to workers.""" @property def destination_mesh_dims(self) -> tuple[int, int]: """`dst_mesh_dims`, or a flat mesh over every inference worker.""" if self.dst_mesh_dims is not None: - return self.dst_mesh_dims + return cast(tuple[int, int], tuple(self.dst_mesh_dims)) return (self.world_size - self.num_trainer_ranks, 1) def __post_init__(self) -> None: """Reject rank counts or destination meshes that cannot form a group.""" - _ = self.nccl_unique_id_bytes + if not isinstance(self.master_address, str) or not self.master_address: + raise ValueError("`master_address` must name the metadata rendezvous") + if ( + not isinstance(self.master_port, int) + or isinstance(self.master_port, bool) + or not 0 < self.master_port < 65536 + ): + raise ValueError("`master_port` must be a valid metadata port") if not 0 < self.num_trainer_ranks < self.world_size: raise ValueError( f"`num_trainer_ranks` ({self.num_trainer_ranks}) must leave at " @@ -101,7 +121,11 @@ class M2NTrainerInitInfo(TrainerInitInfo): num_workers = self.world_size - self.num_trainer_ranks if self.dst_mesh_dims is not None: dims = tuple(self.dst_mesh_dims) - if len(dims) != 2 or dims[0] * dims[1] != num_workers: + try: + mesh = M2NMesh(cast(tuple[int, int], dims), self.num_trainer_ranks) + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid `dst_mesh_dims` {dims}: {exc}") from exc + if mesh.size != num_workers: raise ValueError( f"`dst_mesh_dims` {dims} must be two dims covering the " f"{num_workers} inference workers" @@ -111,14 +135,23 @@ class M2NTrainerInitInfo(TrainerInitInfo): f"trainer `rank` ({self.rank}) must be below " f"`num_trainer_ranks` ({self.num_trainer_ranks})" ) + if self.listen_socket is not None: + if self.rank != 0: + raise ValueError("only trainer rank zero may reserve the TCPStore") + if not isinstance(self.listen_socket, socket.socket): + raise TypeError("listen_socket must be a socket.socket") + if self.listen_socket.fileno() < 0: + raise ValueError("listen_socket is already closed") @property def nccl_unique_id_bytes(self) -> bytes | None: + if self.nccl_unique_id_b64 is None: + return None return decode_nccl_unique_id( - master_address=self.master_address, - master_port=self.master_port, + master_address=None, + master_port=None, nccl_unique_id_b64=self.nccl_unique_id_b64, - ctx="M2NTrainerInitInfo", + ctx="M2NTrainerInitInfo data plane", ) @@ -136,22 +169,20 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI *, client: VLLMWeightSyncClient, source: M2NWeightSource, - is_controller: bool, - num_trainer_ranks: int, + is_sender: bool, ) -> None: """Hold the client and source; `trainer_init` fills in the group.""" - super().__init__(client=client, source=source, is_sender=is_controller) - self.source: M2NWeightSource = source - self.is_controller = is_controller - self.num_trainer_ranks = num_trainer_ranks - self.group: PyNcclCommunicator | None = None + super().__init__(client=client, source=source, is_sender=is_sender) + self.group: Any = None self._m2n: Any = None self._handle: Any = None self._metas: list[M2NParamMeta] = [] - self._src_mesh: M2NMesh | None = None + self._wire_params: list[M2NWireParam] = [] + self._source_digest = "" self._dst_mesh: M2NMesh | None = None self._dst_placements: list[Placements | None] = [] - self._executor: ThreadPoolExecutor | None = None + self._metadata_group: StatelessProcessGroup | None = None + self._num_trainer_ranks = 0 @classmethod def trainer_init( @@ -168,204 +199,351 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI build local state and join the communicator. Every rank participates in every reshard and sends its local shard. - The ordering here is not arbitrary -- see the comments inline. In - short: validate before anything blocks; start the init RPC without - waiting (workers join the communicator inside that RPC, so it cannot - return until this side has joined too); then join the NCCL - communicator; then create the M2N handle; and only then wait for the RPC - to finish. + Build a local validation envelope, start the worker RPC without + waiting, join the metadata rendezvous, and require unanimous source, + data-plane, and runtime readiness before constructing NCCL. """ - # The base WeightSource carries name/dtype/shape but not layout, which - # is not enough to plan a reshard. Fail here rather than at the first - # missing attribute deep inside the send loop. - if not isinstance(source, M2NWeightSource): - raise TypeError( - "nccl_m2n needs per-parameter layouts, so its source must be an " - f"M2NWeightSource (e.g. DTensorModuleSource); got " - f"{type(source).__name__}" - ) - - # Trainer-side weight transfer engine for this rank. `__init__` only - # holds the client and source; the communicator, meshes, and M2N handle - # are filled in below. engine = cls( client=client, - source=source, - is_controller=init_info.rank == 0, - num_trainer_ranks=init_info.num_trainer_ranks, + source=cast(M2NWeightSource, source), + is_sender=init_info.is_sender, ) - engine._m2n = import_m2n() - - engine._src_mesh = source.mesh() - if engine._src_mesh.size != init_info.num_trainer_ranks: - raise ValueError( - f"source mesh covers {engine._src_mesh.size} ranks, but there " - f"are {init_info.num_trainer_ranks} trainer ranks" - ) - - # Validate the whole plan before any rendezvous. A layout m2n cannot - # express is a mismatched collective at send time, which hangs rather - # than raises -- so every precondition is checked while an exception can - # still propagate to the caller. - engine._metas = cast(list[M2NParamMeta], source.metadata()) - for meta in engine._metas: - check_transferable(meta.name, meta.dtype, meta.shape) - mesh, placements = resolve_layout(engine._src_mesh, meta.placements) - validate_layout(mesh, placements, meta.shape, "source") - - rendezvous: Future | None = None - if engine.is_controller: - engine._executor = ThreadPoolExecutor( - max_workers=1, thread_name_prefix="m2n-weight-sync" - ) - # Only trainer rank 0 starts the control-plane RPC. Workers join - # the communicator inside it, so submit it first, join below, and - # wait for it at the end. - rendezvous = engine._executor.submit( - client.init_weight_transfer_engine, - engine._worker_init_info(init_info), - ) - - # Every trainer rank joins -- reshard is collective over the whole - # communicator. Trainer ranks take - # [0, num_trainer_ranks) and the workers follow, which is the - # contiguous-interval layout m2n requires of both meshes. - unique_id_bytes = init_info.nccl_unique_id_bytes - if unique_id_bytes is not None: - engine.group = uid_init_process_group( - unique_id_bytes, - rank=init_info.rank, - world_size=init_info.world_size, - device=torch.accelerator.current_device_index(), - ) - else: - assert init_info.master_address is not None - assert init_info.master_port is not None - engine.group = trainer_init( - { - "master_address": init_info.master_address, - "master_port": init_info.master_port, - "world_size": init_info.world_size, - }, - rank=init_info.rank, - ) - - # The trainer declares the inference topology; the worker uses exactly - # this, so neither side has to infer the other's factorization. - engine._dst_mesh = M2NMesh( - init_info.destination_mesh_dims, init_info.num_trainer_ranks - ) - engine._dst_placements = engine._receive_destination_plan( - init_info.num_trainer_ranks - ) - logger.info( - "nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), " - "inference ranks [%d, %d)", - len(engine._metas), - init_info.num_trainer_ranks, - init_info.num_trainer_ranks, - init_info.world_size, - ) - - # Created after the group so a failed rendezvous does not leave a live - # handle behind; shutdown() releases it once caller streams are synced. - engine._handle = engine._m2n.Handle.create( - engine._m2n.Config(max_cta=init_info.max_cta) - ) - - # Now safe to collect: both sides have joined, so the RPC has returned - # or raised. Doing this last means a worker-side init failure surfaces - # here, as an exception from trainer_init, rather than at the first send. - if rendezvous is not None: - rendezvous.result() + engine._initialize(init_info, source) return engine - def _receive_destination_plan( - self, first_worker_rank: int - ) -> list[Placements | None]: - """Receive and validate the inference-side per-parameter placements.""" - if self.group is None or self._dst_mesh is None: - raise RuntimeError( - "destination plan requires the communicator and destination mesh" + def _initialize( + self, + init_info: M2NTrainerInitInfo, + source: WeightSource | None, + ) -> None: + """Run source validation, metadata agreement, and NCCL rendezvous.""" + self._num_trainer_ranks = init_info.num_trainer_ranks + local_source_error: str | None = None + try: + if not isinstance(source, M2NWeightSource): + raise TypeError( + "nccl_m2n needs per-parameter layouts, so its source must " + f"be an M2NWeightSource; got {type(source).__name__}" + ) + self._m2n = import_m2n() + self._prepare_source_plan(source, init_info.num_trainer_ranks) + except Exception as exc: + local_source_error = f"{type(exc).__name__}: {exc}" + self._metas = [] + self._wire_params = [] + self._source_digest = source_plan_digest([]) + + executor: ThreadPoolExecutor | None = None + init_future: Future[None] | None = None + if self.is_sender: + executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="m2n-weight-sync" + ) + init_future = executor.submit( + self.client.init_weight_transfer_engine, + self._worker_init_info(init_info), + ) + handle = None + group = None + try: + reservation = init_info.listen_socket + try: + metadata_group = stateless_init_metadata_group( + init_info.master_address, + init_info.master_port, + init_info.rank, + init_info.world_size, + listen_socket=reservation, + ) + finally: + init_info.listen_socket = None + if reservation is not None: + reservation.close() + plan_digest = check_source_plan_agreement( + metadata_group, + self._wire_params, + expected_digest=self._source_digest, + local_error=local_source_error, ) + # Retain the common destination mesh and the merged NCCL publication + # of per-parameter placements. + self._dst_mesh = M2NMesh( + init_info.destination_mesh_dims, init_info.num_trainer_ranks + ) + unique_id_bytes: bytes | None = None + data_plane_error: str | None = None + try: + unique_id_bytes = init_info.nccl_unique_id_bytes + except Exception as exc: + data_plane_error = f"{type(exc).__name__}: {exc}" + check_data_plane_agreement( + metadata_group, + unique_id_bytes, + data_plane_error, + ) + + runtime_error: str | None = None + nccl = None + device = None + try: + nccl, device, handle = prepare_m2n_local_runtime( + self._m2n, init_info.max_cta + ) + except Exception as exc: + runtime_error = f"{type(exc).__name__}: {exc}" + check_runtime_ready_agreement(metadata_group, runtime_error) + assert nccl is not None and device is not None and handle is not None + group = create_m2n_communicator(nccl, metadata_group, unique_id_bytes) + destination_placements = self._receive_destination_plan( + group, + init_info.num_trainer_ranks, + torch.device("cuda", device), + ) + + complete_plan_error: str | None = None + try: + assert self._dst_mesh is not None + validate_transfer_plan( + self._metas, self._dst_mesh, destination_placements + ) + except Exception as exc: + complete_plan_error = f"{type(exc).__name__}: {exc}" + assert self._dst_mesh is not None + check_complete_plan_agreement( + metadata_group, + plan_digest, + self._dst_mesh, + destination_placements, + complete_plan_error, + ) + + if init_future is not None: + init_future.result() + + logger.info( + "nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), " + "inference ranks [%d, %d)", + len(self._metas), + init_info.num_trainer_ranks, + init_info.num_trainer_ranks, + init_info.world_size, + ) + logger.info("nccl_m2n source plan digest: %s", plan_digest) + except BaseException: + if group is not None: + try: + # Transport cleanup only; reload validity belongs to the + # model-update lifecycle, not this backend. + group.abort() + except BaseException: + logger.exception( + "failed to abort partially initialized nccl_m2n communicator" + ) + if handle is not None: + try: + handle.destroy() + except BaseException: + logger.exception( + "failed to destroy partially initialized nccl_m2n handle" + ) + raise + finally: + if executor is not None: + executor.shutdown(wait=False, cancel_futures=True) + + self._metadata_group = metadata_group + self._handle = handle + self.group = group + self._dst_placements = destination_placements + + def _receive_destination_plan( + self, group: Any, first_worker_rank: int, device: torch.device + ) -> list[Placements | None]: + """Receive the inference-side per-parameter placements.""" # The first inference worker publishes one placement per parameter. # Every trainer rank receives the complete plan in this single, # initialization-time broadcast. placements = publish_destination_placements( - self.group, first_worker_rank, None, len(self._metas) + group, + device, + torch.cuda.current_stream(), + first_worker_rank, + None, + len(self._metas), ) - # The send loop matches placements to metadata by position, so reject - # an incomplete or oversized plan before entering any M2N collective. - if len(placements) != len(self._metas): - raise RuntimeError( - f"inference side planned {len(placements)} parameters, but " - f"this trainer announced {len(self._metas)}" - ) - # Check every worker-selected destination against the corresponding - # global parameter shape before the first transfer. - for meta, placement in zip(self._metas, placements): - mesh, resolved = resolve_layout(self._dst_mesh, placement) - validate_layout(mesh, resolved, meta.shape, "destination") return placements def _worker_init_info(self, init_info: M2NTrainerInitInfo) -> dict[str, Any]: """Handshake payload: rendezvous, both meshes, and the transfer plan.""" - assert self._src_mesh is not None payload: dict[str, Any] = { + "schema_version": M2N_WIRE_SCHEMA_VERSION, "master_address": init_info.master_address, "master_port": init_info.master_port, "nccl_unique_id_b64": init_info.nccl_unique_id_b64, "rank_offset": init_info.num_trainer_ranks, "world_size": init_info.world_size, - "src_mesh_dims": list(self._src_mesh.dims), "dst_mesh_dims": list(init_info.destination_mesh_dims), - "names": [m.name for m in self._metas], - "dtype_names": [_dtype_name(m.dtype) for m in self._metas], - "shapes": [list(m.shape) for m in self._metas], - "src_placements": [ - None if m.placements is None else list(m.placements) - for m in self._metas - ], + "source_digest": self._source_digest, + "params": [param.to_dict() for param in self._wire_params], "max_cta": init_info.max_cta, } return {key: value for key, value in payload.items() if value is not None} + def _prepare_source_plan( + self, source: M2NWeightSource, num_trainer_ranks: int + ) -> None: + """Build and validate this rank's immutable source plan.""" + self._metas = source.metadata() + names: set[str] = set() + for meta in self._metas: + if not isinstance(meta, M2NParamMeta): + raise TypeError( + "nccl_m2n source metadata must contain M2NParamMeta " + f"entries, got {type(meta).__name__}" + ) + if meta.name in names: + raise ValueError( + "nccl_m2n source metadata contains duplicate parameter " + f"'{meta.name}'" + ) + names.add(meta.name) + check_transferable(meta.name, meta.dtype, meta.shape) + layout = meta.source_layout + if layout.mesh.start_rank != 0: + raise ValueError( + f"parameter '{meta.name}' source mesh must start at rank 0, " + f"got {layout.mesh.start_rank}" + ) + if layout.mesh.size != num_trainer_ranks: + raise ValueError( + f"parameter '{meta.name}' source mesh covers " + f"{layout.mesh.size} ranks, but there are " + f"{num_trainer_ranks} trainer ranks" + ) + expected_local_shape(meta) + self._wire_params = [ + M2NWireParam( + name=meta.name, + dtype_name=str(meta.dtype).split(".")[-1], + shape=meta.shape, + src_mesh_dims=meta.source_layout.mesh.dims, + src_placements=meta.source_layout.placements, + ) + for meta in self._metas + ] + self._source_digest = source_plan_digest(self._wire_params) + def send_weights(self) -> None: """Drive one update round: start, reshard concurrently, then finish.""" if self._handle is None or self.group is None: raise RuntimeError("nccl_m2n trainer engine is not initialized") - update: Future | None = None - if self.is_controller: - self.client.start_weight_update() - # The workers reshard from inside `update_weights`, concurrently - # with the sends below. - assert self._executor is not None - update = self._executor.submit( - self.client.update_weights, - {"names": [m.name for m in self._metas]}, - ) + if not self.is_sender: + try: + self._send() + except BaseException: + try: + self._receive_round_outcome() + except BaseException: + logger.exception( + "failed to consume nccl_m2n round outcome after local failure" + ) + raise + self._receive_round_outcome() + return try: - self._send() - finally: - if update is not None: + self.client.start_weight_update() + executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="m2n-weight-sync" + ) + try: + update = executor.submit( + self.client.update_weights, + {"names": [m.name for m in self._metas]}, + ) + # Surface an immediate worker-side rejection before trainers + # enter the collective sequence. + if update.done(): + update.result() + self._send() update.result() - - if self.is_controller: + finally: + executor.shutdown(wait=False, cancel_futures=True) self.client.finish_weight_update() + except BaseException as exc: + try: + self._publish_round_outcome(exc) + except BaseException: + logger.exception("failed to publish nccl_m2n round failure") + raise + + self._publish_round_outcome(None) + + def _publish_round_outcome(self, failure: BaseException | None) -> None: + """Tell non-sender trainer ranks whether worker finalization passed.""" + if self._num_trainer_ranks <= 1: + return + if self._metadata_group is None: + raise RuntimeError("nccl_m2n metadata group is unavailable") + envelope = { + "ok": failure is None, + "error": ( + None if failure is None else f"{type(failure).__name__}: {failure}" + ), + } + for rank in range(1, self._num_trainer_ranks): + self._metadata_group.send_obj(envelope, rank) + # Keep rank zero's TCPStore alive until every trainer consumed its + # outcome. This also keeps the per-rank message counters aligned. + expected_ack = {"phase": "round_outcome_consumed"} + for rank in range(1, self._num_trainer_ranks): + ack = self._metadata_group.recv_obj(rank) + if ack != expected_ack: + raise RuntimeError( + f"trainer rank {rank} sent an invalid round-outcome acknowledgement" + ) + + def _receive_round_outcome(self) -> None: + """Wait until trainer rank zero observed every worker's finish RPC.""" + if self._metadata_group is None: + raise RuntimeError("nccl_m2n metadata group is unavailable") + outcome = self._metadata_group.recv_obj(0) + self._metadata_group.send_obj({"phase": "round_outcome_consumed"}, 0) + if not isinstance(outcome, Mapping) or set(outcome) != {"ok", "error"}: + raise RuntimeError("trainer rank zero sent an invalid round outcome") + ok = outcome["ok"] + error = outcome["error"] + if not isinstance(ok, bool) or ok != (error is None): + raise RuntimeError("trainer rank zero sent an invalid round outcome") + if not ok: + if not isinstance(error, str) or not error: + raise RuntimeError("trainer rank zero sent an invalid round error") + raise RuntimeError( + "nccl_m2n worker update failed on trainer rank zero: " + error + ) def _send(self) -> None: """Reshard each local shard into its worker-planned destination.""" assert self.group is not None - assert self._src_mesh is not None assert self._dst_mesh is not None + assert self.source is not None m2n = self._m2n - comm = comm_ptr(self.group) stream = torch.cuda.current_stream() - for index, (name, tensor) in enumerate(self.source): + device_index = torch.accelerator.current_device_index() + # Optimizer work may have used another CUDA stream. Make prior writes + # visible before a provider materializes current-stream local shards. + torch.accelerator.synchronize(device_index) + inflight_buffers: list[Any] = [] + sent = 0 + for index, (name, buffer) in enumerate(self.source): + if index >= len(self._metas): + raise RuntimeError( + "weight source yielded more parameters than metadata() " + f"declared (first extra parameter: '{name}')" + ) meta = self._metas[index] if name != meta.name: raise RuntimeError( @@ -373,28 +551,47 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI f"its metadata declared '{meta.name}'; the source must " "iterate in a stable order" ) - shard = tensor.detach().contiguous() - src_mesh, src_placements = resolve_layout(self._src_mesh, meta.placements) + validate_local_buffer(m2n, meta, buffer) + src_mesh, src_placements = resolve_layout( + meta.source_layout.mesh, meta.source_layout.placements + ) dst_mesh, dst_placements = resolve_layout( self._dst_mesh, self._dst_placements[index] ) + raw_pointer_args = {} + if isinstance(buffer, int): + raw_pointer_args = { + "src_local_shape": expected_local_shape(meta), + "src_dtype": meta.dtype, + } m2n.reshard( - shard, + buffer, None, - comm, + self.group, stream, src_mesh=to_mesh(m2n, src_mesh), src_placements=to_placements(m2n, src_placements), dst_mesh=to_mesh(m2n, dst_mesh), dst_placements=to_placements(m2n, dst_placements), handle=self._handle, + **raw_pointer_args, + ) + # Let allocator-aware buffers recycle storage in stream order. + # Other object owners must stay alive until the final synchronize; + # raw-pointer lifetime remains the source's responsibility. + record_stream = getattr(buffer, "record_stream", None) + if callable(record_stream): + record_stream(stream) + elif not isinstance(buffer, int): + inflight_buffers.append(buffer) + sent += 1 + + if sent != len(self._metas): + missing = self._metas[sent].name + raise RuntimeError( + f"weight source yielded {sent} parameters, but metadata() " + f"declared {len(self._metas)} (first missing: '{missing}')" ) - # A custom source may yield temporary storage allocated on another - # stream. M2N retains only its device pointer for asynchronous work, - # so tell the allocator not to reuse that storage until this stream - # has finished with it. For storage allocated on this stream, - # record_stream is effectively a no-op. - shard.record_stream(stream) # As in the other weight-transfer backends, wait before returning # rather than relying on same-stream ordering. The caller may then @@ -402,14 +599,30 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI stream.synchronize() def shutdown(self) -> None: - """Destroy the m2n handle, join the helper thread, and drop the group.""" - if self._handle is not None: - torch.accelerator.synchronize() - self._handle.destroy() - self._handle = None - if self._executor is not None: - self._executor.shutdown() - self._executor = None - if self.group is not None: - self.group.destroy() - self.group = None + """Finish pending GPU work and release M2N communication state.""" + handle, self._handle = self._handle, None + group, self.group = self.group, None + self._metadata_group = None + self._dst_placements = [] + self._dst_mesh = None + self._m2n = None + + first_error: BaseException | None = None + + def attempt(action: Any, message: str) -> None: + nonlocal first_error + try: + action() + except BaseException as exc: + if first_error is None: + first_error = exc + logger.exception(message) + + if handle is not None or group is not None: + attempt(torch.accelerator.synchronize, "failed to synchronize nccl_m2n") + if handle is not None: + attempt(handle.destroy, "failed to destroy nccl_m2n handle") + if group is not None: + attempt(group.abort, "failed to abort nccl_m2n communicator") + if first_error is not None: + raise first_error diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index 1cd026fb07..b0ac285cf1 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -8,6 +8,7 @@ The dense (`NCCLWeightTransferEngine`) and sparse sparse engine does not have to subclass the dense one. """ +import socket from dataclasses import asdict, dataclass, field from typing import TYPE_CHECKING, Protocol @@ -17,6 +18,7 @@ import torch if TYPE_CHECKING: from vllm.config.parallel import ParallelConfig from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup from vllm.distributed.device_communicators.pynccl_wrapper import ( NCCL_UNIQUE_ID_BYTES, @@ -136,6 +138,26 @@ class NCCLRendezvous(Protocol): world_size: int +def stateless_init_metadata_group( + master_address: str, + master_port: int, + rank: int, + world_size: int, + *, + listen_socket: socket.socket | None = None, +) -> "StatelessProcessGroup": + """Create only the TCPStore group used for metadata agreement.""" + from vllm.distributed.utils import StatelessProcessGroup + + return StatelessProcessGroup.create( + host=master_address, + port=master_port, + rank=rank, + world_size=world_size, + listen_socket=listen_socket, + ) + + def stateless_init_process_group( master_address: str, master_port: int, @@ -181,6 +203,29 @@ def uid_init_process_group( ) +def _worker_group_rank(rank_offset: int, parallel_config: "ParallelConfig") -> int: + """Return this worker's rank in a trainer-to-worker process group.""" + dp_rank = parallel_config.data_parallel_index + world_size_per_dp = parallel_config.world_size # TP * PP + rank_within_dp = parallel_config.rank + return rank_offset + dp_rank * world_size_per_dp + rank_within_dp + + +def worker_init_metadata_group( + init_info: NCCLWeightTransferInitInfo, + parallel_config: "ParallelConfig", +) -> "StatelessProcessGroup": + """Join the trainer-to-worker TCP metadata group without constructing NCCL.""" + assert init_info.master_address is not None + assert init_info.master_port is not None + return stateless_init_metadata_group( + init_info.master_address, + init_info.master_port, + _worker_group_rank(init_info.rank_offset, parallel_config), + init_info.world_size, + ) + + def worker_init_process_group( init_info: NCCLWeightTransferInitInfo, parallel_config: "ParallelConfig", @@ -190,15 +235,7 @@ def worker_init_process_group( Computes a unique rank for this worker across all data-parallel groups and joins the trainer via whichever rendezvous mode `init_info` carries. """ - # Calculate the global rank in the trainer-worker process group. - # Must account for data parallel to get unique ranks across all workers. - dp_rank = parallel_config.data_parallel_index - world_size_per_dp = parallel_config.world_size # TP * PP - rank_within_dp = parallel_config.rank - - # Unique rank across all DP groups - worker_rank = dp_rank * world_size_per_dp + rank_within_dp - rank = worker_rank + init_info.rank_offset + rank = _worker_group_rank(init_info.rank_offset, parallel_config) device = torch.accelerator.current_device_index() unique_id_bytes = init_info.nccl_unique_id_bytes