@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b6d2b4fc25

[RL] Add sharding-aware NCCL M2N weight transfer (#51520)

Signed-off-by: Ke Wen <kwen@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>

Ke Wen committed Sep 25, 2026 at 18:38 UTC b6d2b4fc2575e89d5be7d1c26b191735932793e7
13 files changed +2812 -6
.buildkite/test_areas/distributed.yaml
+1
@@ -373,6 +373,7 @@ steps:
373 - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
374 - pytest -v -s tests/v1/distributed/test_dbo.py
375 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
376 + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer_m2n.py
377 - pytest -v -s tests/distributed/test_packed_tensor.py
378 mirror:
379 amd:
docs/training/weight_transfer/README.md
+3 -1
@@ -35,6 +35,7 @@ engine drives on your behalf:
35 | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
36 | [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches |
37 | [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) |
38 +| [nccl_m2n](m2n.md) | NCCL M2N reshard | Trainer and inference use different sharding layouts |
39
40 ## Quickstart
41
@@ -49,7 +50,8 @@ from vllm.config import WeightTransferConfig
50
51 llm = LLM(
52 model="my-model",
52 - weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc", "sparse_nccl", "sharded_rdt"
53 + # Other backends: "ipc", "sparse_nccl", "sharded_rdt", "nccl_m2n".
54 + weight_transfer_config=WeightTransferConfig(backend="nccl"),
55 )
56 ```
57
docs/training/weight_transfer/m2n.md new
+101
@@ -0,0 +1,101 @@
1 +# NCCL M2N Engine
2 +
3 +The `nccl_m2n` weight transfer engine uses [NCCL M2N](https://github.com/NVIDIA/nccl-extensions) to **reshard** weights between two disjoint meshes of GPUs — the trainer and the inference workers — in a single collective per parameter. Unlike the broadcast-based [NCCL engine](nccl.md), the trainer sends its local shard as-is: no full tensor is ever materialized on the training side, whatever its parallelism layout.
4 +
5 +## When to Use NCCL M2N
6 +
7 +- The trainer and the inference engine use **different parallelism layouts** (e.g. trainer FSDP/EP, inference TP) and gathering full tensors on the trainer is expensive
8 +- Training and inference are on **separate GPUs**, possibly across nodes
9 +- You want weight sync off the critical path of an async RL loop
10 +
11 +## Requirements
12 +
13 +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:
14 +
15 +```bash
16 +export VLLM_NCCL_SO_PATH=/path/to/the/same/libnccl.so
17 +```
18 +
19 +The backend imports the runtime lazily, so vLLM is unaffected unless you select it.
20 +
21 +## How It Works
22 +
23 +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.
24 +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.
25 +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.
26 +
27 +### Destination resolution
28 +
29 +A parameter is resharded **directly into the live vLLM parameter** when its checkpoint name maps 1:1 onto a model parameter whose shape is either identical to the checkpoint shape (replicated) or differs on exactly one dim by the tensor-parallel factor. That rank then only ever receives its own shard.
30 +
31 +Anything else falls back to receiving the full tensor and letting `load_weights` do the sharding, exactly as the NCCL engine does for every parameter. The fallback covers:
32 +
33 +- fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`), whose checkpoint names do not name a model parameter
34 +- quantized models, where the weight loader does more than a sharded copy
35 +- pipeline-parallel deployments
36 +
37 +Correctness never depends on a parameter resolving. The engine logs how many parameters took each path at init.
38 +
39 +## Configuration
40 +
41 +```python
42 +from vllm import LLM
43 +from vllm.config import WeightTransferConfig
44 +
45 +llm = LLM(
46 + model="my-model",
47 + weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"),
48 +)
49 +```
50 +
51 +```bash
52 +vllm serve my-model --weight-transfer-config '{"backend": "nccl_m2n"}'
53 +```
54 +
55 +## Trainer Side
56 +
57 +The trainer uses the stateful trainer engine. m2n needs each parameter's source layout, which the base `WeightSource` does not carry, so it takes an `M2NWeightSource`. `DTensorModuleSource` covers the common case: it reads each parameter's device mesh and placements and yields **local shards** — it never calls `full_tensor()`.
58 +
59 +```python
60 +from vllm.distributed.weight_transfer import (
61 + RayVLLMWeightSyncClient,
62 + WeightTransferTrainerFactory,
63 +)
64 +from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource
65 +from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
66 +
67 +# Called on every trainer rank; rank 0 is the sender.
68 +engine = WeightTransferTrainerFactory.trainer_init(
69 + init_info=M2NTrainerInitInfo(
70 + master_address=master_address,
71 + master_port=master_port,
72 + world_size=num_trainer_ranks + num_inference_workers,
73 + num_trainer_ranks=num_trainer_ranks,
74 + rank=my_rank,
75 + ),
76 + client=RayVLLMWeightSyncClient(llm_handle),
77 + source=DTensorModuleSource(model, num_trainer_ranks),
78 +)
79 +
80 +# Each round, on every trainer rank:
81 +engine.send_weights()
82 +```
83 +
84 +`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.
85 +
86 +Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass `M2NWeightSource` and supply layouts explicitly.
87 +
88 +See [`examples/rl/rlhf_m2n.py`](../../../examples/rl/rlhf_m2n.py) for a runnable FSDP → TP example.
89 +
90 +## Limitations
91 +
92 +- Tensor rank 1..3; 4-D parameters are rejected.
93 +- dtypes: int8/uint8, fp8 e4m3/e5m2, fp16, bf16, int32/uint32, fp32, int64/uint64, fp64. No fp4.
94 +- The trainer's device mesh must cover a contiguous rank interval starting at 0, and be 1-D or 2-D.
95 +- `Partial` placements are not supported.
96 +- Models with known model-level checkpoint preprocessing use the full-tensor
97 + fallback for the affected module.
98 +- A reshard cannot be issued from inside a CUDA graph capture (weight updates run outside capture, so this only matters for custom callers).
99 +- 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).
100 +- 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.
101 +- Each communicator holds a cached staging pool (~256 MiB by default: 4 channels x 64 MiB).
examples/rl/rlhf_m2n.py new
+197
@@ -0,0 +1,197 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""RLHF with FSDP2 training and vLLM tensor-parallel inference using **NCCL M2N**
4 +sharding-aware weight transfer.
5 +
6 +Layout (4 GPUs, no colocation):
7 + * GPUs 0-1: two FSDP2 training workers, one per GPU.
8 + * GPUs 2-3: one vLLM ``LLM`` actor with ``tensor_parallel_size=2``.
9 +
10 +The trainer and the inference workers share one NCCL communicator of 4 ranks:
11 +trainer ranks ``[0, 2)`` and inference ranks ``[2, 4)``. Every parameter moves
12 +with a single ``reshard`` that redistributes it from the FSDP layout to the
13 +inference layout — the trainer sends its local shards and never all-gathers a
14 +full tensor, which broadcast-based weight sync would force it to do.
15 +
16 +Every FSDP rank builds an ``M2NTrainerWeightTransferEngine`` and calls
17 +``send_weights()``; all ranks run every reshard, and only rank 0 drives the
18 +inference side through its ``RayVLLMWeightSyncClient``.
19 +
20 +Requires the ``nccl-extensions`` package (NCCL 2.30.5+) and a
21 +``VLLM_NCCL_SO_PATH`` pointing at the same ``libnccl.so`` that
22 +``libnccl_m2n.so`` was linked against.
23 +
24 +This example was written for 4xH100.
25 +"""
26 +
27 +from __future__ import annotations
28 +
29 +import os
30 +
31 +import pybase64 as base64
32 +import ray
33 +import torch
34 +import torch.distributed as dist
35 +from huggingface_hub import snapshot_download
36 +from torch.distributed.fsdp import fully_shard
37 +from transformers import AutoModelForCausalLM
38 +
39 +from vllm import LLM, SamplingParams
40 +from vllm.config import WeightTransferConfig
41 +from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
42 +from vllm.distributed.weight_transfer import (
43 + RayVLLMWeightSyncClient,
44 + WeightTransferTrainerFactory,
45 +)
46 +from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource
47 +from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
48 +from vllm.utils.network_utils import get_ip, get_open_port
49 +
50 +MODEL_NAME = "Qwen/Qwen3-0.6B"
51 +
52 +FSDP_WORLD_SIZE = 2
53 +INFERENCE_TP_SIZE = 2
54 +
55 +
56 +class MyLLM(LLM):
57 + """LLM subclass that keeps Ray from pinning it to a single device."""
58 +
59 + def __init__(self, *args, **kwargs):
60 + os.environ.pop("CUDA_VISIBLE_DEVICES", None)
61 + super().__init__(*args, **kwargs)
62 +
63 + def ready(self):
64 + return True
65 +
66 +
67 +@ray.remote(num_cpus=1, num_gpus=1)
68 +class FSDPTrainWorker:
69 + """One FSDP2 worker per GPU. Rank 0 is the weight-transfer sender."""
70 +
71 + def __init__(
72 + self,
73 + model_name: str,
74 + rank: int,
75 + fsdp_master_addr: str,
76 + fsdp_master_port: int,
77 + ):
78 + self.rank = rank
79 +
80 + os.environ["MASTER_ADDR"] = fsdp_master_addr
81 + os.environ["MASTER_PORT"] = str(fsdp_master_port)
82 +
83 + dist.init_process_group(backend="nccl", rank=rank, world_size=FSDP_WORLD_SIZE)
84 + torch.accelerator.set_device_index(0)
85 +
86 + model = AutoModelForCausalLM.from_pretrained(
87 + model_name, torch_dtype=torch.bfloat16
88 + ).cuda()
89 + for layer in model.model.layers:
90 + fully_shard(layer)
91 + fully_shard(model)
92 + self.model = model
93 +
94 + def ready(self):
95 + return True
96 +
97 + def setup_engine(self, llm_handle, nccl_unique_id_b64, world_size, num_workers):
98 + """Build the trainer engine on every FSDP rank.
99 +
100 + `DTensorModuleSource` reads each parameter's FSDP device mesh and
101 + placements, so the engine knows the source layout without gathering
102 + anything. Rank 0 additionally drives the inference-side handshake.
103 + """
104 + self.engine = WeightTransferTrainerFactory.trainer_init(
105 + init_info=M2NTrainerInitInfo(
106 + nccl_unique_id_b64=nccl_unique_id_b64,
107 + world_size=world_size,
108 + num_trainer_ranks=FSDP_WORLD_SIZE,
109 + # One DP group of TP=2 workers; declared so both sides
110 + # describe the destination the same way.
111 + dst_mesh_dims=(num_workers // INFERENCE_TP_SIZE, INFERENCE_TP_SIZE),
112 + rank=self.rank, # FSDP rank; sender is 0
113 + ),
114 + client=RayVLLMWeightSyncClient(llm_handle),
115 + source=DTensorModuleSource(self.model, FSDP_WORLD_SIZE),
116 + )
117 +
118 + def send_weights(self):
119 + """Called on all ranks concurrently; every rank runs every reshard."""
120 + self.engine.send_weights()
121 +
122 +
123 +def main():
124 + ray.init()
125 +
126 + local_model_path = snapshot_download(MODEL_NAME)
127 + print(f"[init] Model downloaded to {local_model_path}")
128 +
129 + fsdp_master_addr = get_ip()
130 + fsdp_master_port = get_open_port()
131 +
132 + fsdp_workers = [
133 + FSDPTrainWorker.remote(
134 + local_model_path, rank, fsdp_master_addr, fsdp_master_port
135 + )
136 + for rank in range(FSDP_WORLD_SIZE)
137 + ]
138 + ray.get([w.ready.remote() for w in fsdp_workers])
139 + print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")
140 +
141 + llm = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote(
142 + model=local_model_path,
143 + enforce_eager=True,
144 + tensor_parallel_size=INFERENCE_TP_SIZE,
145 + distributed_executor_backend="ray",
146 + weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"),
147 + load_format="dummy",
148 + )
149 + ray.get(llm.ready.remote())
150 + num_workers = ray.get(llm.get_world_size.remote())
151 + print(f"[init] vLLM ready with {num_workers} inference workers.")
152 +
153 + prompts = [
154 + "Hello, my name is",
155 + "The capital of France is",
156 + ]
157 + sampling_params = SamplingParams(temperature=0)
158 +
159 + outputs = ray.get(llm.generate.remote(prompts, sampling_params))
160 + print("-" * 60)
161 + print("BEFORE weight sync (dummy weights):")
162 + for output in outputs:
163 + print(f"Prompt: {output.prompt!r}")
164 + print(f"Generated: {output.outputs[0].text!r}")
165 + print("-" * 60)
166 +
167 + # Trainer ranks [0, FSDP_WORLD_SIZE) and inference ranks after them share
168 + # one communicator, so the two meshes are contiguous rank intervals.
169 + nccl = NCCLLibrary()
170 + nccl_unique_id_b64 = base64.b64encode(
171 + bytes(nccl.ncclGetUniqueId().internal)
172 + ).decode()
173 + world_size = FSDP_WORLD_SIZE + num_workers
174 +
175 + print("[transfer] Initializing nccl_m2n weight transfer (all FSDP ranks)...")
176 + ray.get(
177 + [
178 + w.setup_engine.remote(llm, nccl_unique_id_b64, world_size, num_workers)
179 + for w in fsdp_workers
180 + ]
181 + )
182 +
183 + print("[sync] Resharding FSDP -> vLLM...")
184 + ray.get([w.send_weights.remote() for w in fsdp_workers])
185 + print("[sync] Weight transfer complete.")
186 +
187 + outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
188 + print("-" * 60)
189 + print("AFTER weight sync (real weights):")
190 + for output in outputs_updated:
191 + print(f"Prompt: {output.prompt!r}")
192 + print(f"Generated: {output.outputs[0].text!r}")
193 + print("-" * 60)
194 +
195 +
196 +if __name__ == "__main__":
197 + main()
tests/distributed/test_weight_transfer_m2n.py new
+868
@@ -0,0 +1,868 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the NCCL M2N weight transfer backend.
4 +
5 +These cover how a transfer is *described* — layout encoding and the wire-type
6 +validation that keeps a bad plan from reaching a collective, since a mismatch
7 +there is a hang rather than an error. The transfer itself needs the m2n runtime
8 +and multiple GPUs, so it is exercised separately.
9 +"""
10 +
11 +import logging
12 +from functools import wraps
13 +from types import MethodType
14 +from unittest.mock import Mock
15 +
16 +import pybase64 as base64
17 +import pytest
18 +import ray
19 +import torch
20 +
21 +from vllm.distributed.weight_transfer import (
22 + WeightTransferEngineFactory,
23 + WeightTransferTrainerFactory,
24 +)
25 +from vllm.distributed.weight_transfer.m2n_common import (
26 + REPLICATE,
27 + REPLICATED,
28 + M2NMesh,
29 + M2NParamMeta,
30 + check_placements,
31 + check_transferable,
32 + publish_destination_placements,
33 + resolve_layout,
34 + validate_layout,
35 +)
36 +from vllm.distributed.weight_transfer.m2n_engine import (
37 + M2NWeightTransferEngine,
38 + M2NWeightTransferInitInfo,
39 + M2NWeightTransferUpdateInfo,
40 +)
41 +from vllm.distributed.weight_transfer.m2n_layout import (
42 + M2NDestination,
43 + resolve_parameter_destinations,
44 +)
45 +from vllm.distributed.weight_transfer.m2n_source import (
46 + mesh_from_tensor,
47 + placements_from_tensor,
48 +)
49 +from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
50 +from vllm.model_executor.layers.linear import (
51 + ColumnParallelLinear,
52 + RowParallelLinear,
53 +)
54 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
55 +from vllm.platforms import current_platform
56 +
57 +VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode()
58 +
59 +
60 +class TestLayout:
61 + def test_replicated_splits_nothing(self):
62 + """A replicated tensor imposes no divisibility constraint. The shape
63 + dims are deliberately coprime with the mesh size: a layout that actually
64 + split the tensor would reject them."""
65 + mesh = M2NMesh((2, 2), start_rank=1)
66 + indivisible_shape = (7, 13) # neither dim divisible by any mesh axis
67 +
68 + resolved_mesh, placements = resolve_layout(mesh, REPLICATED)
69 + validate_layout(resolved_mesh, placements, indivisible_shape, "destination")
70 +
71 + def test_replicated_keeps_the_same_ranks(self):
72 + """Replication re-factors the mesh to get a size-1 axis for its no-op
73 + shard. That is only sound if it still covers exactly the same GPUs."""
74 + mesh = M2NMesh((2, 3), start_rank=4)
75 + resolved_mesh, _ = resolve_layout(mesh, REPLICATED)
76 + assert resolved_mesh.size == mesh.size
77 + assert resolved_mesh.start_rank == mesh.start_rank
78 +
79 + def test_sharded_keeps_its_own_factorization(self):
80 + """Rank order decides who owns which shard, so a sharded tensor must
81 + not be re-factored the way a replicated one is."""
82 + mesh = M2NMesh((2, 3), start_rank=4)
83 + resolved_mesh, placements = resolve_layout(mesh, (REPLICATE, 0))
84 + assert resolved_mesh == mesh
85 + assert placements == (REPLICATE, 0)
86 +
87 + def test_two_shard_axes_rejected(self):
88 + """One axis has to replicate; a 2-D mesh that shards both is not
89 + something a single reshard can express."""
90 + with pytest.raises(ValueError, match="shards both"):
91 + check_placements((0, 1))
92 +
93 + def test_placement_below_replicate_rejected(self):
94 + with pytest.raises(ValueError, match=r"parameter 'w'.*-2"):
95 + check_placements((-2, REPLICATE), "parameter 'w' source placements")
96 +
97 + def test_shard_dim_must_exist(self):
98 + with pytest.raises(ValueError, match="rank 2"):
99 + validate_layout(M2NMesh((1, 2), 0), (REPLICATE, 2), (8, 16), "source")
100 +
101 + def test_shard_must_divide_evenly(self):
102 + with pytest.raises(ValueError, match="does not divide evenly"):
103 + validate_layout(M2NMesh((1, 3), 0), (REPLICATE, 0), (8, 16), "source")
104 +
105 +
106 +class TestSourceLayout:
107 + def test_plain_tensor_is_replicated_across_trainer_ranks(self):
108 + """A tensor with no DTensor metadata is the same on every trainer rank,
109 + so the source describes it as replicated over all of them."""
110 + assert placements_from_tensor(torch.zeros(4)) is REPLICATED
111 + assert mesh_from_tensor(torch.zeros(4), 4) == M2NMesh((4, 1), 0)
112 +
113 +
114 +class TestTransferable:
115 + def test_unsupported_dtype_names_the_parameter(self):
116 + with pytest.raises(ValueError, match="'w'"):
117 + check_transferable("w", torch.complex64, (4,))
118 +
119 + def test_rank_four_rejected(self):
120 + with pytest.raises(ValueError, match="rank 4"):
121 + check_transferable("w", torch.bfloat16, (2, 2, 2, 2))
122 +
123 +
124 +class TestWireTypes:
125 + def _init_info(self, **overrides):
126 + fields = dict(
127 + master_address="127.0.0.1",
128 + master_port=1234,
129 + rank_offset=1,
130 + world_size=3,
131 + src_mesh_dims=[1, 1],
132 + dst_mesh_dims=[2, 1],
133 + names=["w"],
134 + dtype_names=["bfloat16"],
135 + shapes=[[16, 16]],
136 + src_placements=[None],
137 + )
138 + fields.update(overrides)
139 + return M2NWeightTransferInitInfo(**fields)
140 +
141 + def test_accepts_a_consistent_plan(self):
142 + assert self._init_info().names == ["w"]
143 +
144 + def test_accepts_pre_shared_nccl_unique_id(self):
145 + info = self._init_info(
146 + master_address=None,
147 + master_port=None,
148 + nccl_unique_id_b64=VALID_UID_B64,
149 + )
150 + assert info.nccl_unique_id_bytes == b"\x00" * 128
151 + assert VALID_UID_B64 not in repr(info)
152 +
153 + def test_destination_plan_uses_uid_communicator_without_group(self):
154 + comm = Mock(rank=0, device=torch.device("cpu"), group=None)
155 +
156 + def receive_plan(tensor, src):
157 + assert src == 1
158 + tensor.copy_(torch.tensor([[-1, 0], [-2, -2]], dtype=torch.int8))
159 +
160 + comm.broadcast.side_effect = receive_plan
161 +
162 + placements = publish_destination_placements(comm, 1, None, 2)
163 +
164 + assert placements == [(REPLICATE, 0), REPLICATED]
165 +
166 + def test_rejects_both_rendezvous_modes(self):
167 + with pytest.raises(ValueError, match="not both"):
168 + self._init_info(nccl_unique_id_b64=VALID_UID_B64)
169 +
170 + def test_trainer_accepts_pre_shared_nccl_unique_id(self):
171 + info = M2NTrainerInitInfo(
172 + nccl_unique_id_b64=VALID_UID_B64,
173 + world_size=3,
174 + num_trainer_ranks=1,
175 + rank=0,
176 + )
177 + assert info.nccl_unique_id_bytes == b"\x00" * 128
178 + assert VALID_UID_B64 not in repr(info)
179 +
180 + def test_ragged_plan_rejected(self):
181 + with pytest.raises(ValueError, match="`shapes`"):
182 + self._init_info(shapes=[])
183 +
184 + def test_destination_mesh_must_cover_the_workers(self):
185 + """The trainer declares the inference mesh, so one that does not cover
186 + the workers is a config error — and it has to fail the init RPC, since
187 + a mismatched mesh would otherwise surface as a hung collective."""
188 + with pytest.raises(ValueError, match="dst_mesh_dims"):
189 + self._init_info(dst_mesh_dims=[3, 1]) # 3 != the 2 workers
190 +
191 + def test_world_must_hold_a_trainer_and_a_worker(self):
192 + with pytest.raises(ValueError, match="rank_offset"):
193 + self._init_info(rank_offset=3, world_size=3)
194 +
195 + @pytest.mark.parametrize("dtype_name", ["not_a_dtype", "Tensor"])
196 + def test_invalid_dtype_name_names_parameter(self, monkeypatch, dtype_name):
197 + monkeypatch.setattr(
198 + "vllm.distributed.weight_transfer.m2n_engine.import_m2n",
199 + lambda: object(),
200 + )
201 + engine = object.__new__(M2NWeightTransferEngine)
202 +
203 + with pytest.raises(ValueError, match=r"parameter 'w'.*dtype"):
204 + engine.init_transfer_engine(self._init_info(dtype_names=[dtype_name]))
205 +
206 + def _init_32_to_4_plan(self, monkeypatch, destination):
207 + m2n = Mock()
208 + m2n.Handle.create.return_value = object()
209 + monkeypatch.setattr(
210 + "vllm.distributed.weight_transfer.m2n_engine.import_m2n", lambda: m2n
211 + )
212 + monkeypatch.setattr(
213 + "vllm.distributed.weight_transfer.m2n_engine."
214 + "resolve_parameter_destinations",
215 + lambda *args, **kwargs: [destination],
216 + )
217 + monkeypatch.setattr(
218 + "vllm.distributed.weight_transfer.m2n_engine.worker_init_process_group",
219 + lambda *args, **kwargs: object(),
220 + )
221 + monkeypatch.setattr(
222 + "vllm.distributed.weight_transfer.m2n_engine."
223 + "publish_destination_placements",
224 + lambda *args: args[2],
225 + )
226 +
227 + engine = object.__new__(M2NWeightTransferEngine)
228 + engine.parallel_config = Mock(pipeline_parallel_size=1)
229 + engine.model_config = Mock(quantization=None)
230 + engine.model = torch.nn.Module()
231 + engine.init_transfer_engine(
232 + self._init_info(
233 + rank_offset=32,
234 + world_size=36,
235 + src_mesh_dims=[1, 32],
236 + dst_mesh_dims=[1, 4],
237 + shapes=[[32, 8]],
238 + src_placements=[[REPLICATE, 0]],
239 + )
240 + )
241 +
242 + def test_limits_use_resolved_sharded_destination(self, monkeypatch):
243 + destination = M2NDestination("w", (REPLICATE, 0), torch.empty(8, 8))
244 +
245 + self._init_32_to_4_plan(monkeypatch, destination)
246 +
247 + def test_limits_still_reject_replicated_destination(self, monkeypatch):
248 + destination = M2NDestination("w", REPLICATED, None)
249 +
250 + with pytest.raises(ValueError, match=r"32 source shards.*MAX_SOURCES=16"):
251 + self._init_32_to_4_plan(monkeypatch, destination)
252 +
253 + def test_update_preflights_all_names_before_reshard(self):
254 + engine = object.__new__(M2NWeightTransferEngine)
255 + engine._handle = object()
256 + engine.model_update_group = object()
257 + engine._index = {"valid": 0}
258 + engine._metas = [M2NParamMeta("valid", torch.float32, (4,), REPLICATED)]
259 + engine._reshard = Mock()
260 +
261 + with pytest.raises(ValueError, match=r"parameter 'unknown'"):
262 + engine.receive_weights(
263 + M2NWeightTransferUpdateInfo(names=["valid", "unknown"])
264 + )
265 +
266 + engine._reshard.assert_not_called()
267 +
268 + def test_fallback_tensors_share_one_load_weights_invocation(self, monkeypatch):
269 + class CoupledModel:
270 + def __init__(self):
271 + self.calls = 0
272 + self.loaded = []
273 +
274 + def load_weights(self, weights):
275 + self.calls += 1
276 + received = list(weights)
277 + if [name for name, _ in received] == ["pair.weight", "pair.scale"]:
278 + self.loaded = [(name, tensor.clone()) for name, tensor in received]
279 +
280 + model = CoupledModel()
281 + direct = torch.zeros(1)
282 + engine = object.__new__(M2NWeightTransferEngine)
283 + engine._handle = object()
284 + engine.model_update_group = object()
285 + engine.device = torch.device("cpu")
286 + engine.model = model
287 + engine._index = {"pair.weight": 0, "direct": 1, "pair.scale": 2}
288 + engine._metas = [
289 + M2NParamMeta(name, torch.float32, (1,), REPLICATED)
290 + for name in engine._index
291 + ]
292 + engine._parameter_destinations = [
293 + M2NDestination("pair.weight", REPLICATED, None),
294 + M2NDestination("direct", REPLICATED, direct),
295 + M2NDestination("pair.scale", REPLICATED, None),
296 + ]
297 + reshard_order = []
298 +
299 + def reshard(comm, stream, meta, placements, buffer):
300 + reshard_order.append(meta.name)
301 + buffer.fill_(len(reshard_order))
302 +
303 + engine._reshard = reshard
304 + stream = Mock()
305 + monkeypatch.setattr(
306 + "vllm.distributed.weight_transfer.m2n_engine.comm_ptr", lambda group: 0
307 + )
308 + monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream)
309 +
310 + engine.receive_weights(
311 + M2NWeightTransferUpdateInfo(names=["pair.weight", "direct", "pair.scale"])
312 + )
313 +
314 + assert model.calls == 1
315 + assert [name for name, _ in model.loaded] == ["pair.weight", "pair.scale"]
316 + assert [tensor.item() for _, tensor in model.loaded] == [1, 3]
317 + assert direct.item() == 2
318 + assert reshard_order == ["pair.weight", "direct", "pair.scale"]
319 + assert stream.synchronize.call_count == 2
320 +
321 + def test_trainer_rank_must_be_a_trainer_rank(self):
322 + """A trainer rank must fall within the trainer portion of the group."""
323 + with pytest.raises(ValueError, match="num_trainer_ranks"):
324 + M2NTrainerInitInfo(
325 + master_address="127.0.0.1",
326 + master_port=1234,
327 + world_size=4,
328 + num_trainer_ranks=2,
329 + rank=2,
330 + )
331 +
332 + def test_trainer_destination_mesh_must_cover_the_workers(self):
333 + """The destination mesh must include every inference worker."""
334 + with pytest.raises(ValueError, match="dst_mesh_dims"):
335 + M2NTrainerInitInfo(
336 + master_address="127.0.0.1",
337 + master_port=1234,
338 + world_size=6,
339 + num_trainer_ranks=2,
340 + dst_mesh_dims=(3, 1), # 3 != the 4 inference workers
341 + rank=0,
342 + )
343 +
344 + def test_destination_mesh_defaults_to_flat(self):
345 + """A replicated destination does not care how the mesh is factored, so
346 + callers that do not shard it need not supply one."""
347 + info = M2NTrainerInitInfo(
348 + master_address="127.0.0.1",
349 + master_port=1234,
350 + world_size=6,
351 + num_trainer_ranks=2,
352 + rank=0,
353 + )
354 + assert info.destination_mesh_dims == (4, 1)
355 +
356 + def test_sender_is_trainer_rank_zero(self):
357 + """Trainer rank 0 drives the inference control plane."""
358 + info = M2NTrainerInitInfo(
359 + master_address="127.0.0.1", master_port=1234, world_size=4, rank=0
360 + )
361 + assert info.is_sender
362 +
363 +
364 +class TestRegistration:
365 + def test_both_registries_expose_the_backend(self):
366 + """Both worker and trainer factories register nccl_m2n."""
367 + assert "nccl_m2n" in WeightTransferEngineFactory._registry
368 + assert "nccl_m2n" in WeightTransferTrainerFactory._registry
369 +
370 + def test_init_info_dispatches_to_the_backend(self):
371 + """Trainer init info selects the nccl_m2n factory entry."""
372 + assert M2NTrainerInitInfo.backend == "nccl_m2n"
373 +
374 +
375 +# ---------------------------------------------------------------------------
376 +# End-to-end transfer
377 +#
378 +# Transport-only, in the style of the NCCL/sparse tests in
379 +# test_weight_transfer.py: two Ray tasks with one GPU each drive the two engines
380 +# directly, so no HTTP server or LLM instance is involved. The worker is not an
381 +# RPC endpoint here, so the trainer engine gets a no-op control-plane client --
382 +# the NCCL rendezvous and the reshard itself are the real thing.
383 +# ---------------------------------------------------------------------------
384 +
385 +SHAPE = [64, 32]
386 +DTYPE = "float32"
387 +
388 +
389 +def _init_ray() -> None:
390 + if ray.is_initialized():
391 + return
392 + ray.init(
393 + ignore_reinit_error=True,
394 + runtime_env={
395 + "env_vars": {
396 + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
397 + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1",
398 + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1",
399 + }
400 + },
401 + )
402 +
403 +
404 +def _assigned_device() -> "torch.device":
405 + gpu_ids = ray.get_gpu_ids()
406 + device = torch.device(f"cuda:{int(gpu_ids[0])}" if gpu_ids else "cuda:0")
407 + current_platform.set_device(device)
408 + return device
409 +
410 +
411 +@ray.remote(num_gpus=1)
412 +def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool:
413 + """Send one parameter through the real trainer engine."""
414 + device = _assigned_device()
415 +
416 + from vllm.distributed.weight_transfer import WeightTransferTrainerFactory
417 + from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource
418 + from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
419 +
420 + class NoopClient:
421 + def init_weight_transfer_engine(self, init_info):
422 + pass
423 +
424 + def start_weight_update(self):
425 + pass
426 +
427 + def update_weights(self, update_info):
428 + pass
429 +
430 + def finish_weight_update(self, weight_version=None):
431 + pass
432 +
433 + class Tiny(torch.nn.Module):
434 + def __init__(self):
435 + super().__init__()
436 + self.weight = torch.nn.Parameter(
437 + torch.arange(
438 + SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device
439 + ).reshape(SHAPE)
440 + )
441 +
442 + engine = WeightTransferTrainerFactory.trainer_init(
443 + init_info=M2NTrainerInitInfo(
444 + nccl_unique_id_b64=nccl_unique_id_b64,
445 + world_size=world_size,
446 + num_trainer_ranks=1,
447 + dst_mesh_dims=(1, world_size - 1),
448 + rank=0,
449 + ),
450 + client=NoopClient(),
451 + source=DTensorModuleSource(Tiny(), num_trainer_ranks=1),
452 + )
453 + engine.send_weights()
454 + torch.accelerator.synchronize()
455 + engine.shutdown()
456 + return True
457 +
458 +
459 +@ray.remote(num_gpus=1)
460 +def _m2n_worker_receive(
461 + nccl_unique_id_b64: str,
462 + world_size: int,
463 + worker_rank: int = 0,
464 +) -> dict:
465 + """Receive that parameter through the real worker engine."""
466 + import contextlib
467 + from unittest.mock import MagicMock
468 +
469 + device = _assigned_device()
470 +
471 + from vllm.config.parallel import ParallelConfig
472 + from vllm.config.weight_transfer import WeightTransferConfig
473 + from vllm.distributed.weight_transfer.m2n_engine import (
474 + M2NWeightTransferEngine,
475 + M2NWeightTransferInitInfo,
476 + M2NWeightTransferUpdateInfo,
477 + )
478 +
479 + class Recorder(torch.nn.Module):
480 + def __init__(self):
481 + super().__init__()
482 + self.received = []
483 + if world_size > 2:
484 + self.weight = torch.nn.Parameter(
485 + torch.zeros(
486 + SHAPE[0] // (world_size - 1),
487 + SHAPE[1],
488 + dtype=torch.float32,
489 + device=device,
490 + )
491 + )
492 + self.weight.output_dim = 0
493 + self.weight.weight_loader = MethodType(
494 + ColumnParallelLinear.weight_loader,
495 + Mock(tp_size=world_size - 1),
496 + )
497 +
498 + def load_weights(self, weights):
499 + for name, tensor in weights:
500 + self.received.append((name, tensor.clone()))
501 +
502 + parallel_config = MagicMock(spec=ParallelConfig)
503 + parallel_config.rank = worker_rank
504 + parallel_config.world_size = world_size - 1
505 + parallel_config.data_parallel_rank = 0
506 + parallel_config.data_parallel_index = 0
507 + parallel_config.tensor_parallel_size = world_size - 1
508 + parallel_config.pipeline_parallel_size = 1
509 + vllm_config = MagicMock()
510 + vllm_config.parallel_config = parallel_config
511 + vllm_config.model_config = MagicMock()
512 + vllm_config.model_config.quantization = None
513 +
514 + recorder = Recorder()
515 + engine = M2NWeightTransferEngine(
516 + WeightTransferConfig(backend="nccl_m2n"), vllm_config, device, recorder
517 + )
518 + # Transport-only: receive_weights enters set_current_vllm_config, and
519 + # vllm_config here is a mock.
520 + import vllm.config as _vllm_config_mod
521 +
522 + _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
523 +
524 + engine.init_transfer_engine(
525 + M2NWeightTransferInitInfo(
526 + nccl_unique_id_b64=nccl_unique_id_b64,
527 + rank_offset=1, # trainer occupies rank 0
528 + world_size=world_size,
529 + src_mesh_dims=[1, 1],
530 + dst_mesh_dims=[1, world_size - 1],
531 + names=["weight"],
532 + dtype_names=[DTYPE],
533 + shapes=[SHAPE],
534 + src_placements=[None], # replicated on the single trainer rank
535 + )
536 + )
537 + engine.receive_weights(M2NWeightTransferUpdateInfo(names=["weight"]))
538 + torch.accelerator.synchronize()
539 +
540 + full = torch.arange(
541 + SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device
542 + ).reshape(SHAPE)
543 + direct = world_size > 2
544 + name: str | None
545 + got: torch.Tensor | None
546 + if direct:
547 + name, got = "weight", recorder.weight
548 + expected = full.chunk(world_size - 1, dim=0)[worker_rank]
549 + else:
550 + name, got = recorder.received[0] if recorder.received else (None, None)
551 + expected = full
552 + result = {
553 + "count": len(recorder.received),
554 + "name": name,
555 + "shape": list(got.shape) if got is not None else None,
556 + "exact": bool(torch.equal(got, expected)) if got is not None else False,
557 + "direct": direct,
558 + }
559 + engine.shutdown()
560 + return result
561 +
562 +
563 +@pytest.mark.skipif(
564 + torch.accelerator.device_count() < 2,
565 + reason="Need at least 2 GPUs: one trainer rank and one inference worker.",
566 +)
567 +def test_m2n_weight_transfer_between_processes():
568 + """A parameter survives a real reshard from a trainer process to a worker.
569 +
570 + This is the only test here that moves bytes: it builds both engines, joins
571 + one NCCL communicator across two processes, and reshards. Everything else
572 + in this file checks how a transfer is *described*.
573 + """
574 + pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime")
575 + _init_ray()
576 +
577 + from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
578 +
579 + nccl = NCCLLibrary()
580 + nccl_unique_id_b64 = base64.b64encode(
581 + bytes(nccl.ncclGetUniqueId().internal)
582 + ).decode()
583 + world_size = 2 # 1 trainer + 1 inference worker
584 +
585 + worker = _m2n_worker_receive.remote(nccl_unique_id_b64, world_size)
586 + trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size)
587 + trainer_ok, result = ray.get([trainer, worker])
588 +
589 + assert trainer_ok, "trainer engine did not complete"
590 + assert result["count"] == 1, f"expected one parameter, got {result['count']}"
591 + assert result["name"] == "weight"
592 + assert result["shape"] == SHAPE
593 + assert result["exact"], "resharded tensor does not match what the trainer sent"
594 +
595 +
596 +@pytest.mark.skipif(
597 + torch.accelerator.device_count() < 3,
598 + reason="Need at least 3 GPUs: one trainer rank and two inference workers.",
599 +)
600 +def test_m2n_weight_transfer_to_tp2_shards():
601 + """Two inference workers receive their own TP shard in live model storage."""
602 + pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime")
603 + _init_ray()
604 +
605 + from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
606 +
607 + nccl = NCCLLibrary()
608 + nccl_unique_id_b64 = base64.b64encode(
609 + bytes(nccl.ncclGetUniqueId().internal)
610 + ).decode()
611 + world_size = 3
612 +
613 + workers = [
614 + _m2n_worker_receive.remote(nccl_unique_id_b64, world_size, rank)
615 + for rank in range(2)
616 + ]
617 + trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size)
618 + trainer_ok, *results = ray.get([trainer, *workers])
619 +
620 + assert trainer_ok, "trainer engine did not complete"
621 + for rank, result in enumerate(results):
622 + assert result["count"] == 0, f"worker {rank} unexpectedly called load_weights"
623 + assert result["name"] == "weight"
624 + assert result["shape"] == [SHAPE[0] // 2, SHAPE[1]]
625 + assert result["direct"]
626 + assert result["exact"], f"worker {rank} received the wrong shard"
627 +
628 +
629 +class _Model(torch.nn.Module):
630 + """Stand-in for a TP=2 shard of a tiny model.
631 +
632 + `column` is sharded on dim 0 and `row` on dim 1 (the two shapes a TP linear
633 + layer produces), `norm` is replicated, and `fused` has no checkpoint-name
634 + counterpart — the shape a fused parameter presents to the resolver.
635 + """
636 +
637 + def __init__(self) -> None:
638 + super().__init__()
639 + self.column = torch.nn.Parameter(torch.zeros(8, 16))
640 + self.column.output_dim = 0
641 + self.column.weight_loader = MethodType(
642 + ColumnParallelLinear.weight_loader, Mock(tp_size=2)
643 + )
644 + self.row = torch.nn.Parameter(torch.zeros(16, 8))
645 + self.row.input_dim = 1
646 + self.row.weight_loader = MethodType(
647 + RowParallelLinear.weight_loader, Mock(tp_size=2)
648 + )
649 + self.norm = torch.nn.Parameter(torch.zeros(16))
650 + self.norm.weight_loader = default_weight_loader
651 + self.fused = torch.nn.Parameter(torch.zeros(24, 16))
652 +
653 +
654 +class _MixedModel(torch.nn.Module):
655 + def __init__(self) -> None:
656 + super().__init__()
657 + self.proj = torch.nn.Module()
658 + self.proj.weight = torch.nn.Parameter(torch.zeros(16, 8))
659 + self.proj.weight.input_dim = 1
660 + self.proj.weight.weight_loader = MethodType(
661 + RowParallelLinear.weight_loader, Mock(tp_size=2)
662 + )
663 + self.proj.bias = torch.nn.Parameter(torch.zeros(16))
664 +
665 +
666 +def _resolve(names, dtypes, shapes, **kwargs):
667 + defaults = dict(num_workers=2, shard_axis_size=2, allow_direct=True)
668 + defaults.update(kwargs)
669 + return resolve_parameter_destinations(_Model(), names, dtypes, shapes, **defaults)
670 +
671 +
672 +class TestDestinationResolution:
673 + @pytest.mark.parametrize(
674 + "name,expected_dim",
675 + [("column", 0), ("row", 1)],
676 + )
677 + def test_sharded_parameter_resolves_to_its_tp_dim(self, name, expected_dim):
678 + [destination] = _resolve([name], [torch.float32], [(16, 16)])
679 + assert destination.direct
680 + assert destination.placements == (REPLICATE, expected_dim)
681 +
682 + def test_replicated_parameter_needs_no_placement(self):
683 + """A parameter every rank holds in full is REPLICATED, so it works
684 + whatever way the inference mesh happens to be factored."""
685 + [destination] = _resolve(["norm"], [torch.float32], [(16,)])
686 + assert destination.direct
687 + assert destination.placements is REPLICATED
688 +
689 + def test_direct_destination_is_the_live_parameter(self):
690 + model = _Model()
691 + [destination] = resolve_parameter_destinations(
692 + model,
693 + ["column"],
694 + [torch.float32],
695 + [(16, 16)],
696 + num_workers=2,
697 + shard_axis_size=2,
698 + allow_direct=True,
699 + )
700 + assert destination.tensor.data_ptr() == model.column.data_ptr()
701 +
702 + def test_known_transforming_model_submodule_falls_back(self):
703 + class GPT2Model(_Model):
704 + pass
705 +
706 + GPT2Model.__module__ = "vllm.model_executor.models.gpt2"
707 + model = torch.nn.Module()
708 + model.transforming = GPT2Model()
709 + [destination] = resolve_parameter_destinations(
710 + model,
711 + ["transforming.row"],
712 + [torch.float32],
713 + [(16, 16)],
714 + num_workers=2,
715 + shard_axis_size=2,
716 + allow_direct=True,
717 + )
718 +
719 + assert not destination.direct
720 + assert destination.placements is REPLICATED
721 +
722 + def test_unknown_name_falls_back(self):
723 + """Fused parameters reach the worker under checkpoint names that do not
724 + exist in the model; those must take the full-tensor path."""
725 + [destination] = _resolve(["mlp.gate_proj.weight"], [torch.float32], [(16, 16)])
726 + assert not destination.direct
727 + assert destination.placements is REPLICATED
728 +
729 + def test_fallback_parameter_demotes_direct_sibling_in_same_module(self):
730 + """Layerwise finalization must not overwrite a directly loaded weight
731 + when a sibling bias takes the fallback path."""
732 + destinations = resolve_parameter_destinations(
733 + _MixedModel(),
734 + ["proj.weight", "proj.bias"],
735 + [torch.float32, torch.float32],
736 + [(16, 16), (16,)],
737 + num_workers=2,
738 + shard_axis_size=2,
739 + allow_direct=True,
740 + )
741 +
742 + assert not any(destination.direct for destination in destinations)
743 +
744 + def test_unknown_loader_falls_back_despite_matching_name_and_shape(self):
745 + model = _Model()
746 + model.norm.weight_loader = lambda param, weight: param.data.copy_(weight)
747 + [destination] = resolve_parameter_destinations(
748 + model,
749 + ["norm"],
750 + [torch.float32],
751 + [(16,)],
752 + num_workers=2,
753 + shard_axis_size=2,
754 + allow_direct=True,
755 + )
756 + assert not destination.direct
757 +
758 + def test_missing_loader_is_not_assumed_to_be_a_plain_copy(self):
759 + model = _Model()
760 + del model.norm.weight_loader
761 + [destination] = resolve_parameter_destinations(
762 + model,
763 + ["norm"],
764 + [torch.float32],
765 + [(16,)],
766 + num_workers=2,
767 + shard_axis_size=2,
768 + allow_direct=True,
769 + )
770 + assert not destination.direct
771 +
772 + def test_wrapped_known_loader_falls_back(self):
773 + model = _Model()
774 +
775 + @wraps(default_weight_loader)
776 + def wrapped_loader(param, weight):
777 + default_weight_loader(param, weight)
778 +
779 + model.norm.weight_loader = wrapped_loader
780 + [destination] = resolve_parameter_destinations(
781 + model,
782 + ["norm"],
783 + [torch.float32],
784 + [(16,)],
785 + num_workers=2,
786 + shard_axis_size=2,
787 + allow_direct=True,
788 + )
789 + assert not destination.direct
790 +
791 + def test_replicated_loader_with_transform_metadata_falls_back(self):
792 + model = _Model()
793 + model.norm.is_transposed = True
794 + [destination] = resolve_parameter_destinations(
795 + model,
796 + ["norm"],
797 + [torch.float32],
798 + [(16,)],
799 + num_workers=2,
800 + shard_axis_size=2,
801 + allow_direct=True,
802 + )
803 + assert not destination.direct
804 +
805 + def test_packed_dim_zero_falls_back(self):
806 + model = _Model()
807 + model.column.packed_dim = 0
808 + [destination] = resolve_parameter_destinations(
809 + model,
810 + ["column"],
811 + [torch.float32],
812 + [(16, 16)],
813 + num_workers=2,
814 + shard_axis_size=2,
815 + allow_direct=True,
816 + )
817 + assert not destination.direct
818 +
819 + def test_loader_declared_dim_wins_over_shape_guessing(self):
820 + model = _Model()
821 + model.column.output_dim = 1
822 + [destination] = resolve_parameter_destinations(
823 + model,
824 + ["column"],
825 + [torch.float32],
826 + [(16, 16)],
827 + num_workers=2,
828 + shard_axis_size=2,
829 + allow_direct=True,
830 + )
831 + assert not destination.direct
832 +
833 + def test_shape_the_tp_factor_cannot_explain_falls_back(self):
834 + [destination] = _resolve(["fused"], [torch.float32], [(16, 16)])
835 + assert not destination.direct
836 +
837 + def test_dtype_mismatch_falls_back(self):
838 + """A parameter stored in a different dtype than the wire dtype means a
839 + quantized or otherwise transformed layout — never write into it."""
840 + [destination] = _resolve(["column"], [torch.bfloat16], [(16, 16)])
841 + assert not destination.direct
842 +
843 + def test_allow_direct_false_forces_every_parameter_to_fall_back(self):
844 + destinations = _resolve(
845 + ["column", "norm"],
846 + [torch.float32] * 2,
847 + [(16, 16), (16,)],
848 + allow_direct=False,
849 + )
850 + assert not any(d.direct for d in destinations)
851 +
852 + def test_plan_reports_byte_coverage(self, caplog_vllm):
853 + """A small direct parameter must not hide a large fallback tensor."""
854 + with caplog_vllm.at_level(
855 + logging.INFO,
856 + logger="vllm.distributed.weight_transfer.m2n_layout",
857 + ):
858 + _resolve(
859 + ["norm", "missing_fused_weight"],
860 + [torch.float32, torch.bfloat16],
861 + [(16,), (1024, 1024)],
862 + )
863 +
864 + assert (
865 + "1/2 parameters resharded directly into the model, 1 via "
866 + "full-tensor fallback; direct byte coverage: 64/2097216 bytes "
867 + "(0.0%)" in caplog_vllm.text
868 + )
vllm/config/weight_transfer.py
+3 -1
@@ -9,7 +9,9 @@ from vllm.config.utils import config
9 class WeightTransferConfig:
10 """Configuration for weight transfer during RL training."""
11
12 - backend: Literal["nccl", "ipc", "sparse_nccl", "sharded_rdt"] | str = "nccl"
12 + backend: Literal["nccl", "ipc", "sparse_nccl", "sharded_rdt", "nccl_m2n"] | str = (
13 + "nccl"
14 + )
15 """The backend to use for weight transfer. Validated against the
16 `WeightTransferEngineFactory` registry at engine creation time.
17 """
vllm/distributed/weight_transfer/factory.py
+12
@@ -243,6 +243,12 @@ WeightTransferEngineFactory.register_engine(
243 "ShardedRDTWeightTransferEngine",
244 )
245
246 +WeightTransferEngineFactory.register_engine(
247 + "nccl_m2n",
248 + "vllm.distributed.weight_transfer.m2n_engine",
249 + "M2NWeightTransferEngine",
250 +)
251 +
252
253 # Trainer-side engines, parallel to the worker registry above.
254 WeightTransferTrainerFactory.register_engine(
@@ -268,3 +274,9 @@ WeightTransferTrainerFactory.register_engine(
274 "vllm.distributed.weight_transfer.sharded_rdt_trainer",
275 "ShardedRDTTrainerWeightTransferEngine",
276 )
277 +
278 +WeightTransferTrainerFactory.register_engine(
279 + "nccl_m2n",
280 + "vllm.distributed.weight_transfer.m2n_trainer",
281 + "M2NTrainerWeightTransferEngine",
282 +)
vllm/distributed/weight_transfer/m2n_common.py new
+338
@@ -0,0 +1,338 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared helpers for the NCCL M2N (`nccl_m2n`) weight transfer backend.
4 +
5 +M2N reshards a tensor between two disjoint meshes of ranks that live in one
6 +communicator: the trainer occupies ranks `[0, T)` and the inference workers
7 +`[T, T + N)`. This module holds everything both sides need — the optional
8 +runtime import, layout descriptors, and the conversion to `nccl.m2n` types —
9 +so the engine module stays about the transfer itself.
10 +"""
11 +
12 +from collections.abc import Sequence
13 +from dataclasses import dataclass
14 +from typing import TYPE_CHECKING, Any, cast
15 +
16 +import torch
17 +
18 +from vllm.distributed.weight_transfer.base import ParamMeta
19 +
20 +if TYPE_CHECKING:
21 + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
22 +
23 +# Placement code for a mesh axis that replicates. Any other (non-negative) code
24 +# is the tensor dim that axis shards. Ints rather than `nccl.m2n` objects so
25 +# layouts survive the JSON init handshake without a custom encoder.
26 +REPLICATE = -1
27 +
28 +# `NCCL_RESHARD_MAX_TENSOR_DIMS`; 4-D and higher are rejected by the library.
29 +MAX_TENSOR_DIMS = 3
30 +
31 +MESH_NDIMS = 2
32 +
33 +# Destination mesh convention. Keeping these roles named avoids scattering
34 +# positional axis assumptions through the worker-side planner.
35 +DESTINATION_REPLICA_AXIS = 0
36 +DESTINATION_SHARD_AXIS = 1
37 +
38 +# m2n bounds how many source shards may feed one destination shard, and how many
39 +# destination shards one source shard may feed, with compile-time arrays in
40 +# `reshard_limits.h`. The bindings do not expose them, so they are mirrored
41 +# here; a build with larger arrays makes these conservative, never wrong.
42 +MAX_SOURCE_SHARDS = 16
43 +MAX_DEST_SHARDS = 64
44 +
45 +# Wire dtypes `ncclReshard` accepts. Notably excludes fp4 and any packed
46 +# sub-byte type, so quantized checkpoints are out of scope for now.
47 +SUPPORTED_DTYPES: frozenset[torch.dtype] = frozenset(
48 + dtype
49 + for dtype in (
50 + getattr(torch, name, None)
51 + for name in (
52 + "int8",
53 + "uint8",
54 + "float8_e4m3fn",
55 + "float8_e5m2",
56 + "float16",
57 + "bfloat16",
58 + "int32",
59 + "uint32",
60 + "float32",
61 + "int64",
62 + "uint64",
63 + "float64",
64 + )
65 + )
66 + if dtype is not None
67 +)
68 +
69 +_IMPORT_HINT = (
70 + "The nccl_m2n weight transfer backend requires the `nccl-extensions` "
71 + "package (and its `nccl4py` dependency), which is not installed by vLLM. "
72 + "See https://github.com/NVIDIA/nccl-extensions for build "
73 + "and install instructions. It needs NCCL 2.30.5 or newer, and "
74 + "VLLM_NCCL_SO_PATH must point at the same libnccl.so that libnccl_m2n.so "
75 + "was linked against."
76 +)
77 +
78 +
79 +def import_m2n() -> Any:
80 + """Import `nccl.m2n` lazily, with an actionable error when it is missing.
81 +
82 + Deferred so that importing vLLM — or any other weight transfer backend —
83 + never requires the m2n runtime to be present.
84 + """
85 + try:
86 + import nccl.m2n as m2n
87 + except ImportError as e:
88 + raise ImportError(f"{_IMPORT_HINT} (import failed: {e})") from e
89 + return m2n
90 +
91 +
92 +@dataclass(frozen=True)
93 +class M2NMesh:
94 + """One side's rank topology — pure topology, no tensor placement.
95 +
96 + Mirrors `ncclMesh_t`: a 2-axis mesh owning the contiguous rank interval
97 + `[start_rank, start_rank + dims[0] * dims[1])`. There is no 1-D mesh; a
98 + single-axis topology is spelled with a second axis of size 1.
99 +
100 + One mesh describes every tensor on its side, so it is exchanged once at the
101 + init handshake rather than per parameter.
102 + """
103 +
104 + dims: tuple[int, int]
105 + start_rank: int
106 +
107 + def __post_init__(self) -> None:
108 + if len(self.dims) != MESH_NDIMS or any(d <= 0 for d in self.dims):
109 + raise ValueError(f"mesh dims must be {MESH_NDIMS} positive ints")
110 + if self.start_rank < 0:
111 + raise ValueError(
112 + f"mesh start_rank must be non-negative, got {self.start_rank}"
113 + )
114 +
115 + @property
116 + def size(self) -> int:
117 + return self.dims[0] * self.dims[1]
118 +
119 +
120 +# A tensor's placement over its side's mesh: one code per mesh axis
121 +# (`REPLICATE`, or the tensor dim that axis shards), or `REPLICATED` for a
122 +# tensor every rank holds in full.
123 +Placements = tuple[int, int]
124 +REPLICATED: Placements | None = None
125 +
126 +
127 +def check_placements(placements: Placements, context: str = "placements") -> None:
128 + """Reject placement pairs m2n cannot express.
129 +
130 + The header requires exactly one SHARD axis and one REPLICATE axis.
131 + `{REPLICATE, REPLICATE}` hits a degenerate prepare branch, which is why
132 + full replication is carried as `REPLICATED` and resolved separately;
133 + sharding both axes is not expressible at all.
134 + """
135 + if len(placements) != MESH_NDIMS:
136 + raise ValueError(f"{context} must have {MESH_NDIMS} entries")
137 + invalid = [code for code in placements if code < REPLICATE]
138 + if invalid:
139 + raise ValueError(
140 + f"{context} contains invalid placement code {invalid[0]}; valid "
141 + f"codes are {REPLICATE} (Replicate) and non-negative tensor dimensions"
142 + )
143 + num_sharded = sum(code != REPLICATE for code in placements)
144 + if num_sharded == 0:
145 + raise ValueError(
146 + f"{context}: a fully replicated tensor is carried as REPLICATED, "
147 + "not as two "
148 + "REPLICATE axes"
149 + )
150 + if num_sharded == MESH_NDIMS:
151 + raise ValueError(
152 + f"{context}: nccl_m2n needs one REPLICATE mesh axis, but "
153 + f"{placements} shards both"
154 + )
155 +
156 +
157 +def resolve_layout(
158 + mesh: M2NMesh,
159 + placements: Placements | None,
160 + context: str = "placements",
161 +) -> tuple[M2NMesh, Placements]:
162 + """Pair one tensor's placement with the mesh m2n should see for it.
163 +
164 + A replicated tensor needs a size-1 mesh axis to carry a no-op shard, since
165 + m2n has no `{REPLICATE, REPLICATE}`. It is therefore described over the
166 + *same rank interval* re-factored as `(size, 1)`. That re-factoring is sound
167 + precisely because replication is order-independent — every rank holds the
168 + whole tensor, so it does not matter that `(a, b)` and `(size, 1)` walk the
169 + interval in a different order. A sharded tensor keeps its side's own
170 + factorization, where rank order decides who owns which shard.
171 + """
172 + if placements is None:
173 + return M2NMesh((mesh.size, 1), mesh.start_rank), (REPLICATE, 0)
174 + check_placements(placements, context)
175 + return mesh, placements
176 +
177 +
178 +def shard_count(mesh: M2NMesh, placements: Placements) -> int:
179 + """How many pieces this layout splits the tensor into."""
180 + return next(
181 + (mesh.dims[axis] for axis, code in enumerate(placements) if code != REPLICATE),
182 + 1,
183 + )
184 +
185 +
186 +def check_plan_limits(
187 + src: tuple[M2NMesh, Placements],
188 + dst: tuple[M2NMesh, Placements],
189 + name: str,
190 +) -> None:
191 + """Reject plans that exceed m2n's static per-shard fan-in / fan-out arrays.
192 +
193 + Only the unambiguous cases are checked: when one side is a single shard it
194 + is fed by (or feeds) every shard on the other side, so the count is exactly
195 + the other side's shard count. In the general sharded-to-sharded case the
196 + overlap depends on the library's chunking, and reproducing that arithmetic
197 + here would duplicate internals that can change under us.
198 +
199 + m2n re-checks authoritatively and, because the plan is derived from the
200 + shared descriptors, fails identically on every rank -- so this is about
201 + reporting at init with a message that names the parameter, not about
202 + avoiding a hang.
203 + """
204 + src_shards = shard_count(*src)
205 + dst_shards = shard_count(*dst)
206 + if dst_shards == 1 and src_shards > MAX_SOURCE_SHARDS:
207 + raise ValueError(
208 + f"parameter '{name}' would feed {src_shards} source shards into one "
209 + f"replicated destination, over m2n's MAX_SOURCES={MAX_SOURCE_SHARDS}. "
210 + "Shard the destination, replicate the source, or rebuild m2n with "
211 + "larger arrays."
212 + )
213 + if src_shards == 1 and dst_shards > MAX_DEST_SHARDS:
214 + raise ValueError(
215 + f"parameter '{name}' would fan one source shard out to {dst_shards} "
216 + f"destination shards, over m2n's MAX_TARGETS={MAX_DEST_SHARDS}. "
217 + "Rebuild m2n with larger arrays to raise the bound."
218 + )
219 +
220 +
221 +def validate_layout(
222 + mesh: M2NMesh, placements: Placements, shape: Sequence[int], side: str
223 +) -> None:
224 + """Check a resolved layout can describe `shape` before any collective runs.
225 +
226 + Called on both sides at init so a bad plan surfaces as an error from the
227 + init RPC rather than as a hang inside the first reshard.
228 + """
229 + for axis, code in enumerate(placements):
230 + if code == REPLICATE:
231 + continue
232 + if code >= len(shape):
233 + raise ValueError(
234 + f"{side} layout shards tensor dim {code}, but the tensor "
235 + f"has rank {len(shape)}"
236 + )
237 + factor = mesh.dims[axis]
238 + if shape[code] % factor:
239 + raise ValueError(
240 + f"{side} layout shards dim {code} (size {shape[code]}) over "
241 + f"{factor} ranks, which does not divide evenly"
242 + )
243 +
244 +
245 +@dataclass(frozen=True)
246 +class M2NParamMeta(ParamMeta):
247 + """`ParamMeta` extended with how the trainer places this tensor.
248 +
249 + The base class carries only name / dtype / full shape, which is not enough
250 + to plan a reshard. `placements` is relative to its side's `M2NMesh`, or
251 + `REPLICATED` when every rank holds the whole tensor.
252 + """
253 +
254 + placements: Placements | None
255 +
256 +
257 +def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> None:
258 + """Reject tensors m2n cannot move at all, with the parameter named."""
259 + if dtype not in SUPPORTED_DTYPES:
260 + raise ValueError(
261 + f"parameter '{name}' has dtype {dtype}, which nccl_m2n does not "
262 + f"support. Supported: {sorted(str(d) for d in SUPPORTED_DTYPES)}"
263 + )
264 + if not 1 <= len(shape) <= MAX_TENSOR_DIMS:
265 + raise ValueError(
266 + f"parameter '{name}' has rank {len(shape)}; nccl_m2n supports "
267 + f"rank 1..{MAX_TENSOR_DIMS}"
268 + )
269 +
270 +
271 +def to_mesh(m2n: Any, mesh: M2NMesh) -> Any:
272 + return m2n.Mesh(mesh.dims, start_rank=mesh.start_rank)
273 +
274 +
275 +def to_placements(m2n: Any, placements: Placements) -> list[Any]:
276 + return [
277 + m2n.Replicate() if code == REPLICATE else m2n.Shard(code) for code in placements
278 + ]
279 +
280 +
281 +def publish_destination_placements(
282 + comm: "PyNcclCommunicator",
283 + first_worker_rank: int,
284 + placements: "Sequence[Placements | None] | None",
285 + num_parameters: int,
286 +) -> list[Placements | None]:
287 + """Share the worker-side destination plan with every rank in the group.
288 +
289 + The trainer must issue each reshard with the same destination the workers
290 + use, but once destinations are per-parameter they depend on the inference
291 + model, which only the workers can see. The first worker publishes them here,
292 + over the shared NCCL communicator; trainer ranks pass `None` and receive
293 + them, and the other workers pass their own so a disagreement is caught
294 + rather than deadlocking later.
295 +
296 + Only placements travel: the destination *mesh* is still derived from the
297 + rank split, identically on both sides. Use the NCCL communicator directly
298 + because unique-id rendezvous deliberately has no bootstrap process group.
299 + """
300 + replicated_sentinel = REPLICATE - 1
301 + encoded = torch.full(
302 + (num_parameters, MESH_NDIMS),
303 + replicated_sentinel,
304 + dtype=torch.int8,
305 + device=comm.device,
306 + )
307 + if comm.rank == first_worker_rank:
308 + if placements is None or len(placements) != num_parameters:
309 + raise ValueError(
310 + f"publishing rank needs {num_parameters} destination placements"
311 + )
312 + rows = [
313 + [replicated_sentinel] * MESH_NDIMS if placement is REPLICATED else placement
314 + for placement in placements
315 + ]
316 + if rows:
317 + encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=comm.device))
318 +
319 + comm.broadcast(encoded, src=first_worker_rank)
320 + return [
321 + REPLICATED
322 + if all(code == replicated_sentinel for code in row)
323 + else cast(Placements, tuple(row))
324 + for row in encoded.tolist()
325 + ]
326 +
327 +
328 +def comm_ptr(comm: "PyNcclCommunicator") -> int:
329 + """Raw `ncclComm_t` behind vLLM's `PyNcclCommunicator`.
330 +
331 + m2n links its own NCCL, so the handle only means anything if vLLM loaded
332 + the same `libnccl.so` — set `VLLM_NCCL_SO_PATH` accordingly.
333 + """
334 + handle = comm.comm
335 + ptr = getattr(handle, "value", handle)
336 + if not ptr:
337 + raise RuntimeError("PyNcclCommunicator has no live NCCL communicator")
338 + return int(ptr)
vllm/distributed/weight_transfer/m2n_engine.py new
+453
@@ -0,0 +1,453 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-side weight transfer engine built on NCCL M2N (`nccl_m2n`).
4 +
5 +The trainer and the inference workers share one NCCL communicator: trainer ranks
6 +occupy `[0, T)`, workers `[T, T + N)`. Each parameter is moved with a single
7 +`nccl.m2n.reshard`, which redistributes it from the trainer's layout (FSDP / EP
8 +/ arbitrary DTensor sharding) to the inference layout — so the trainer sends its
9 +local shards and never all-gathers a full tensor, which is what the broadcast
10 +NCCL backend forces it to do.
11 +
12 +When an incoming checkpoint parameter maps directly to a live vLLM parameter,
13 +M2N reshards into that worker's local model storage. Parameters whose names or
14 +layouts cannot be resolved receive a full tensor and fall back to
15 +`load_weights`, preserving the existing backend's loading behavior.
16 +
17 +Both meshes participate in every reshard, in the same order, so this backend has
18 +the same concurrency shape as the broadcast NCCL backend: the worker must be
19 +inside `receive_weights` while the trainer is sending. Driving that from the
20 +trainer is the trainer engine's job.
21 +"""
22 +
23 +from collections.abc import Iterator
24 +from dataclasses import dataclass, field
25 +from typing import TYPE_CHECKING, Any, cast
26 +
27 +import torch
28 +
29 +from vllm.config.weight_transfer import WeightTransferConfig
30 +from vllm.distributed.weight_transfer.base import (
31 + WeightTransferEngine,
32 + WeightTransferInitInfo,
33 + WeightTransferUpdateInfo,
34 +)
35 +from vllm.distributed.weight_transfer.m2n_common import (
36 + DESTINATION_SHARD_AXIS,
37 + MESH_NDIMS,
38 + REPLICATED,
39 + M2NMesh,
40 + M2NParamMeta,
41 + Placements,
42 + check_plan_limits,
43 + check_transferable,
44 + comm_ptr,
45 + import_m2n,
46 + publish_destination_placements,
47 + resolve_layout,
48 + to_mesh,
49 + to_placements,
50 + validate_layout,
51 +)
52 +from vllm.distributed.weight_transfer.m2n_layout import (
53 + M2NDestination,
54 + resolve_parameter_destinations,
55 +)
56 +from vllm.distributed.weight_transfer.nccl_common import (
57 + NCCLWeightTransferInitInfo,
58 + decode_nccl_unique_id,
59 + worker_init_process_group,
60 +)
61 +
62 +if TYPE_CHECKING:
63 + from vllm.config import VllmConfig
64 + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
65 +
66 +__all__ = [
67 + "M2NWeightTransferInitInfo",
68 + "M2NWeightTransferUpdateInfo",
69 + "M2NWeightTransferEngine",
70 +]
71 +
72 +
73 +# ---------------------------------------------------------------------------
74 +# Wire types
75 +# ---------------------------------------------------------------------------
76 +
77 +
78 +@dataclass(kw_only=True)
79 +class M2NWeightTransferInitInfo(WeightTransferInitInfo):
80 + """Worker-side init info: the rendezvous plus the full transfer plan.
81 +
82 + Layouts are static for the whole run, so they ride the one-time init
83 + handshake and the per-round update info stays a list of names. Everything is
84 + plain JSON so the HTTP control plane carries it unchanged.
85 + """
86 +
87 + master_address: str | None = None
88 + master_port: int | None = None
89 + rank_offset: int
90 + """First worker rank, i.e. the number of trainer ranks."""
91 + world_size: int
92 + """Trainer ranks + all inference workers."""
93 + src_mesh_dims: list[int]
94 + """The trainer's mesh, shared by every parameter (`start_rank` is 0: the
95 + trainer occupies the front of the communicator)."""
96 + dst_mesh_dims: list[int]
97 + """The inference mesh, starting at `rank_offset`. Declared rather than
98 + derived so both sides describe the destination identically."""
99 + names: list[str]
100 + dtype_names: list[str]
101 + shapes: list[list[int]]
102 + src_placements: list[list[int] | None]
103 + """Per parameter, relative to `src_mesh_dims`; `None` means replicated."""
104 + nccl_unique_id_b64: str | None = field(default=None, repr=False)
105 + max_cta: int | None = None
106 +
107 + def __post_init__(self) -> None:
108 + _ = self.nccl_unique_id_bytes
109 + num_params = len(self.names)
110 + for label, values in (
111 + ("dtype_names", self.dtype_names),
112 + ("shapes", self.shapes),
113 + ("src_placements", self.src_placements),
114 + ):
115 + if len(values) != num_params:
116 + raise ValueError(
117 + f"`{label}` should be of the same size as `names`: "
118 + f"got {len(values)} and {num_params}"
119 + )
120 + if self.rank_offset < 1 or self.rank_offset >= self.world_size:
121 + raise ValueError(
122 + f"`rank_offset` ({self.rank_offset}) must leave at least one "
123 + f"trainer rank and one worker in world_size {self.world_size}"
124 + )
125 + num_workers = self.world_size - self.rank_offset
126 + dst = tuple(self.dst_mesh_dims)
127 + if len(dst) != MESH_NDIMS or dst[0] * dst[1] != num_workers:
128 + raise ValueError(
129 + f"`dst_mesh_dims` {dst} must be {MESH_NDIMS} dims covering the "
130 + f"{num_workers} inference workers"
131 + )
132 +
133 + @property
134 + def nccl_unique_id_bytes(self) -> bytes | None:
135 + return decode_nccl_unique_id(
136 + master_address=self.master_address,
137 + master_port=self.master_port,
138 + nccl_unique_id_b64=self.nccl_unique_id_b64,
139 + ctx="M2NWeightTransferInitInfo",
140 + )
141 +
142 +
143 +@dataclass
144 +class M2NWeightTransferUpdateInfo(WeightTransferUpdateInfo):
145 + """Per-round update info: which parameters this chunk carries, in order.
146 +
147 + Shapes, dtypes and layouts were fixed at init, so a round only needs to say
148 + what is coming and in what order — both sides must issue their reshards in
149 + exactly that order.
150 + """
151 +
152 + names: list[str]
153 +
154 +
155 +# ---------------------------------------------------------------------------
156 +# Worker engine
157 +# ---------------------------------------------------------------------------
158 +
159 +
160 +class M2NWeightTransferEngine(
161 + WeightTransferEngine[M2NWeightTransferInitInfo, M2NWeightTransferUpdateInfo]
162 +):
163 + """Inference-side engine: receives each parameter with one reshard.
164 +
165 + Resolvable parameters are resharded from the trainer layout directly into
166 + each worker's live local model storage. Unresolvable parameters receive a
167 + full tensor and use `load_weights`. In both cases, the trainer sends its
168 + local shards without materializing a full tensor.
169 + """
170 +
171 + init_info_cls = M2NWeightTransferInitInfo
172 + update_info_cls = M2NWeightTransferUpdateInfo
173 + # The destination plan is resolved against the model this engine was built
174 + # for, so it cannot be retargeted at a draft model mid-flight.
175 + supports_draft_weight_update = False
176 +
177 + def __init__(
178 + self,
179 + config: WeightTransferConfig,
180 + vllm_config: "VllmConfig",
181 + device: torch.device,
182 + model: torch.nn.Module,
183 + ) -> None:
184 + super().__init__(config, vllm_config, device, model)
185 + self.model_update_group: PyNcclCommunicator | None = None
186 + self._m2n: Any = None
187 + self._handle: Any = None
188 + self._metas: list[M2NParamMeta] = []
189 + self._src_mesh: M2NMesh | None = None
190 + self._dst_mesh: M2NMesh | None = None
191 + self._parameter_destinations: list[M2NDestination] = []
192 + self._index: dict[str, int] = {}
193 + self._uses_load_weights = False
194 +
195 + def init_transfer_engine(self, init_info: M2NWeightTransferInitInfo) -> None:
196 + """Join the trainer's communicator and rebuild the transfer plan.
197 +
198 + Every precondition (dtype, tensor rank, divisibility) is checked here so
199 + a bad plan fails the init RPC instead of hanging the first collective.
200 + """
201 + self._m2n = import_m2n()
202 + self._metas = []
203 +
204 + self._src_mesh = M2NMesh(
205 + cast(tuple[int, int], tuple(init_info.src_mesh_dims)), 0
206 + )
207 + if self._src_mesh.size != init_info.rank_offset:
208 + raise ValueError(
209 + f"source mesh covers {self._src_mesh.size} ranks, but there are "
210 + f"{init_info.rank_offset} trainer ranks"
211 + )
212 + # The inference topology, as declared by the trainer. This is the mesh
213 + # a sharded destination is placed over; a replicated one is described
214 + # by the alternate descriptor `resolve_layout` derives from it, which
215 + # covers the same ranks but is not this mesh.
216 + self._dst_mesh = M2NMesh(
217 + cast(tuple[int, int], tuple(init_info.dst_mesh_dims)),
218 + init_info.rank_offset,
219 + )
220 +
221 + for name, name_dtype, shape, placements in zip(
222 + init_info.names,
223 + init_info.dtype_names,
224 + init_info.shapes,
225 + init_info.src_placements,
226 + ):
227 + dtype = getattr(torch, name_dtype, None)
228 + if not isinstance(dtype, torch.dtype):
229 + raise ValueError(
230 + f"parameter '{name}' has invalid dtype name "
231 + f"'{name_dtype}'; expected the name of a torch.dtype"
232 + )
233 + check_transferable(name, dtype, shape)
234 + codes = (
235 + REPLICATED
236 + if placements is None
237 + else cast(Placements, tuple(placements))
238 + )
239 + src_mesh, src_placements = resolve_layout(
240 + self._src_mesh, codes, f"parameter '{name}' source placements"
241 + )
242 + validate_layout(src_mesh, src_placements, shape, "source")
243 + # Destination checks are deferred until loader-aware resolution;
244 + # assuming replication here can reject a valid sharded plan.
245 + self._metas.append(M2NParamMeta(name, dtype, tuple(shape), codes))
246 +
247 + parallel_config = self.parallel_config
248 + shard_axis_size = self._dst_mesh.dims[DESTINATION_SHARD_AXIS]
249 + quantization = getattr(self.model_config, "quantization", None)
250 + # A quantized weight loader may dequantize / transpose / requantize, and
251 + # under PP a rank's local shape is not the checkpoint shape split over
252 + # the shard axis. Both break shape-derived resolution, so take the
253 + # fallback.
254 + allow_direct = (
255 + quantization is None and parallel_config.pipeline_parallel_size == 1
256 + )
257 + # Match every incoming checkpoint parameter to its rank-local target.
258 + # Resolvable parameters can be resharded directly into live model
259 + # storage; the rest receive a full tensor for `load_weights`.
260 + self._parameter_destinations = resolve_parameter_destinations(
261 + self.model,
262 + [m.name for m in self._metas],
263 + [m.dtype for m in self._metas],
264 + [m.shape for m in self._metas],
265 + num_workers=self._dst_mesh.size,
266 + shard_axis_size=shard_axis_size,
267 + allow_direct=allow_direct,
268 + )
269 + # Check the resolved layout rather than a provisional replicated one.
270 + # For example, 32 source shards feeding 4 destination shards may need
271 + # only 8 sources per destination, while a replicated plan needs all 32.
272 + for meta, destination in zip(self._metas, self._parameter_destinations):
273 + src_mesh, src_placements = resolve_layout(
274 + self._src_mesh,
275 + meta.placements,
276 + f"parameter '{meta.name}' source placements",
277 + )
278 + dst_mesh, dst_placements = resolve_layout(
279 + self._dst_mesh,
280 + destination.placements,
281 + f"parameter '{meta.name}' destination placements",
282 + )
283 + validate_layout(dst_mesh, dst_placements, meta.shape, "destination")
284 + check_plan_limits(
285 + (src_mesh, src_placements),
286 + (dst_mesh, dst_placements),
287 + meta.name,
288 + )
289 +
290 + # Update requests carry names only, so cache their plan indices. Track
291 + # whether any fallback entry requires the `load_weights` lifecycle.
292 + self._index = {meta.name: i for i, meta in enumerate(self._metas)}
293 + self._uses_load_weights = any(
294 + not destination.direct for destination in self._parameter_destinations
295 + )
296 +
297 + self.model_update_group = worker_init_process_group(
298 + NCCLWeightTransferInitInfo(
299 + master_address=init_info.master_address,
300 + master_port=init_info.master_port,
301 + nccl_unique_id_b64=init_info.nccl_unique_id_b64,
302 + rank_offset=init_info.rank_offset,
303 + world_size=init_info.world_size,
304 + ),
305 + parallel_config,
306 + )
307 + # Destinations are per-parameter and depend on the inference model, so
308 + # the trainer cannot derive them. The first worker publishes the plan
309 + # and every other rank checks its own against it — a silent
310 + # disagreement between workers would mean mismatched collectives.
311 + mine = [destination.placements for destination in self._parameter_destinations]
312 + agreed = publish_destination_placements(
313 + self.model_update_group, init_info.rank_offset, mine, len(self._metas)
314 + )
315 + if agreed != mine:
316 + mismatched = next(
317 + meta.name for meta, a, b in zip(self._metas, mine, agreed) if a != b
318 + )
319 + raise RuntimeError(
320 + "inference workers resolved different nccl_m2n destinations "
321 + f"(first mismatch: '{mismatched}'); all workers must run the "
322 + "same model and parallel config"
323 + )
324 +
325 + self._handle = self._m2n.Handle.create(
326 + self._m2n.Config(max_cta=init_info.max_cta)
327 + )
328 +
329 + def start_weight_update(self) -> None:
330 + """Set up layerwise reloading, but only if some parameter needs it.
331 +
332 + Directly-resharded parameters are written in place and never go through
333 + `load_weights`, so a plan with no fallback entries has nothing to
334 + reload.
335 + """
336 + if not self._uses_load_weights:
337 + return
338 + from vllm.model_executor.model_loader.reload import initialize_layerwise_reload
339 +
340 + initialize_layerwise_reload(self.model)
341 +
342 + def finish_weight_update(self) -> None:
343 + """Finalize layerwise reloading when the plan uses fallback entries."""
344 + if not self._uses_load_weights:
345 + return
346 + from vllm.model_executor.model_loader.reload import finalize_layerwise_reload
347 +
348 + finalize_layerwise_reload(self.model, self.model_config)
349 +
350 + def receive_weights(self, update_info: M2NWeightTransferUpdateInfo) -> None:
351 + """Receive each requested parameter using its initialization-time plan."""
352 + if self._handle is None or self.model_update_group is None:
353 + raise RuntimeError(
354 + "nccl_m2n weight transfer not initialized. "
355 + "Call init_transfer_engine() first."
356 + )
357 +
358 + requested = []
359 + for name in update_info.names:
360 + index = self._index.get(name)
361 + if index is None:
362 + raise ValueError(
363 + f"parameter '{name}' was not declared at init; the "
364 + "trainer must send the same parameter set it announced"
365 + )
366 + requested.append((name, index))
367 +
368 + from vllm.model_executor.model_loader.mtp_validation import (
369 + disable_mtp_completeness_check,
370 + )
371 +
372 + comm = comm_ptr(self.model_update_group)
373 + stream = torch.cuda.current_stream()
374 +
375 + # Reshard lazily so one `load_weights` invocation sees the complete
376 + # fallback sequence without staging every full tensor at once. Direct
377 + # destinations still execute in request order as the iterator advances.
378 + def reshard_requested_weights() -> Iterator[tuple[str, torch.Tensor]]:
379 + for name, index in requested:
380 + meta = self._metas[index]
381 + destination = self._parameter_destinations[index]
382 + buffer = destination.tensor
383 + if buffer is None:
384 + buffer = torch.empty(
385 + meta.shape, dtype=meta.dtype, device=self.device
386 + )
387 +
388 + self._reshard(comm, stream, meta, destination.placements, buffer)
389 +
390 + if destination.tensor is None:
391 + # `load_weights` reads on the host stream, so the transfer
392 + # has to have landed before it runs.
393 + stream.synchronize()
394 + yield name, buffer
395 +
396 + has_fallback = any(
397 + self._parameter_destinations[index].tensor is None for _, index in requested
398 + )
399 + with disable_mtp_completeness_check():
400 + received_weights = reshard_requested_weights()
401 + if has_fallback:
402 + # Some model loaders keep invocation-local state to combine
403 + # related checkpoint tensors, so preserve one iterable scope.
404 + self.model.load_weights(received_weights)
405 + else:
406 + # The iterator performs the reshards. With nothing to yield to
407 + # `load_weights`, exhaust it here to execute direct transfers.
408 + for _ in received_weights:
409 + pass
410 +
411 + def _reshard(
412 + self,
413 + comm: int,
414 + stream: torch.cuda.Stream,
415 + src_meta: M2NParamMeta,
416 + dst_placements: Placements | None,
417 + dst_buffer: torch.Tensor,
418 + ) -> None:
419 + """Reshard one parameter from its trainer layout into a worker buffer.
420 +
421 + `dst_placements` describes the buffer's logical placement over the
422 + worker mesh; `dst_buffer` is this rank's physical storage.
423 + """
424 + assert self._src_mesh is not None
425 + assert self._dst_mesh is not None
426 + m2n = self._m2n
427 + src_mesh, resolved_src_placements = resolve_layout(
428 + self._src_mesh, src_meta.placements
429 + )
430 + dst_mesh, resolved_dst_placements = resolve_layout(
431 + self._dst_mesh, dst_placements
432 + )
433 + m2n.reshard(
434 + None,
435 + dst_buffer,
436 + comm,
437 + stream,
438 + src_mesh=to_mesh(m2n, src_mesh),
439 + src_placements=to_placements(m2n, resolved_src_placements),
440 + dst_mesh=to_mesh(m2n, dst_mesh),
441 + dst_placements=to_placements(m2n, resolved_dst_placements),
442 + handle=self._handle,
443 + )
444 +
445 + def shutdown(self) -> None:
446 + """Finish pending GPU work and release M2N communication state."""
447 + if self._handle is not None:
448 + # M2N does not synchronize caller streams on finalize.
449 + torch.accelerator.synchronize()
450 + self._handle.destroy()
451 + self._handle = None
452 + self.model_update_group = None
453 + self._parameter_destinations = []
vllm/distributed/weight_transfer/m2n_layout.py new
+261
@@ -0,0 +1,261 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Destination-layout resolution for the NCCL M2N weight transfer backend.
4 +
5 +For every incoming checkpoint parameter the worker needs a destination buffer
6 +plus its placement over the inference mesh. Two outcomes are possible:
7 +
8 +* **direct** — when the model is not quantized, the parameter maps 1:1 onto a
9 + live vLLM parameter and its loader is a small, known-safe copy or tensor-
10 + parallel loader. The loader's declared input/output dimension determines the
11 + placement; shapes are validation, not an inference mechanism. The reshard
12 + writes straight into the live parameter, so each rank receives only its own
13 + shard and nothing is copied afterwards.
14 +* **fallback** — anything else. The reshard delivers the whole tensor to every
15 + rank and `load_weights` does the sharding, exactly as the broadcast NCCL
16 + backend does. Fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`)
17 + take this path: the checkpoint name does not name a vLLM parameter, so there
18 + is nothing to resolve against.
19 +
20 +Correctness never depends on a parameter resolving — the fallback is always
21 +available and is the same path the existing backend uses.
22 +"""
23 +
24 +import inspect
25 +import math
26 +from collections.abc import Sequence
27 +from dataclasses import dataclass
28 +from functools import partial
29 +from typing import cast
30 +
31 +import torch
32 +
33 +from vllm.distributed.weight_transfer.m2n_common import (
34 + DESTINATION_REPLICA_AXIS,
35 + DESTINATION_SHARD_AXIS,
36 + MESH_NDIMS,
37 + REPLICATE,
38 + REPLICATED,
39 + Placements,
40 +)
41 +from vllm.logger import init_logger
42 +
43 +logger = init_logger(__name__)
44 +
45 +_DIRECT_WEIGHT_LOADING_BLOCKLIST = {
46 + "vllm.model_executor.models.gpt2.GPT2Model",
47 +}
48 +
49 +
50 +def _module_type_name(module: torch.nn.Module) -> str:
51 + cls = type(module)
52 + return f"{cls.__module__}.{cls.__name__}"
53 +
54 +
55 +def _destination_placements(shard_dim: int) -> Placements:
56 + """Place one tensor shard on the configured destination mesh axis."""
57 + placements = [REPLICATE] * MESH_NDIMS
58 + placements[DESTINATION_REPLICA_AXIS] = REPLICATE
59 + placements[DESTINATION_SHARD_AXIS] = shard_dim
60 + return cast(Placements, tuple(placements))
61 +
62 +
63 +@dataclass
64 +class M2NDestination:
65 + """Where one checkpoint parameter lands on this worker."""
66 +
67 + name: str
68 + placements: Placements | None
69 + """Placement over the inference mesh, or `REPLICATED` for the fallback."""
70 + tensor: torch.Tensor | None
71 + """The live parameter view to reshard into, or None for the fallback path
72 + (the engine allocates a full replica per round and calls `load_weights`)."""
73 +
74 + @property
75 + def direct(self) -> bool:
76 + return self.tensor is not None
77 +
78 +
79 +def _base_loader_shard_dim(
80 + param: torch.nn.Parameter,
81 + shard_axis_size: int,
82 +) -> int | None:
83 + """Return a known-safe loader's declared shard dim, or reject it.
84 +
85 + Loader identity is deliberately fail-closed. A missing loader is not
86 + treated as ``default_weight_loader`` because a model-level ``load_weights``
87 + implementation may transform the tensor before it reaches that default.
88 + Wrappers and partials are also rejected: recognizing their wrapped callable
89 + would ignore behavior added by the wrapper.
90 + """
91 + loader = getattr(param, "weight_loader", None)
92 + if loader is None or isinstance(loader, partial):
93 + return None
94 +
95 + owner = loader.__self__ if inspect.ismethod(loader) else None
96 + loader_fn = loader.__func__ if inspect.ismethod(loader) else loader
97 +
98 + # Lazy imports avoid making the weight-transfer registry import all model
99 + # executor layers during normal vLLM startup.
100 + from vllm.model_executor.layers.linear import (
101 + ColumnParallelLinear,
102 + ReplicatedLinear,
103 + RowParallelLinear,
104 + )
105 + from vllm.model_executor.model_loader.weight_utils import default_weight_loader
106 +
107 + replicated_loaders = {
108 + default_weight_loader,
109 + ReplicatedLinear.weight_loader,
110 + }
111 + column_loaders = {
112 + ColumnParallelLinear.weight_loader,
113 + ColumnParallelLinear.weight_loader_v2,
114 + }
115 + row_loaders = {
116 + RowParallelLinear.weight_loader,
117 + RowParallelLinear.weight_loader_v2,
118 + }
119 +
120 + # These attributes signal packing, a transform, or pre-sharded checkpoint
121 + # storage. Such semantics cannot be reproduced by a plain M2N placement.
122 + # Check them before accepting replicated loaders too: model-level loading
123 + # may transform a tensor before handing it to an otherwise plain copier.
124 + metadata_attrs = ("packed_dim", "packed_factor", "marlin_tile_size")
125 + flag_attrs = ("needs_scalar_to_array", "is_sharded_weight", "is_transposed")
126 + if any(getattr(param, attr, None) is not None for attr in metadata_attrs):
127 + return None
128 + if any(bool(getattr(param, attr, False)) for attr in flag_attrs):
129 + return None
130 +
131 + if loader_fn in replicated_loaders:
132 + return REPLICATE
133 + if loader_fn not in column_loaders | row_loaders:
134 + return None
135 +
136 + if owner is None or getattr(owner, "tp_size", None) != shard_axis_size:
137 + return None
138 +
139 + attr = "output_dim" if loader_fn in column_loaders else "input_dim"
140 + dim = getattr(param, attr, None)
141 + return dim if isinstance(dim, int) and dim >= 0 else None
142 +
143 +
144 +def _validated_shard_dim(
145 + global_shape: Sequence[int],
146 + local_shape: Sequence[int],
147 + declared_dim: int,
148 + shard_axis_size: int,
149 +) -> int | None:
150 + """Validate a loader-declared shard dimension against both shapes."""
151 + if len(global_shape) != len(local_shape):
152 + return None
153 +
154 + if declared_dim == REPLICATE:
155 + return REPLICATE if tuple(global_shape) == tuple(local_shape) else None
156 + if declared_dim >= len(global_shape):
157 + return None
158 + for dim, (whole, local) in enumerate(zip(global_shape, local_shape)):
159 + expected = local * shard_axis_size if dim == declared_dim else local
160 + if whole != expected:
161 + return None
162 + return declared_dim
163 +
164 +
165 +def resolve_parameter_destinations(
166 + model: torch.nn.Module,
167 + names: Sequence[str],
168 + dtypes: Sequence[torch.dtype],
169 + shapes: Sequence[Sequence[int]],
170 + *,
171 + num_workers: int,
172 + shard_axis_size: int,
173 + allow_direct: bool,
174 +) -> list[M2NDestination]:
175 + """Build this worker's destination plan, one entry per checkpoint parameter.
176 +
177 + Placements are relative to the inference mesh: axis 0 replicates and axis 1
178 + shards. `allow_direct=False` forces every parameter onto the fallback path;
179 + the engine sets it for pipeline-parallel or quantized deployments, where a
180 + parameter's local shape is not simply the checkpoint shape split across the
181 + shard axis. Parameters below a model type with known model-level checkpoint
182 + preprocessing also fall back, since bypassing `load_weights` would skip the
183 + transformation.
184 + """
185 + params = dict(model.named_parameters()) if allow_direct else {}
186 + blocked_param_ids = {
187 + id(param)
188 + for module in model.modules()
189 + if _module_type_name(module) in _DIRECT_WEIGHT_LOADING_BLOCKLIST
190 + for param in module.parameters()
191 + }
192 +
193 + destinations: list[M2NDestination] = []
194 + for name, dtype, shape in zip(names, dtypes, shapes):
195 + param = params.get(name)
196 + dim = None
197 + if (
198 + param is not None
199 + and id(param) not in blocked_param_ids
200 + and param.dtype == dtype
201 + and param.data.is_contiguous()
202 + and num_workers % shard_axis_size == 0
203 + ):
204 + declared_dim = _base_loader_shard_dim(param, shard_axis_size)
205 + if declared_dim is not None:
206 + dim = _validated_shard_dim(
207 + shape, param.shape, declared_dim, shard_axis_size
208 + )
209 +
210 + if dim is None:
211 + # A replicated parameter is identical on every rank, so it needs no
212 + # placement of its own — REPLICATED lets resolve_layout spread it
213 + # over all the workers regardless of how the mesh is factored.
214 + destinations.append(M2NDestination(name, REPLICATED, None))
215 + elif dim == REPLICATE:
216 + assert param is not None
217 + destinations.append(M2NDestination(name, REPLICATED, param.data))
218 + else:
219 + assert param is not None
220 + destinations.append(
221 + M2NDestination(name, _destination_placements(dim), param.data)
222 + )
223 +
224 + # Layerwise reload finalizes a module as a unit. Mixing direct and fallback
225 + # parameters within one module can copy stale storage over the direct load.
226 + fallback_modules = {
227 + name.rpartition(".")[0]
228 + for name, destination in zip(names, destinations)
229 + if not destination.direct and name in params
230 + }
231 + destinations = [
232 + M2NDestination(destination.name, REPLICATED, None)
233 + if destination.direct
234 + and destination.name.rpartition(".")[0] in fallback_modules
235 + else destination
236 + for destination in destinations
237 + ]
238 +
239 + num_direct = sum(d.direct for d in destinations)
240 + parameter_bytes = [
241 + math.prod(shape) * dtype.itemsize for dtype, shape in zip(dtypes, shapes)
242 + ]
243 + total_bytes = sum(parameter_bytes)
244 + direct_bytes = sum(
245 + nbytes
246 + for destination, nbytes in zip(destinations, parameter_bytes)
247 + if destination.direct
248 + )
249 + direct_byte_percentage = 100 * direct_bytes / total_bytes if total_bytes else 0
250 + logger.info(
251 + "nccl_m2n destination plan: %d/%d parameters resharded directly into "
252 + "the model, %d via full-tensor fallback; direct byte coverage: "
253 + "%d/%d bytes (%.1f%%)",
254 + num_direct,
255 + len(destinations),
256 + len(destinations) - num_direct,
257 + direct_bytes,
258 + total_bytes,
259 + direct_byte_percentage,
260 + )
261 + return destinations
vllm/distributed/weight_transfer/m2n_source.py new
+155
@@ -0,0 +1,155 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Trainer-side weight sources for the NCCL M2N backend.
4 +
5 +m2n plans a transfer from both sides' layouts, so the trainer has to say how it
6 +holds each parameter — which the base `WeightSource` / `ParamMeta` pair does not
7 +express. This module supplies the m2n flavor of both, and the DTensor-backed
8 +implementation that covers the common trainer.
9 +
10 +A source declares its mesh once (`mesh()`) and a placement per parameter, since
11 +one topology describes every tensor on the side.
12 +"""
13 +
14 +from collections.abc import Iterator
15 +from typing import Any
16 +
17 +import torch
18 +
19 +from vllm.distributed.weight_transfer.base import ParamMeta, WeightSource
20 +from vllm.distributed.weight_transfer.m2n_common import (
21 + REPLICATE,
22 + REPLICATED,
23 + M2NMesh,
24 + M2NParamMeta,
25 + Placements,
26 +)
27 +
28 +__all__ = [
29 + "M2NWeightSource",
30 + "DTensorModuleSource",
31 + "mesh_from_tensor",
32 + "placements_from_tensor",
33 +]
34 +
35 +
36 +class M2NWeightSource(WeightSource):
37 + """A `WeightSource` that also describes how the trainer holds its weights.
38 +
39 + Unlike `ModuleSource`, iteration yields each rank's **local shard**, not a
40 + materialized full tensor: gathering is exactly the cost m2n removes.
41 + """
42 +
43 + def mesh(self) -> M2NMesh:
44 + """The trainer's rank topology, shared by every parameter."""
45 + raise NotImplementedError
46 +
47 + def metadata(self) -> list[ParamMeta]:
48 + """Name, dtype, full shape, and trainer placement for each parameter."""
49 + raise NotImplementedError
50 +
51 + def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
52 + """Yield `(name, local shard)` pairs in the same order as `metadata()`."""
53 + raise NotImplementedError
54 +
55 +
56 +def _placement_code(placement: Any) -> int:
57 + """Map a `torch.distributed` placement onto an m2n placement code."""
58 + name = type(placement).__name__
59 + if name == "Replicate":
60 + return REPLICATE
61 + if name == "Shard":
62 + return int(placement.dim)
63 + raise ValueError(
64 + f"nccl_m2n cannot express the {name} placement; only Replicate and "
65 + "Shard are supported"
66 + )
67 +
68 +
69 +def mesh_from_tensor(tensor: torch.Tensor, num_trainer_ranks: int) -> M2NMesh:
70 + """The mesh a parameter lives on, as an `M2NMesh`.
71 +
72 + A plain tensor is identical on every trainer rank, so it spans all of them.
73 + """
74 + device_mesh = getattr(tensor, "device_mesh", None)
75 + if device_mesh is None:
76 + return M2NMesh((num_trainer_ranks, 1), 0)
77 +
78 + grid = device_mesh.mesh
79 + ranks = grid.flatten().tolist()
80 + if ranks != list(range(len(ranks))):
81 + raise ValueError(
82 + "nccl_m2n requires the trainer's device mesh to cover the "
83 + f"contiguous rank interval [0, {len(ranks)}); got {ranks}"
84 + )
85 + if grid.ndim > 2:
86 + raise ValueError(
87 + f"nccl_m2n supports 1-D and 2-D device meshes, got {grid.ndim}-D"
88 + )
89 + dims = tuple(grid.shape)
90 + return M2NMesh(dims if len(dims) == 2 else (dims[0], 1), 0)
91 +
92 +
93 +def placements_from_tensor(tensor: torch.Tensor) -> Placements | None:
94 + """How a parameter is placed over its mesh, or `REPLICATED`.
95 +
96 + `REPLICATED` is returned rather than a `(REPLICATE, REPLICATE)` pair: m2n
97 + cannot express that, and the size-1-axis encoding it needs instead is
98 + applied later by `resolve_layout`, which owns that workaround.
99 + """
100 + placements = getattr(tensor, "placements", None)
101 + if placements is None:
102 + return REPLICATED
103 + codes = [_placement_code(p) for p in placements]
104 + if all(code == REPLICATE for code in codes):
105 + return REPLICATED
106 + if len(codes) == 1:
107 + return (codes[0], REPLICATE)
108 + return (codes[0], codes[1])
109 +
110 +
111 +class DTensorModuleSource(M2NWeightSource):
112 + """`M2NWeightSource` over `module.named_parameters()`.
113 +
114 + Covers both the FSDP/DTensor trainer (placement read off each parameter,
115 + local shard yielded via `to_local()`) and the plain replicated trainer, with
116 + no special casing. Trainers with a custom producer (a Megatron export, MoE
117 + re-fusing) subclass `M2NWeightSource` instead.
118 + """
119 +
120 + def __init__(self, module: torch.nn.Module, num_trainer_ranks: int) -> None:
121 + """Wrap `module`; `num_trainer_ranks` sizes the mesh for plain tensors."""
122 + self._module = module
123 + self._num_trainer_ranks = num_trainer_ranks
124 +
125 + def mesh(self) -> M2NMesh:
126 + """The mesh every sharded parameter agrees on.
127 +
128 + Replicated parameters span the trainer without a device mesh of their
129 + own, so they never decide this; a model whose *sharded* parameters
130 + disagree cannot be described by one side mesh and is rejected.
131 + """
132 + meshes = {
133 + mesh_from_tensor(p, self._num_trainer_ranks)
134 + for _, p in self._module.named_parameters()
135 + if placements_from_tensor(p) is not REPLICATED
136 + }
137 + if len(meshes) > 1:
138 + raise ValueError(
139 + "nccl_m2n needs one mesh per side, but this module's sharded "
140 + f"parameters span several: {sorted(m.dims for m in meshes)}"
141 + )
142 + return meshes.pop() if meshes else M2NMesh((self._num_trainer_ranks, 1), 0)
143 +
144 + def metadata(self) -> list[ParamMeta]:
145 + """Read global shape/dtype and placements without gathering shards."""
146 + return [
147 + M2NParamMeta(name, p.dtype, tuple(p.shape), placements_from_tensor(p))
148 + for name, p in self._module.named_parameters()
149 + ]
150 +
151 + def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
152 + """Yield each parameter's local shard (`to_local()`), or the tensor itself."""
153 + for name, param in self._module.named_parameters():
154 + to_local = getattr(param, "to_local", None)
155 + yield name, (to_local() if callable(to_local) else param)
vllm/distributed/weight_transfer/m2n_trainer.py new
+414
@@ -0,0 +1,414 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Trainer-side weight transfer engine for the NCCL M2N backend.
4 +
5 +Symmetric to `M2NWeightTransferEngine` but in the training process. Every
6 +trainer rank joins the shared communicator and runs every reshard, sending its
7 +own local shard; only rank 0 touches the inference control plane.
8 +
9 +Workers join the communicator during init and run reshard from inside
10 +`update_weights`, so those two RPCs overlap with work on this side. Both run on
11 +a helper thread and are joined afterwards, which keeps that overlap inside the
12 +engine -- where the `TrainerWeightTransferEngine` contract puts it -- rather than
13 +in every caller.
14 +"""
15 +
16 +from concurrent.futures import Future, ThreadPoolExecutor
17 +from dataclasses import dataclass, field
18 +from typing import TYPE_CHECKING, Any, ClassVar, cast
19 +
20 +import torch
21 +from typing_extensions import Self
22 +
23 +from vllm.distributed.weight_transfer.base import (
24 + TrainerInitInfo,
25 + TrainerWeightTransferEngine,
26 + VLLMWeightSyncClient,
27 + WeightSource,
28 +)
29 +from vllm.distributed.weight_transfer.m2n_common import (
30 + M2NMesh,
31 + M2NParamMeta,
32 + Placements,
33 + check_transferable,
34 + comm_ptr,
35 + import_m2n,
36 + publish_destination_placements,
37 + resolve_layout,
38 + to_mesh,
39 + to_placements,
40 + validate_layout,
41 +)
42 +from vllm.distributed.weight_transfer.m2n_source import M2NWeightSource
43 +from vllm.distributed.weight_transfer.nccl_common import (
44 + decode_nccl_unique_id,
45 + trainer_init,
46 + uid_init_process_group,
47 +)
48 +from vllm.logger import init_logger
49 +
50 +if TYPE_CHECKING:
51 + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
52 +
53 +logger = init_logger(__name__)
54 +
55 +__all__ = ["M2NTrainerInitInfo", "M2NTrainerWeightTransferEngine"]
56 +
57 +
58 +def _dtype_name(dtype: torch.dtype) -> str:
59 + """Short name for a dtype, e.g. `bfloat16`, for the worker-side init dict."""
60 + return str(dtype).split(".")[-1]
61 +
62 +
63 +@dataclass(kw_only=True)
64 +class M2NTrainerInitInfo(TrainerInitInfo):
65 + """Trainer-side init info for nccl_m2n.
66 +
67 + `rank` (from `TrainerInitInfo`) is this trainer process's rank; it is also
68 + its rank in the shared communicator, since the trainer occupies
69 + `[0, num_trainer_ranks)`. Rank 0 drives the control plane.
70 + """
71 +
72 + backend: ClassVar[str] = "nccl_m2n"
73 +
74 + master_address: str | None = None
75 + master_port: int | None = None
76 + world_size: int
77 + """Trainer ranks + all inference workers."""
78 + num_trainer_ranks: int = 1
79 + dst_mesh_dims: tuple[int, int] | None = None
80 + """How the inference ranks are laid out: axis 0 replicates and axis 1
81 + shards. The trainer declares it so both sides describe the destination
82 + identically; defaults to a flat `(num_workers, 1)`, which is all a
83 + replicated destination needs."""
84 + max_cta: int | None = None
85 + nccl_unique_id_b64: str | None = field(default=None, repr=False)
86 +
87 + @property
88 + def destination_mesh_dims(self) -> tuple[int, int]:
89 + """`dst_mesh_dims`, or a flat mesh over every inference worker."""
90 + if self.dst_mesh_dims is not None:
91 + return self.dst_mesh_dims
92 + return (self.world_size - self.num_trainer_ranks, 1)
93 +
94 + def __post_init__(self) -> None:
95 + """Reject rank counts or destination meshes that cannot form a group."""
96 + _ = self.nccl_unique_id_bytes
97 + if not 0 < self.num_trainer_ranks < self.world_size:
98 + raise ValueError(
99 + f"`num_trainer_ranks` ({self.num_trainer_ranks}) must leave at "
100 + f"least one worker in world_size {self.world_size}"
101 + )
102 + num_workers = self.world_size - self.num_trainer_ranks
103 + if self.dst_mesh_dims is not None:
104 + dims = tuple(self.dst_mesh_dims)
105 + if len(dims) != 2 or dims[0] * dims[1] != num_workers:
106 + raise ValueError(
107 + f"`dst_mesh_dims` {dims} must be two dims covering the "
108 + f"{num_workers} inference workers"
109 + )
110 + if not 0 <= self.rank < self.num_trainer_ranks:
111 + raise ValueError(
112 + f"trainer `rank` ({self.rank}) must be below "
113 + f"`num_trainer_ranks` ({self.num_trainer_ranks})"
114 + )
115 +
116 + @property
117 + def nccl_unique_id_bytes(self) -> bytes | None:
118 + return decode_nccl_unique_id(
119 + master_address=self.master_address,
120 + master_port=self.master_port,
121 + nccl_unique_id_b64=self.nccl_unique_id_b64,
122 + ctx="M2NTrainerInitInfo",
123 + )
124 +
125 +
126 +class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitInfo]):
127 + """Trainer-side engine: sends every rank's local shard, once per parameter.
128 +
129 + Called on every trainer rank. All ranks join the communicator and run every
130 + reshard; only rank 0 touches the control plane.
131 + """
132 +
133 + init_info_cls = M2NTrainerInitInfo
134 +
135 + def __init__(
136 + self,
137 + *,
138 + client: VLLMWeightSyncClient,
139 + source: M2NWeightSource,
140 + is_controller: bool,
141 + num_trainer_ranks: int,
142 + ) -> None:
143 + """Hold the client and source; `trainer_init` fills in the group."""
144 + super().__init__(client=client, source=source, is_sender=is_controller)
145 + self.source: M2NWeightSource = source
146 + self.is_controller = is_controller
147 + self.num_trainer_ranks = num_trainer_ranks
148 + self.group: PyNcclCommunicator | None = None
149 + self._m2n: Any = None
150 + self._handle: Any = None
151 + self._metas: list[M2NParamMeta] = []
152 + self._src_mesh: M2NMesh | None = None
153 + self._dst_mesh: M2NMesh | None = None
154 + self._dst_placements: list[Placements | None] = []
155 + self._executor: ThreadPoolExecutor | None = None
156 +
157 + @classmethod
158 + def trainer_init(
159 + cls,
160 + init_info: M2NTrainerInitInfo,
161 + *,
162 + client: VLLMWeightSyncClient,
163 + source: WeightSource | None = None,
164 + ) -> Self:
165 + """Build the engine and rendezvous with the inference side.
166 +
167 + Runs on *every* trainer rank. Rank 0 additionally ships the transfer
168 + plan to the workers and drives the control plane; the other ranks only
169 + build local state and join the communicator. Every rank participates
170 + in every reshard and sends its local shard.
171 +
172 + The ordering here is not arbitrary -- see the comments inline. In
173 + short: validate before anything blocks; start the init RPC without
174 + waiting (workers join the communicator inside that RPC, so it cannot
175 + return until this side has joined too); then join the NCCL
176 + communicator; then create the M2N handle; and only then wait for the RPC
177 + to finish.
178 + """
179 + # The base WeightSource carries name/dtype/shape but not layout, which
180 + # is not enough to plan a reshard. Fail here rather than at the first
181 + # missing attribute deep inside the send loop.
182 + if not isinstance(source, M2NWeightSource):
183 + raise TypeError(
184 + "nccl_m2n needs per-parameter layouts, so its source must be an "
185 + f"M2NWeightSource (e.g. DTensorModuleSource); got "
186 + f"{type(source).__name__}"
187 + )
188 +
189 + # Trainer-side weight transfer engine for this rank. `__init__` only
190 + # holds the client and source; the communicator, meshes, and M2N handle
191 + # are filled in below.
192 + engine = cls(
193 + client=client,
194 + source=source,
195 + is_controller=init_info.rank == 0,
196 + num_trainer_ranks=init_info.num_trainer_ranks,
197 + )
198 + engine._m2n = import_m2n()
199 +
200 + engine._src_mesh = source.mesh()
201 + if engine._src_mesh.size != init_info.num_trainer_ranks:
202 + raise ValueError(
203 + f"source mesh covers {engine._src_mesh.size} ranks, but there "
204 + f"are {init_info.num_trainer_ranks} trainer ranks"
205 + )
206 +
207 + # Validate the whole plan before any rendezvous. A layout m2n cannot
208 + # express is a mismatched collective at send time, which hangs rather
209 + # than raises -- so every precondition is checked while an exception can
210 + # still propagate to the caller.
211 + engine._metas = cast(list[M2NParamMeta], source.metadata())
212 + for meta in engine._metas:
213 + check_transferable(meta.name, meta.dtype, meta.shape)
214 + mesh, placements = resolve_layout(engine._src_mesh, meta.placements)
215 + validate_layout(mesh, placements, meta.shape, "source")
216 +
217 + rendezvous: Future | None = None
218 + if engine.is_controller:
219 + engine._executor = ThreadPoolExecutor(
220 + max_workers=1, thread_name_prefix="m2n-weight-sync"
221 + )
222 + # Only trainer rank 0 starts the control-plane RPC. Workers join
223 + # the communicator inside it, so submit it first, join below, and
224 + # wait for it at the end.
225 + rendezvous = engine._executor.submit(
226 + client.init_weight_transfer_engine,
227 + engine._worker_init_info(init_info),
228 + )
229 +
230 + # Every trainer rank joins -- reshard is collective over the whole
231 + # communicator. Trainer ranks take
232 + # [0, num_trainer_ranks) and the workers follow, which is the
233 + # contiguous-interval layout m2n requires of both meshes.
234 + unique_id_bytes = init_info.nccl_unique_id_bytes
235 + if unique_id_bytes is not None:
236 + engine.group = uid_init_process_group(
237 + unique_id_bytes,
238 + rank=init_info.rank,
239 + world_size=init_info.world_size,
240 + device=torch.accelerator.current_device_index(),
241 + )
242 + else:
243 + assert init_info.master_address is not None
244 + assert init_info.master_port is not None
245 + engine.group = trainer_init(
246 + {
247 + "master_address": init_info.master_address,
248 + "master_port": init_info.master_port,
249 + "world_size": init_info.world_size,
250 + },
251 + rank=init_info.rank,
252 + )
253 +
254 + # The trainer declares the inference topology; the worker uses exactly
255 + # this, so neither side has to infer the other's factorization.
256 + engine._dst_mesh = M2NMesh(
257 + init_info.destination_mesh_dims, init_info.num_trainer_ranks
258 + )
259 + engine._dst_placements = engine._receive_destination_plan(
260 + init_info.num_trainer_ranks
261 + )
262 + logger.info(
263 + "nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), "
264 + "inference ranks [%d, %d)",
265 + len(engine._metas),
266 + init_info.num_trainer_ranks,
267 + init_info.num_trainer_ranks,
268 + init_info.world_size,
269 + )
270 +
271 + # Created after the group so a failed rendezvous does not leave a live
272 + # handle behind; shutdown() releases it once caller streams are synced.
273 + engine._handle = engine._m2n.Handle.create(
274 + engine._m2n.Config(max_cta=init_info.max_cta)
275 + )
276 +
277 + # Now safe to collect: both sides have joined, so the RPC has returned
278 + # or raised. Doing this last means a worker-side init failure surfaces
279 + # here, as an exception from trainer_init, rather than at the first send.
280 + if rendezvous is not None:
281 + rendezvous.result()
282 + return engine
283 +
284 + def _receive_destination_plan(
285 + self, first_worker_rank: int
286 + ) -> list[Placements | None]:
287 + """Receive and validate the inference-side per-parameter placements."""
288 + if self.group is None or self._dst_mesh is None:
289 + raise RuntimeError(
290 + "destination plan requires the communicator and destination mesh"
291 + )
292 +
293 + # The first inference worker publishes one placement per parameter.
294 + # Every trainer rank receives the complete plan in this single,
295 + # initialization-time broadcast.
296 + placements = publish_destination_placements(
297 + self.group, first_worker_rank, None, len(self._metas)
298 + )
299 + # The send loop matches placements to metadata by position, so reject
300 + # an incomplete or oversized plan before entering any M2N collective.
301 + if len(placements) != len(self._metas):
302 + raise RuntimeError(
303 + f"inference side planned {len(placements)} parameters, but "
304 + f"this trainer announced {len(self._metas)}"
305 + )
306 + # Check every worker-selected destination against the corresponding
307 + # global parameter shape before the first transfer.
308 + for meta, placement in zip(self._metas, placements):
309 + mesh, resolved = resolve_layout(self._dst_mesh, placement)
310 + validate_layout(mesh, resolved, meta.shape, "destination")
311 + return placements
312 +
313 + def _worker_init_info(self, init_info: M2NTrainerInitInfo) -> dict[str, Any]:
314 + """Handshake payload: rendezvous, both meshes, and the transfer plan."""
315 + assert self._src_mesh is not None
316 + payload: dict[str, Any] = {
317 + "master_address": init_info.master_address,
318 + "master_port": init_info.master_port,
319 + "nccl_unique_id_b64": init_info.nccl_unique_id_b64,
320 + "rank_offset": init_info.num_trainer_ranks,
321 + "world_size": init_info.world_size,
322 + "src_mesh_dims": list(self._src_mesh.dims),
323 + "dst_mesh_dims": list(init_info.destination_mesh_dims),
324 + "names": [m.name for m in self._metas],
325 + "dtype_names": [_dtype_name(m.dtype) for m in self._metas],
326 + "shapes": [list(m.shape) for m in self._metas],
327 + "src_placements": [
328 + None if m.placements is None else list(m.placements)
329 + for m in self._metas
330 + ],
331 + "max_cta": init_info.max_cta,
332 + }
333 + return {key: value for key, value in payload.items() if value is not None}
334 +
335 + def send_weights(self) -> None:
336 + """Drive one update round: start, reshard concurrently, then finish."""
337 + if self._handle is None or self.group is None:
338 + raise RuntimeError("nccl_m2n trainer engine is not initialized")
339 +
340 + update: Future | None = None
341 + if self.is_controller:
342 + self.client.start_weight_update()
343 + # The workers reshard from inside `update_weights`, concurrently
344 + # with the sends below.
345 + assert self._executor is not None
346 + update = self._executor.submit(
347 + self.client.update_weights,
348 + {"names": [m.name for m in self._metas]},
349 + )
350 +
351 + try:
352 + self._send()
353 + finally:
354 + if update is not None:
355 + update.result()
356 +
357 + if self.is_controller:
358 + self.client.finish_weight_update()
359 +
360 + def _send(self) -> None:
361 + """Reshard each local shard into its worker-planned destination."""
362 + assert self.group is not None
363 + assert self._src_mesh is not None
364 + assert self._dst_mesh is not None
365 + m2n = self._m2n
366 + comm = comm_ptr(self.group)
367 + stream = torch.cuda.current_stream()
368 +
369 + for index, (name, tensor) in enumerate(self.source):
370 + meta = self._metas[index]
371 + if name != meta.name:
372 + raise RuntimeError(
373 + f"weight source yielded '{name}' at position {index}, but "
374 + f"its metadata declared '{meta.name}'; the source must "
375 + "iterate in a stable order"
376 + )
377 + shard = tensor.detach().contiguous()
378 + src_mesh, src_placements = resolve_layout(self._src_mesh, meta.placements)
379 + dst_mesh, dst_placements = resolve_layout(
380 + self._dst_mesh, self._dst_placements[index]
381 + )
382 + m2n.reshard(
383 + shard,
384 + None,
385 + comm,
386 + stream,
387 + src_mesh=to_mesh(m2n, src_mesh),
388 + src_placements=to_placements(m2n, src_placements),
389 + dst_mesh=to_mesh(m2n, dst_mesh),
390 + dst_placements=to_placements(m2n, dst_placements),
391 + handle=self._handle,
392 + )
393 + # A custom source may yield temporary storage allocated on another
394 + # stream. M2N retains only its device pointer for asynchronous work,
395 + # so tell the allocator not to reuse that storage until this stream
396 + # has finished with it. For storage allocated on this stream,
397 + # record_stream is effectively a no-op.
398 + shard.record_stream(stream)
399 +
400 + # As in the other weight-transfer backends, wait before returning
401 + # rather than relying on same-stream ordering. The caller may then
402 + # safely mutate parameters or continue work on another stream.
403 + stream.synchronize()
404 +
405 + def shutdown(self) -> None:
406 + """Destroy the m2n handle, join the helper thread, and drop the group."""
407 + if self._handle is not None:
408 + torch.accelerator.synchronize()
409 + self._handle.destroy()
410 + self._handle = None
411 + if self._executor is not None:
412 + self._executor.shutdown()
413 + self._executor = None
414 + self.group = None
vllm/distributed/weight_transfer/nccl_common.py
+6 -4
@@ -220,11 +220,11 @@ def worker_init_process_group(
220
221 def trainer_init(
222 init_info: NCCLRendezvous | dict,
223 + rank: int = 0,
224 ) -> "PyNcclCommunicator":
225 """Initialize NCCL process group for trainer-side weight transfer.
226
226 - The trainer is always rank 0 in the process group. Uses the current
227 - CUDA device (torch.accelerator.current_device_index()).
227 + Uses the current CUDA device (torch.accelerator.current_device_index()).
228
229 Args:
230 init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or
@@ -232,6 +232,9 @@ def trainer_init(
232 - master_address: str
233 - master_port: int
234 - world_size: int
235 + rank: This trainer process's rank in the group. The broadcast backends
236 + have a single trainer at rank 0; multi-rank trainers (m2n) occupy
237 + ranks `[0, num_trainer_ranks)` and the workers start after them.
238
239 Returns:
240 PyNcclCommunicator for weight transfer.
@@ -246,12 +249,11 @@ def trainer_init(
249 master_port = init_info.master_port
250 world_size = init_info.world_size
251
249 - # Trainer is always rank 0
252 device = torch.accelerator.current_device_index()
253 return stateless_init_process_group(
254 master_address,
255 master_port,
254 - 0,
256 + rank,
257 world_size,
258 device,
259 )