[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>

.buildkite/test_areas/distributed.yaml

@@ -373,6 +373,7 @@ steps:
     - 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
     - pytest -v -s tests/v1/distributed/test_dbo.py
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
+    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer_m2n.py
     - pytest -v -s tests/distributed/test_packed_tensor.py
   mirror:
     amd:

docs/training/weight_transfer/README.md

@@ -35,6 +35,7 @@ engine drives on your behalf:
 | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
 | [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches |
 | [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) |
+| [nccl_m2n](m2n.md) | NCCL M2N reshard | Trainer and inference use different sharding layouts |
 
 ## Quickstart
 
@@ -49,7 +50,8 @@ from vllm.config import WeightTransferConfig
 
 llm = LLM(
     model="my-model",
-    weight_transfer_config=WeightTransferConfig(backend="nccl"),  # or "ipc", "sparse_nccl", "sharded_rdt"
+    # Other backends: "ipc", "sparse_nccl", "sharded_rdt", "nccl_m2n".
+    weight_transfer_config=WeightTransferConfig(backend="nccl"),
 )
 ```
 

docs/training/weight_transfer/m2n.md

@@ -0,0 +1,101 @@
+# NCCL M2N Engine
+
+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.
+
+## When to Use NCCL M2N
+
+- 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
+- Training and inference are on **separate GPUs**, possibly across nodes
+- You want weight sync off the critical path of an async RL loop
+
+## Requirements
+
+This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. It needs NCCL 2.30.5 or newer, and `libnccl_m2n.so` is linked against a specific NCCL build, so vLLM must load the same library:
+
+```bash
+export VLLM_NCCL_SO_PATH=/path/to/the/same/libnccl.so
+```
+
+The backend imports the runtime lazily, so vLLM is unaffected unless you select it.
+
+## How It Works
+
+1. The trainer ranks and all inference workers join **one** NCCL communicator via `StatelessProcessGroup`. Trainer ranks occupy `[0, T)` and workers `[T, T + N)` — m2n meshes are contiguous rank intervals, which is exactly what vLLM's existing `rank_offset` convention produces.
+2. At the init handshake the trainer ships the full transfer plan: both meshes — one per side, shared by every parameter — then each parameter's name, dtype, shape and placement. The workers resolve the matching **destination placements** and publish them back over the shared group; the meshes themselves are declared, not negotiated.
+3. Each round, both sides issue one `reshard` per parameter in the same order. Every rank in the communicator participates in every reshard, so — as with the NCCL engine — the worker must be inside `update_weights` while the trainer sends. The trainer engine handles that concurrency itself.
+
+### Destination resolution
+
+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.
+
+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:
+
+- fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`), whose checkpoint names do not name a model parameter
+- quantized models, where the weight loader does more than a sharded copy
+- pipeline-parallel deployments
+
+Correctness never depends on a parameter resolving. The engine logs how many parameters took each path at init.
+
+## Configuration
+
+```python
+from vllm import LLM
+from vllm.config import WeightTransferConfig
+
+llm = LLM(
+    model="my-model",
+    weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"),
+)
+```
+
+```bash
+vllm serve my-model --weight-transfer-config '{"backend": "nccl_m2n"}'
+```
+
+## Trainer Side
+
+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()`.
+
+```python
+from vllm.distributed.weight_transfer import (
+    RayVLLMWeightSyncClient,
+    WeightTransferTrainerFactory,
+)
+from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource
+from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
+
+# Called on every trainer rank; rank 0 is the sender.
+engine = WeightTransferTrainerFactory.trainer_init(
+    init_info=M2NTrainerInitInfo(
+        master_address=master_address,
+        master_port=master_port,
+        world_size=num_trainer_ranks + num_inference_workers,
+        num_trainer_ranks=num_trainer_ranks,
+        rank=my_rank,
+    ),
+    client=RayVLLMWeightSyncClient(llm_handle),
+    source=DTensorModuleSource(model, num_trainer_ranks),
+)
+
+# Each round, on every trainer rank:
+engine.send_weights()
+```
+
+`send_weights()` drives the whole round — `start_weight_update`, `update_weights`, `finish_weight_update`. Every rank in the communicator must be inside the same reshard, so `update_weights` has to be in flight while the trainer sends; the engine runs it on a helper thread so the trainer loop does not have to open-code that.
+
+Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass `M2NWeightSource` and supply layouts explicitly.
+
+See [`examples/rl/rlhf_m2n.py`](../../../examples/rl/rlhf_m2n.py) for a runnable FSDP → TP example.
+
+## Limitations
+
+- Tensor rank 1..3; 4-D parameters are rejected.
+- dtypes: int8/uint8, fp8 e4m3/e5m2, fp16, bf16, int32/uint32, fp32, int64/uint64, fp64. No fp4.
+- The trainer's device mesh must cover a contiguous rank interval starting at 0, and be 1-D or 2-D.
+- `Partial` placements are not supported.
+- Models with known model-level checkpoint preprocessing use the full-tensor
+  fallback for the affected module.
+- A reshard cannot be issued from inside a CUDA graph capture (weight updates run outside capture, so this only matters for custom callers).
+- The reshard plan is capped per *shard*, not per mesh: at most 16 source shards may feed one destination shard, and one source shard may feed at most 64 destination shards (`MAX_SOURCES` / `MAX_TARGETS` in the m2n build; raising either needs a rebuild).
+- A directly-resharded parameter spreads over `tp_size` destination shards, so each one is fed by fewer source shards. A parameter on the fallback path is replicated, i.e. a single destination shard fed by every source shard, so a trainer that **sharded** its parameters is still limited to 16 ranks whenever any parameter falls back — which fused parameters always do. A replicated trainer is unaffected.
+- Each communicator holds a cached staging pool (~256 MiB by default: 4 channels x 64 MiB).

examples/rl/rlhf_m2n.py

@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""RLHF with FSDP2 training and vLLM tensor-parallel inference using **NCCL M2N**
+sharding-aware weight transfer.
+
+Layout (4 GPUs, no colocation):
+  * GPUs 0-1: two FSDP2 training workers, one per GPU.
+  * GPUs 2-3: one vLLM ``LLM`` actor with ``tensor_parallel_size=2``.
+
+The trainer and the inference workers share one NCCL communicator of 4 ranks:
+trainer ranks ``[0, 2)`` and inference ranks ``[2, 4)``. Every parameter moves
+with a single ``reshard`` that redistributes it from the FSDP layout to the
+inference layout — the trainer sends its local shards and never all-gathers a
+full tensor, which broadcast-based weight sync would force it to do.
+
+Every FSDP rank builds an ``M2NTrainerWeightTransferEngine`` and calls
+``send_weights()``; all ranks run every reshard, and only rank 0 drives the
+inference side through its ``RayVLLMWeightSyncClient``.
+
+Requires the ``nccl-extensions`` package (NCCL 2.30.5+) and a
+``VLLM_NCCL_SO_PATH`` pointing at the same ``libnccl.so`` that
+``libnccl_m2n.so`` was linked against.
+
+This example was written for 4xH100.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pybase64 as base64
+import ray
+import torch
+import torch.distributed as dist
+from huggingface_hub import snapshot_download
+from torch.distributed.fsdp import fully_shard
+from transformers import AutoModelForCausalLM
+
+from vllm import LLM, SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
+from vllm.distributed.weight_transfer import (
+    RayVLLMWeightSyncClient,
+    WeightTransferTrainerFactory,
+)
+from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource
+from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
+from vllm.utils.network_utils import get_ip, get_open_port
+
+MODEL_NAME = "Qwen/Qwen3-0.6B"
+
+FSDP_WORLD_SIZE = 2
+INFERENCE_TP_SIZE = 2
+
+
+class MyLLM(LLM):
+    """LLM subclass that keeps Ray from pinning it to a single device."""
+
+    def __init__(self, *args, **kwargs):
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        super().__init__(*args, **kwargs)
+
+    def ready(self):
+        return True
+
+
+@ray.remote(num_cpus=1, num_gpus=1)
+class FSDPTrainWorker:
+    """One FSDP2 worker per GPU. Rank 0 is the weight-transfer sender."""
+
+    def __init__(
+        self,
+        model_name: str,
+        rank: int,
+        fsdp_master_addr: str,
+        fsdp_master_port: int,
+    ):
+        self.rank = rank
+
+        os.environ["MASTER_ADDR"] = fsdp_master_addr
+        os.environ["MASTER_PORT"] = str(fsdp_master_port)
+
+        dist.init_process_group(backend="nccl", rank=rank, world_size=FSDP_WORLD_SIZE)
+        torch.accelerator.set_device_index(0)
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16
+        ).cuda()
+        for layer in model.model.layers:
+            fully_shard(layer)
+        fully_shard(model)
+        self.model = model
+
+    def ready(self):
+        return True
+
+    def setup_engine(self, llm_handle, nccl_unique_id_b64, world_size, num_workers):
+        """Build the trainer engine on every FSDP rank.
+
+        `DTensorModuleSource` reads each parameter's FSDP device mesh and
+        placements, so the engine knows the source layout without gathering
+        anything. Rank 0 additionally drives the inference-side handshake.
+        """
+        self.engine = WeightTransferTrainerFactory.trainer_init(
+            init_info=M2NTrainerInitInfo(
+                nccl_unique_id_b64=nccl_unique_id_b64,
+                world_size=world_size,
+                num_trainer_ranks=FSDP_WORLD_SIZE,
+                # One DP group of TP=2 workers; declared so both sides
+                # describe the destination the same way.
+                dst_mesh_dims=(num_workers // INFERENCE_TP_SIZE, INFERENCE_TP_SIZE),
+                rank=self.rank,  # FSDP rank; sender is 0
+            ),
+            client=RayVLLMWeightSyncClient(llm_handle),
+            source=DTensorModuleSource(self.model, FSDP_WORLD_SIZE),
+        )
+
+    def send_weights(self):
+        """Called on all ranks concurrently; every rank runs every reshard."""
+        self.engine.send_weights()
+
+
+def main():
+    ray.init()
+
+    local_model_path = snapshot_download(MODEL_NAME)
+    print(f"[init] Model downloaded to {local_model_path}")
+
+    fsdp_master_addr = get_ip()
+    fsdp_master_port = get_open_port()
+
+    fsdp_workers = [
+        FSDPTrainWorker.remote(
+            local_model_path, rank, fsdp_master_addr, fsdp_master_port
+        )
+        for rank in range(FSDP_WORLD_SIZE)
+    ]
+    ray.get([w.ready.remote() for w in fsdp_workers])
+    print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")
+
+    llm = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote(
+        model=local_model_path,
+        enforce_eager=True,
+        tensor_parallel_size=INFERENCE_TP_SIZE,
+        distributed_executor_backend="ray",
+        weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"),
+        load_format="dummy",
+    )
+    ray.get(llm.ready.remote())
+    num_workers = ray.get(llm.get_world_size.remote())
+    print(f"[init] vLLM ready with {num_workers} inference workers.")
+
+    prompts = [
+        "Hello, my name is",
+        "The capital of France is",
+    ]
+    sampling_params = SamplingParams(temperature=0)
+
+    outputs = ray.get(llm.generate.remote(prompts, sampling_params))
+    print("-" * 60)
+    print("BEFORE weight sync (dummy weights):")
+    for output in outputs:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+    print("-" * 60)
+
+    # Trainer ranks [0, FSDP_WORLD_SIZE) and inference ranks after them share
+    # one communicator, so the two meshes are contiguous rank intervals.
+    nccl = NCCLLibrary()
+    nccl_unique_id_b64 = base64.b64encode(
+        bytes(nccl.ncclGetUniqueId().internal)
+    ).decode()
+    world_size = FSDP_WORLD_SIZE + num_workers
+
+    print("[transfer] Initializing nccl_m2n weight transfer (all FSDP ranks)...")
+    ray.get(
+        [
+            w.setup_engine.remote(llm, nccl_unique_id_b64, world_size, num_workers)
+            for w in fsdp_workers
+        ]
+    )
+
+    print("[sync] Resharding FSDP -> vLLM...")
+    ray.get([w.send_weights.remote() for w in fsdp_workers])
+    print("[sync] Weight transfer complete.")
+
+    outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
+    print("-" * 60)
+    print("AFTER weight sync (real weights):")
+    for output in outputs_updated:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+    print("-" * 60)
+
+
+if __name__ == "__main__":
+    main()

tests/distributed/test_weight_transfer_m2n.py

@@ -0,0 +1,868 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the NCCL M2N weight transfer backend.
+
+These cover how a transfer is *described* — layout encoding and the wire-type
+validation that keeps a bad plan from reaching a collective, since a mismatch
+there is a hang rather than an error. The transfer itself needs the m2n runtime
+and multiple GPUs, so it is exercised separately.
+"""
+
+import logging
+from functools import wraps
+from types import MethodType
+from unittest.mock import Mock
+
+import pybase64 as base64
+import pytest
+import ray
+import torch
+
+from vllm.distributed.weight_transfer import (
+    WeightTransferEngineFactory,
+    WeightTransferTrainerFactory,
+)
+from vllm.distributed.weight_transfer.m2n_common import (
+    REPLICATE,
+    REPLICATED,
+    M2NMesh,
+    M2NParamMeta,
+    check_placements,
+    check_transferable,
+    publish_destination_placements,
+    resolve_layout,
+    validate_layout,
+)
+from vllm.distributed.weight_transfer.m2n_engine import (
+    M2NWeightTransferEngine,
+    M2NWeightTransferInitInfo,
+    M2NWeightTransferUpdateInfo,
+)
+from vllm.distributed.weight_transfer.m2n_layout import (
+    M2NDestination,
+    resolve_parameter_destinations,
+)
+from vllm.distributed.weight_transfer.m2n_source import (
+    mesh_from_tensor,
+    placements_from_tensor,
+)
+from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.platforms import current_platform
+
+VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode()
+
+
+class TestLayout:
+    def test_replicated_splits_nothing(self):
+        """A replicated tensor imposes no divisibility constraint. The shape
+        dims are deliberately coprime with the mesh size: a layout that actually
+        split the tensor would reject them."""
+        mesh = M2NMesh((2, 2), start_rank=1)
+        indivisible_shape = (7, 13)  # neither dim divisible by any mesh axis
+
+        resolved_mesh, placements = resolve_layout(mesh, REPLICATED)
+        validate_layout(resolved_mesh, placements, indivisible_shape, "destination")
+
+    def test_replicated_keeps_the_same_ranks(self):
+        """Replication re-factors the mesh to get a size-1 axis for its no-op
+        shard. That is only sound if it still covers exactly the same GPUs."""
+        mesh = M2NMesh((2, 3), start_rank=4)
+        resolved_mesh, _ = resolve_layout(mesh, REPLICATED)
+        assert resolved_mesh.size == mesh.size
+        assert resolved_mesh.start_rank == mesh.start_rank
+
+    def test_sharded_keeps_its_own_factorization(self):
+        """Rank order decides who owns which shard, so a sharded tensor must
+        not be re-factored the way a replicated one is."""
+        mesh = M2NMesh((2, 3), start_rank=4)
+        resolved_mesh, placements = resolve_layout(mesh, (REPLICATE, 0))
+        assert resolved_mesh == mesh
+        assert placements == (REPLICATE, 0)
+
+    def test_two_shard_axes_rejected(self):
+        """One axis has to replicate; a 2-D mesh that shards both is not
+        something a single reshard can express."""
+        with pytest.raises(ValueError, match="shards both"):
+            check_placements((0, 1))
+
+    def test_placement_below_replicate_rejected(self):
+        with pytest.raises(ValueError, match=r"parameter 'w'.*-2"):
+            check_placements((-2, REPLICATE), "parameter 'w' source placements")
+
+    def test_shard_dim_must_exist(self):
+        with pytest.raises(ValueError, match="rank 2"):
+            validate_layout(M2NMesh((1, 2), 0), (REPLICATE, 2), (8, 16), "source")
+
+    def test_shard_must_divide_evenly(self):
+        with pytest.raises(ValueError, match="does not divide evenly"):
+            validate_layout(M2NMesh((1, 3), 0), (REPLICATE, 0), (8, 16), "source")
+
+
+class TestSourceLayout:
+    def test_plain_tensor_is_replicated_across_trainer_ranks(self):
+        """A tensor with no DTensor metadata is the same on every trainer rank,
+        so the source describes it as replicated over all of them."""
+        assert placements_from_tensor(torch.zeros(4)) is REPLICATED
+        assert mesh_from_tensor(torch.zeros(4), 4) == M2NMesh((4, 1), 0)
+
+
+class TestTransferable:
+    def test_unsupported_dtype_names_the_parameter(self):
+        with pytest.raises(ValueError, match="'w'"):
+            check_transferable("w", torch.complex64, (4,))
+
+    def test_rank_four_rejected(self):
+        with pytest.raises(ValueError, match="rank 4"):
+            check_transferable("w", torch.bfloat16, (2, 2, 2, 2))
+
+
+class TestWireTypes:
+    def _init_info(self, **overrides):
+        fields = dict(
+            master_address="127.0.0.1",
+            master_port=1234,
+            rank_offset=1,
+            world_size=3,
+            src_mesh_dims=[1, 1],
+            dst_mesh_dims=[2, 1],
+            names=["w"],
+            dtype_names=["bfloat16"],
+            shapes=[[16, 16]],
+            src_placements=[None],
+        )
+        fields.update(overrides)
+        return M2NWeightTransferInitInfo(**fields)
+
+    def test_accepts_a_consistent_plan(self):
+        assert self._init_info().names == ["w"]
+
+    def test_accepts_pre_shared_nccl_unique_id(self):
+        info = self._init_info(
+            master_address=None,
+            master_port=None,
+            nccl_unique_id_b64=VALID_UID_B64,
+        )
+        assert info.nccl_unique_id_bytes == b"\x00" * 128
+        assert VALID_UID_B64 not in repr(info)
+
+    def test_destination_plan_uses_uid_communicator_without_group(self):
+        comm = Mock(rank=0, device=torch.device("cpu"), group=None)
+
+        def receive_plan(tensor, src):
+            assert src == 1
+            tensor.copy_(torch.tensor([[-1, 0], [-2, -2]], dtype=torch.int8))
+
+        comm.broadcast.side_effect = receive_plan
+
+        placements = publish_destination_placements(comm, 1, None, 2)
+
+        assert placements == [(REPLICATE, 0), REPLICATED]
+
+    def test_rejects_both_rendezvous_modes(self):
+        with pytest.raises(ValueError, match="not both"):
+            self._init_info(nccl_unique_id_b64=VALID_UID_B64)
+
+    def test_trainer_accepts_pre_shared_nccl_unique_id(self):
+        info = M2NTrainerInitInfo(
+            nccl_unique_id_b64=VALID_UID_B64,
+            world_size=3,
+            num_trainer_ranks=1,
+            rank=0,
+        )
+        assert info.nccl_unique_id_bytes == b"\x00" * 128
+        assert VALID_UID_B64 not in repr(info)
+
+    def test_ragged_plan_rejected(self):
+        with pytest.raises(ValueError, match="`shapes`"):
+            self._init_info(shapes=[])
+
+    def test_destination_mesh_must_cover_the_workers(self):
+        """The trainer declares the inference mesh, so one that does not cover
+        the workers is a config error — and it has to fail the init RPC, since
+        a mismatched mesh would otherwise surface as a hung collective."""
+        with pytest.raises(ValueError, match="dst_mesh_dims"):
+            self._init_info(dst_mesh_dims=[3, 1])  # 3 != the 2 workers
+
+    def test_world_must_hold_a_trainer_and_a_worker(self):
+        with pytest.raises(ValueError, match="rank_offset"):
+            self._init_info(rank_offset=3, world_size=3)
+
+    @pytest.mark.parametrize("dtype_name", ["not_a_dtype", "Tensor"])
+    def test_invalid_dtype_name_names_parameter(self, monkeypatch, dtype_name):
+        monkeypatch.setattr(
+            "vllm.distributed.weight_transfer.m2n_engine.import_m2n",
+            lambda: object(),
+        )
+        engine = object.__new__(M2NWeightTransferEngine)
+
+        with pytest.raises(ValueError, match=r"parameter 'w'.*dtype"):
+            engine.init_transfer_engine(self._init_info(dtype_names=[dtype_name]))
+
+    def _init_32_to_4_plan(self, monkeypatch, destination):
+        m2n = Mock()
+        m2n.Handle.create.return_value = object()
+        monkeypatch.setattr(
+            "vllm.distributed.weight_transfer.m2n_engine.import_m2n", lambda: m2n
+        )
+        monkeypatch.setattr(
+            "vllm.distributed.weight_transfer.m2n_engine."
+            "resolve_parameter_destinations",
+            lambda *args, **kwargs: [destination],
+        )
+        monkeypatch.setattr(
+            "vllm.distributed.weight_transfer.m2n_engine.worker_init_process_group",
+            lambda *args, **kwargs: object(),
+        )
+        monkeypatch.setattr(
+            "vllm.distributed.weight_transfer.m2n_engine."
+            "publish_destination_placements",
+            lambda *args: args[2],
+        )
+
+        engine = object.__new__(M2NWeightTransferEngine)
+        engine.parallel_config = Mock(pipeline_parallel_size=1)
+        engine.model_config = Mock(quantization=None)
+        engine.model = torch.nn.Module()
+        engine.init_transfer_engine(
+            self._init_info(
+                rank_offset=32,
+                world_size=36,
+                src_mesh_dims=[1, 32],
+                dst_mesh_dims=[1, 4],
+                shapes=[[32, 8]],
+                src_placements=[[REPLICATE, 0]],
+            )
+        )
+
+    def test_limits_use_resolved_sharded_destination(self, monkeypatch):
+        destination = M2NDestination("w", (REPLICATE, 0), torch.empty(8, 8))
+
+        self._init_32_to_4_plan(monkeypatch, destination)
+
+    def test_limits_still_reject_replicated_destination(self, monkeypatch):
+        destination = M2NDestination("w", REPLICATED, None)
+
+        with pytest.raises(ValueError, match=r"32 source shards.*MAX_SOURCES=16"):
+            self._init_32_to_4_plan(monkeypatch, destination)
+
+    def test_update_preflights_all_names_before_reshard(self):
+        engine = object.__new__(M2NWeightTransferEngine)
+        engine._handle = object()
+        engine.model_update_group = object()
+        engine._index = {"valid": 0}
+        engine._metas = [M2NParamMeta("valid", torch.float32, (4,), REPLICATED)]
+        engine._reshard = Mock()
+
+        with pytest.raises(ValueError, match=r"parameter 'unknown'"):
+            engine.receive_weights(
+                M2NWeightTransferUpdateInfo(names=["valid", "unknown"])
+            )
+
+        engine._reshard.assert_not_called()
+
+    def test_fallback_tensors_share_one_load_weights_invocation(self, monkeypatch):
+        class CoupledModel:
+            def __init__(self):
+                self.calls = 0
+                self.loaded = []
+
+            def load_weights(self, weights):
+                self.calls += 1
+                received = list(weights)
+                if [name for name, _ in received] == ["pair.weight", "pair.scale"]:
+                    self.loaded = [(name, tensor.clone()) for name, tensor in received]
+
+        model = CoupledModel()
+        direct = torch.zeros(1)
+        engine = object.__new__(M2NWeightTransferEngine)
+        engine._handle = object()
+        engine.model_update_group = object()
+        engine.device = torch.device("cpu")
+        engine.model = model
+        engine._index = {"pair.weight": 0, "direct": 1, "pair.scale": 2}
+        engine._metas = [
+            M2NParamMeta(name, torch.float32, (1,), REPLICATED)
+            for name in engine._index
+        ]
+        engine._parameter_destinations = [
+            M2NDestination("pair.weight", REPLICATED, None),
+            M2NDestination("direct", REPLICATED, direct),
+            M2NDestination("pair.scale", REPLICATED, None),
+        ]
+        reshard_order = []
+
+        def reshard(comm, stream, meta, placements, buffer):
+            reshard_order.append(meta.name)
+            buffer.fill_(len(reshard_order))
+
+        engine._reshard = reshard
+        stream = Mock()
+        monkeypatch.setattr(
+            "vllm.distributed.weight_transfer.m2n_engine.comm_ptr", lambda group: 0
+        )
+        monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream)
+
+        engine.receive_weights(
+            M2NWeightTransferUpdateInfo(names=["pair.weight", "direct", "pair.scale"])
+        )
+
+        assert model.calls == 1
+        assert [name for name, _ in model.loaded] == ["pair.weight", "pair.scale"]
+        assert [tensor.item() for _, tensor in model.loaded] == [1, 3]
+        assert direct.item() == 2
+        assert reshard_order == ["pair.weight", "direct", "pair.scale"]
+        assert stream.synchronize.call_count == 2
+
+    def test_trainer_rank_must_be_a_trainer_rank(self):
+        """A trainer rank must fall within the trainer portion of the group."""
+        with pytest.raises(ValueError, match="num_trainer_ranks"):
+            M2NTrainerInitInfo(
+                master_address="127.0.0.1",
+                master_port=1234,
+                world_size=4,
+                num_trainer_ranks=2,
+                rank=2,
+            )
+
+    def test_trainer_destination_mesh_must_cover_the_workers(self):
+        """The destination mesh must include every inference worker."""
+        with pytest.raises(ValueError, match="dst_mesh_dims"):
+            M2NTrainerInitInfo(
+                master_address="127.0.0.1",
+                master_port=1234,
+                world_size=6,
+                num_trainer_ranks=2,
+                dst_mesh_dims=(3, 1),  # 3 != the 4 inference workers
+                rank=0,
+            )
+
+    def test_destination_mesh_defaults_to_flat(self):
+        """A replicated destination does not care how the mesh is factored, so
+        callers that do not shard it need not supply one."""
+        info = M2NTrainerInitInfo(
+            master_address="127.0.0.1",
+            master_port=1234,
+            world_size=6,
+            num_trainer_ranks=2,
+            rank=0,
+        )
+        assert info.destination_mesh_dims == (4, 1)
+
+    def test_sender_is_trainer_rank_zero(self):
+        """Trainer rank 0 drives the inference control plane."""
+        info = M2NTrainerInitInfo(
+            master_address="127.0.0.1", master_port=1234, world_size=4, rank=0
+        )
+        assert info.is_sender
+
+
+class TestRegistration:
+    def test_both_registries_expose_the_backend(self):
+        """Both worker and trainer factories register nccl_m2n."""
+        assert "nccl_m2n" in WeightTransferEngineFactory._registry
+        assert "nccl_m2n" in WeightTransferTrainerFactory._registry
+
+    def test_init_info_dispatches_to_the_backend(self):
+        """Trainer init info selects the nccl_m2n factory entry."""
+        assert M2NTrainerInitInfo.backend == "nccl_m2n"
+
+
+# ---------------------------------------------------------------------------
+# End-to-end transfer
+#
+# Transport-only, in the style of the NCCL/sparse tests in
+# test_weight_transfer.py: two Ray tasks with one GPU each drive the two engines
+# directly, so no HTTP server or LLM instance is involved. The worker is not an
+# RPC endpoint here, so the trainer engine gets a no-op control-plane client --
+# the NCCL rendezvous and the reshard itself are the real thing.
+# ---------------------------------------------------------------------------
+
+SHAPE = [64, 32]
+DTYPE = "float32"
+
+
+def _init_ray() -> None:
+    if ray.is_initialized():
+        return
+    ray.init(
+        ignore_reinit_error=True,
+        runtime_env={
+            "env_vars": {
+                "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+                "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1",
+                "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1",
+            }
+        },
+    )
+
+
+def _assigned_device() -> "torch.device":
+    gpu_ids = ray.get_gpu_ids()
+    device = torch.device(f"cuda:{int(gpu_ids[0])}" if gpu_ids else "cuda:0")
+    current_platform.set_device(device)
+    return device
+
+
+@ray.remote(num_gpus=1)
+def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool:
+    """Send one parameter through the real trainer engine."""
+    device = _assigned_device()
+
+    from vllm.distributed.weight_transfer import WeightTransferTrainerFactory
+    from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource
+    from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
+
+    class NoopClient:
+        def init_weight_transfer_engine(self, init_info):
+            pass
+
+        def start_weight_update(self):
+            pass
+
+        def update_weights(self, update_info):
+            pass
+
+        def finish_weight_update(self, weight_version=None):
+            pass
+
+    class Tiny(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(
+                torch.arange(
+                    SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device
+                ).reshape(SHAPE)
+            )
+
+    engine = WeightTransferTrainerFactory.trainer_init(
+        init_info=M2NTrainerInitInfo(
+            nccl_unique_id_b64=nccl_unique_id_b64,
+            world_size=world_size,
+            num_trainer_ranks=1,
+            dst_mesh_dims=(1, world_size - 1),
+            rank=0,
+        ),
+        client=NoopClient(),
+        source=DTensorModuleSource(Tiny(), num_trainer_ranks=1),
+    )
+    engine.send_weights()
+    torch.accelerator.synchronize()
+    engine.shutdown()
+    return True
+
+
+@ray.remote(num_gpus=1)
+def _m2n_worker_receive(
+    nccl_unique_id_b64: str,
+    world_size: int,
+    worker_rank: int = 0,
+) -> dict:
+    """Receive that parameter through the real worker engine."""
+    import contextlib
+    from unittest.mock import MagicMock
+
+    device = _assigned_device()
+
+    from vllm.config.parallel import ParallelConfig
+    from vllm.config.weight_transfer import WeightTransferConfig
+    from vllm.distributed.weight_transfer.m2n_engine import (
+        M2NWeightTransferEngine,
+        M2NWeightTransferInitInfo,
+        M2NWeightTransferUpdateInfo,
+    )
+
+    class Recorder(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.received = []
+            if world_size > 2:
+                self.weight = torch.nn.Parameter(
+                    torch.zeros(
+                        SHAPE[0] // (world_size - 1),
+                        SHAPE[1],
+                        dtype=torch.float32,
+                        device=device,
+                    )
+                )
+                self.weight.output_dim = 0
+                self.weight.weight_loader = MethodType(
+                    ColumnParallelLinear.weight_loader,
+                    Mock(tp_size=world_size - 1),
+                )
+
+        def load_weights(self, weights):
+            for name, tensor in weights:
+                self.received.append((name, tensor.clone()))
+
+    parallel_config = MagicMock(spec=ParallelConfig)
+    parallel_config.rank = worker_rank
+    parallel_config.world_size = world_size - 1
+    parallel_config.data_parallel_rank = 0
+    parallel_config.data_parallel_index = 0
+    parallel_config.tensor_parallel_size = world_size - 1
+    parallel_config.pipeline_parallel_size = 1
+    vllm_config = MagicMock()
+    vllm_config.parallel_config = parallel_config
+    vllm_config.model_config = MagicMock()
+    vllm_config.model_config.quantization = None
+
+    recorder = Recorder()
+    engine = M2NWeightTransferEngine(
+        WeightTransferConfig(backend="nccl_m2n"), vllm_config, device, recorder
+    )
+    # Transport-only: receive_weights enters set_current_vllm_config, and
+    # vllm_config here is a mock.
+    import vllm.config as _vllm_config_mod
+
+    _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
+
+    engine.init_transfer_engine(
+        M2NWeightTransferInitInfo(
+            nccl_unique_id_b64=nccl_unique_id_b64,
+            rank_offset=1,  # trainer occupies rank 0
+            world_size=world_size,
+            src_mesh_dims=[1, 1],
+            dst_mesh_dims=[1, world_size - 1],
+            names=["weight"],
+            dtype_names=[DTYPE],
+            shapes=[SHAPE],
+            src_placements=[None],  # replicated on the single trainer rank
+        )
+    )
+    engine.receive_weights(M2NWeightTransferUpdateInfo(names=["weight"]))
+    torch.accelerator.synchronize()
+
+    full = torch.arange(
+        SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device
+    ).reshape(SHAPE)
+    direct = world_size > 2
+    name: str | None
+    got: torch.Tensor | None
+    if direct:
+        name, got = "weight", recorder.weight
+        expected = full.chunk(world_size - 1, dim=0)[worker_rank]
+    else:
+        name, got = recorder.received[0] if recorder.received else (None, None)
+        expected = full
+    result = {
+        "count": len(recorder.received),
+        "name": name,
+        "shape": list(got.shape) if got is not None else None,
+        "exact": bool(torch.equal(got, expected)) if got is not None else False,
+        "direct": direct,
+    }
+    engine.shutdown()
+    return result
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 2,
+    reason="Need at least 2 GPUs: one trainer rank and one inference worker.",
+)
+def test_m2n_weight_transfer_between_processes():
+    """A parameter survives a real reshard from a trainer process to a worker.
+
+    This is the only test here that moves bytes: it builds both engines, joins
+    one NCCL communicator across two processes, and reshards. Everything else
+    in this file checks how a transfer is *described*.
+    """
+    pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime")
+    _init_ray()
+
+    from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
+
+    nccl = NCCLLibrary()
+    nccl_unique_id_b64 = base64.b64encode(
+        bytes(nccl.ncclGetUniqueId().internal)
+    ).decode()
+    world_size = 2  # 1 trainer + 1 inference worker
+
+    worker = _m2n_worker_receive.remote(nccl_unique_id_b64, world_size)
+    trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size)
+    trainer_ok, result = ray.get([trainer, worker])
+
+    assert trainer_ok, "trainer engine did not complete"
+    assert result["count"] == 1, f"expected one parameter, got {result['count']}"
+    assert result["name"] == "weight"
+    assert result["shape"] == SHAPE
+    assert result["exact"], "resharded tensor does not match what the trainer sent"
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 3,
+    reason="Need at least 3 GPUs: one trainer rank and two inference workers.",
+)
+def test_m2n_weight_transfer_to_tp2_shards():
+    """Two inference workers receive their own TP shard in live model storage."""
+    pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime")
+    _init_ray()
+
+    from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
+
+    nccl = NCCLLibrary()
+    nccl_unique_id_b64 = base64.b64encode(
+        bytes(nccl.ncclGetUniqueId().internal)
+    ).decode()
+    world_size = 3
+
+    workers = [
+        _m2n_worker_receive.remote(nccl_unique_id_b64, world_size, rank)
+        for rank in range(2)
+    ]
+    trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size)
+    trainer_ok, *results = ray.get([trainer, *workers])
+
+    assert trainer_ok, "trainer engine did not complete"
+    for rank, result in enumerate(results):
+        assert result["count"] == 0, f"worker {rank} unexpectedly called load_weights"
+        assert result["name"] == "weight"
+        assert result["shape"] == [SHAPE[0] // 2, SHAPE[1]]
+        assert result["direct"]
+        assert result["exact"], f"worker {rank} received the wrong shard"
+
+
+class _Model(torch.nn.Module):
+    """Stand-in for a TP=2 shard of a tiny model.
+
+    `column` is sharded on dim 0 and `row` on dim 1 (the two shapes a TP linear
+    layer produces), `norm` is replicated, and `fused` has no checkpoint-name
+    counterpart — the shape a fused parameter presents to the resolver.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.column = torch.nn.Parameter(torch.zeros(8, 16))
+        self.column.output_dim = 0
+        self.column.weight_loader = MethodType(
+            ColumnParallelLinear.weight_loader, Mock(tp_size=2)
+        )
+        self.row = torch.nn.Parameter(torch.zeros(16, 8))
+        self.row.input_dim = 1
+        self.row.weight_loader = MethodType(
+            RowParallelLinear.weight_loader, Mock(tp_size=2)
+        )
+        self.norm = torch.nn.Parameter(torch.zeros(16))
+        self.norm.weight_loader = default_weight_loader
+        self.fused = torch.nn.Parameter(torch.zeros(24, 16))
+
+
+class _MixedModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.proj = torch.nn.Module()
+        self.proj.weight = torch.nn.Parameter(torch.zeros(16, 8))
+        self.proj.weight.input_dim = 1
+        self.proj.weight.weight_loader = MethodType(
+            RowParallelLinear.weight_loader, Mock(tp_size=2)
+        )
+        self.proj.bias = torch.nn.Parameter(torch.zeros(16))
+
+
+def _resolve(names, dtypes, shapes, **kwargs):
+    defaults = dict(num_workers=2, shard_axis_size=2, allow_direct=True)
+    defaults.update(kwargs)
+    return resolve_parameter_destinations(_Model(), names, dtypes, shapes, **defaults)
+
+
+class TestDestinationResolution:
+    @pytest.mark.parametrize(
+        "name,expected_dim",
+        [("column", 0), ("row", 1)],
+    )
+    def test_sharded_parameter_resolves_to_its_tp_dim(self, name, expected_dim):
+        [destination] = _resolve([name], [torch.float32], [(16, 16)])
+        assert destination.direct
+        assert destination.placements == (REPLICATE, expected_dim)
+
+    def test_replicated_parameter_needs_no_placement(self):
+        """A parameter every rank holds in full is REPLICATED, so it works
+        whatever way the inference mesh happens to be factored."""
+        [destination] = _resolve(["norm"], [torch.float32], [(16,)])
+        assert destination.direct
+        assert destination.placements is REPLICATED
+
+    def test_direct_destination_is_the_live_parameter(self):
+        model = _Model()
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["column"],
+            [torch.float32],
+            [(16, 16)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+        assert destination.tensor.data_ptr() == model.column.data_ptr()
+
+    def test_known_transforming_model_submodule_falls_back(self):
+        class GPT2Model(_Model):
+            pass
+
+        GPT2Model.__module__ = "vllm.model_executor.models.gpt2"
+        model = torch.nn.Module()
+        model.transforming = GPT2Model()
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["transforming.row"],
+            [torch.float32],
+            [(16, 16)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+
+        assert not destination.direct
+        assert destination.placements is REPLICATED
+
+    def test_unknown_name_falls_back(self):
+        """Fused parameters reach the worker under checkpoint names that do not
+        exist in the model; those must take the full-tensor path."""
+        [destination] = _resolve(["mlp.gate_proj.weight"], [torch.float32], [(16, 16)])
+        assert not destination.direct
+        assert destination.placements is REPLICATED
+
+    def test_fallback_parameter_demotes_direct_sibling_in_same_module(self):
+        """Layerwise finalization must not overwrite a directly loaded weight
+        when a sibling bias takes the fallback path."""
+        destinations = resolve_parameter_destinations(
+            _MixedModel(),
+            ["proj.weight", "proj.bias"],
+            [torch.float32, torch.float32],
+            [(16, 16), (16,)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+
+        assert not any(destination.direct for destination in destinations)
+
+    def test_unknown_loader_falls_back_despite_matching_name_and_shape(self):
+        model = _Model()
+        model.norm.weight_loader = lambda param, weight: param.data.copy_(weight)
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["norm"],
+            [torch.float32],
+            [(16,)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+        assert not destination.direct
+
+    def test_missing_loader_is_not_assumed_to_be_a_plain_copy(self):
+        model = _Model()
+        del model.norm.weight_loader
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["norm"],
+            [torch.float32],
+            [(16,)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+        assert not destination.direct
+
+    def test_wrapped_known_loader_falls_back(self):
+        model = _Model()
+
+        @wraps(default_weight_loader)
+        def wrapped_loader(param, weight):
+            default_weight_loader(param, weight)
+
+        model.norm.weight_loader = wrapped_loader
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["norm"],
+            [torch.float32],
+            [(16,)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+        assert not destination.direct
+
+    def test_replicated_loader_with_transform_metadata_falls_back(self):
+        model = _Model()
+        model.norm.is_transposed = True
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["norm"],
+            [torch.float32],
+            [(16,)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+        assert not destination.direct
+
+    def test_packed_dim_zero_falls_back(self):
+        model = _Model()
+        model.column.packed_dim = 0
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["column"],
+            [torch.float32],
+            [(16, 16)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+        assert not destination.direct
+
+    def test_loader_declared_dim_wins_over_shape_guessing(self):
+        model = _Model()
+        model.column.output_dim = 1
+        [destination] = resolve_parameter_destinations(
+            model,
+            ["column"],
+            [torch.float32],
+            [(16, 16)],
+            num_workers=2,
+            shard_axis_size=2,
+            allow_direct=True,
+        )
+        assert not destination.direct
+
+    def test_shape_the_tp_factor_cannot_explain_falls_back(self):
+        [destination] = _resolve(["fused"], [torch.float32], [(16, 16)])
+        assert not destination.direct
+
+    def test_dtype_mismatch_falls_back(self):
+        """A parameter stored in a different dtype than the wire dtype means a
+        quantized or otherwise transformed layout — never write into it."""
+        [destination] = _resolve(["column"], [torch.bfloat16], [(16, 16)])
+        assert not destination.direct
+
+    def test_allow_direct_false_forces_every_parameter_to_fall_back(self):
+        destinations = _resolve(
+            ["column", "norm"],
+            [torch.float32] * 2,
+            [(16, 16), (16,)],
+            allow_direct=False,
+        )
+        assert not any(d.direct for d in destinations)
+
+    def test_plan_reports_byte_coverage(self, caplog_vllm):
+        """A small direct parameter must not hide a large fallback tensor."""
+        with caplog_vllm.at_level(
+            logging.INFO,
+            logger="vllm.distributed.weight_transfer.m2n_layout",
+        ):
+            _resolve(
+                ["norm", "missing_fused_weight"],
+                [torch.float32, torch.bfloat16],
+                [(16,), (1024, 1024)],
+            )
+
+        assert (
+            "1/2 parameters resharded directly into the model, 1 via "
+            "full-tensor fallback; direct byte coverage: 64/2097216 bytes "
+            "(0.0%)" in caplog_vllm.text
+        )

vllm/config/weight_transfer.py

@@ -9,7 +9,9 @@ from vllm.config.utils import config
 class WeightTransferConfig:
     """Configuration for weight transfer during RL training."""
 
-    backend: Literal["nccl", "ipc", "sparse_nccl", "sharded_rdt"] | str = "nccl"
+    backend: Literal["nccl", "ipc", "sparse_nccl", "sharded_rdt", "nccl_m2n"] | str = (
+        "nccl"
+    )
     """The backend to use for weight transfer. Validated against the
     `WeightTransferEngineFactory` registry at engine creation time.
     """

vllm/distributed/weight_transfer/factory.py

@@ -243,6 +243,12 @@ WeightTransferEngineFactory.register_engine(
     "ShardedRDTWeightTransferEngine",
 )
 
+WeightTransferEngineFactory.register_engine(
+    "nccl_m2n",
+    "vllm.distributed.weight_transfer.m2n_engine",
+    "M2NWeightTransferEngine",
+)
+
 
 # Trainer-side engines, parallel to the worker registry above.
 WeightTransferTrainerFactory.register_engine(
@@ -268,3 +274,9 @@ WeightTransferTrainerFactory.register_engine(
     "vllm.distributed.weight_transfer.sharded_rdt_trainer",
     "ShardedRDTTrainerWeightTransferEngine",
 )
+
+WeightTransferTrainerFactory.register_engine(
+    "nccl_m2n",
+    "vllm.distributed.weight_transfer.m2n_trainer",
+    "M2NTrainerWeightTransferEngine",
+)

vllm/distributed/weight_transfer/m2n_common.py

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared helpers for the NCCL M2N (`nccl_m2n`) weight transfer backend.
+
+M2N reshards a tensor between two disjoint meshes of ranks that live in one
+communicator: the trainer occupies ranks `[0, T)` and the inference workers
+`[T, T + N)`. This module holds everything both sides need — the optional
+runtime import, layout descriptors, and the conversion to `nccl.m2n` types —
+so the engine module stays about the transfer itself.
+"""
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+
+from vllm.distributed.weight_transfer.base import ParamMeta
+
+if TYPE_CHECKING:
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+
+# Placement code for a mesh axis that replicates. Any other (non-negative) code
+# is the tensor dim that axis shards. Ints rather than `nccl.m2n` objects so
+# layouts survive the JSON init handshake without a custom encoder.
+REPLICATE = -1
+
+# `NCCL_RESHARD_MAX_TENSOR_DIMS`; 4-D and higher are rejected by the library.
+MAX_TENSOR_DIMS = 3
+
+MESH_NDIMS = 2
+
+# Destination mesh convention. Keeping these roles named avoids scattering
+# positional axis assumptions through the worker-side planner.
+DESTINATION_REPLICA_AXIS = 0
+DESTINATION_SHARD_AXIS = 1
+
+# m2n bounds how many source shards may feed one destination shard, and how many
+# destination shards one source shard may feed, with compile-time arrays in
+# `reshard_limits.h`. The bindings do not expose them, so they are mirrored
+# here; a build with larger arrays makes these conservative, never wrong.
+MAX_SOURCE_SHARDS = 16
+MAX_DEST_SHARDS = 64
+
+# Wire dtypes `ncclReshard` accepts. Notably excludes fp4 and any packed
+# sub-byte type, so quantized checkpoints are out of scope for now.
+SUPPORTED_DTYPES: frozenset[torch.dtype] = frozenset(
+    dtype
+    for dtype in (
+        getattr(torch, name, None)
+        for name in (
+            "int8",
+            "uint8",
+            "float8_e4m3fn",
+            "float8_e5m2",
+            "float16",
+            "bfloat16",
+            "int32",
+            "uint32",
+            "float32",
+            "int64",
+            "uint64",
+            "float64",
+        )
+    )
+    if dtype is not None
+)
+
+_IMPORT_HINT = (
+    "The nccl_m2n weight transfer backend requires the `nccl-extensions` "
+    "package (and its `nccl4py` dependency), which is not installed by vLLM. "
+    "See https://github.com/NVIDIA/nccl-extensions for build "
+    "and install instructions. It needs NCCL 2.30.5 or newer, and "
+    "VLLM_NCCL_SO_PATH must point at the same libnccl.so that libnccl_m2n.so "
+    "was linked against."
+)
+
+
+def import_m2n() -> Any:
+    """Import `nccl.m2n` lazily, with an actionable error when it is missing.
+
+    Deferred so that importing vLLM — or any other weight transfer backend —
+    never requires the m2n runtime to be present.
+    """
+    try:
+        import nccl.m2n as m2n
+    except ImportError as e:
+        raise ImportError(f"{_IMPORT_HINT} (import failed: {e})") from e
+    return m2n
+
+
+@dataclass(frozen=True)
+class M2NMesh:
+    """One side's rank topology — pure topology, no tensor placement.
+
+    Mirrors `ncclMesh_t`: a 2-axis mesh owning the contiguous rank interval
+    `[start_rank, start_rank + dims[0] * dims[1])`. There is no 1-D mesh; a
+    single-axis topology is spelled with a second axis of size 1.
+
+    One mesh describes every tensor on its side, so it is exchanged once at the
+    init handshake rather than per parameter.
+    """
+
+    dims: tuple[int, int]
+    start_rank: int
+
+    def __post_init__(self) -> None:
+        if len(self.dims) != MESH_NDIMS or any(d <= 0 for d in self.dims):
+            raise ValueError(f"mesh dims must be {MESH_NDIMS} positive ints")
+        if self.start_rank < 0:
+            raise ValueError(
+                f"mesh start_rank must be non-negative, got {self.start_rank}"
+            )
+
+    @property
+    def size(self) -> int:
+        return self.dims[0] * self.dims[1]
+
+
+# A tensor's placement over its side's mesh: one code per mesh axis
+# (`REPLICATE`, or the tensor dim that axis shards), or `REPLICATED` for a
+# tensor every rank holds in full.
+Placements = tuple[int, int]
+REPLICATED: Placements | None = None
+
+
+def check_placements(placements: Placements, context: str = "placements") -> None:
+    """Reject placement pairs m2n cannot express.
+
+    The header requires exactly one SHARD axis and one REPLICATE axis.
+    `{REPLICATE, REPLICATE}` hits a degenerate prepare branch, which is why
+    full replication is carried as `REPLICATED` and resolved separately;
+    sharding both axes is not expressible at all.
+    """
+    if len(placements) != MESH_NDIMS:
+        raise ValueError(f"{context} must have {MESH_NDIMS} entries")
+    invalid = [code for code in placements if code < REPLICATE]
+    if invalid:
+        raise ValueError(
+            f"{context} contains invalid placement code {invalid[0]}; valid "
+            f"codes are {REPLICATE} (Replicate) and non-negative tensor dimensions"
+        )
+    num_sharded = sum(code != REPLICATE for code in placements)
+    if num_sharded == 0:
+        raise ValueError(
+            f"{context}: a fully replicated tensor is carried as REPLICATED, "
+            "not as two "
+            "REPLICATE axes"
+        )
+    if num_sharded == MESH_NDIMS:
+        raise ValueError(
+            f"{context}: nccl_m2n needs one REPLICATE mesh axis, but "
+            f"{placements} shards both"
+        )
+
+
+def resolve_layout(
+    mesh: M2NMesh,
+    placements: Placements | None,
+    context: str = "placements",
+) -> tuple[M2NMesh, Placements]:
+    """Pair one tensor's placement with the mesh m2n should see for it.
+
+    A replicated tensor needs a size-1 mesh axis to carry a no-op shard, since
+    m2n has no `{REPLICATE, REPLICATE}`. It is therefore described over the
+    *same rank interval* re-factored as `(size, 1)`. That re-factoring is sound
+    precisely because replication is order-independent — every rank holds the
+    whole tensor, so it does not matter that `(a, b)` and `(size, 1)` walk the
+    interval in a different order. A sharded tensor keeps its side's own
+    factorization, where rank order decides who owns which shard.
+    """
+    if placements is None:
+        return M2NMesh((mesh.size, 1), mesh.start_rank), (REPLICATE, 0)
+    check_placements(placements, context)
+    return mesh, placements
+
+
+def shard_count(mesh: M2NMesh, placements: Placements) -> int:
+    """How many pieces this layout splits the tensor into."""
+    return next(
+        (mesh.dims[axis] for axis, code in enumerate(placements) if code != REPLICATE),
+        1,
+    )
+
+
+def check_plan_limits(
+    src: tuple[M2NMesh, Placements],
+    dst: tuple[M2NMesh, Placements],
+    name: str,
+) -> None:
+    """Reject plans that exceed m2n's static per-shard fan-in / fan-out arrays.
+
+    Only the unambiguous cases are checked: when one side is a single shard it
+    is fed by (or feeds) every shard on the other side, so the count is exactly
+    the other side's shard count. In the general sharded-to-sharded case the
+    overlap depends on the library's chunking, and reproducing that arithmetic
+    here would duplicate internals that can change under us.
+
+    m2n re-checks authoritatively and, because the plan is derived from the
+    shared descriptors, fails identically on every rank -- so this is about
+    reporting at init with a message that names the parameter, not about
+    avoiding a hang.
+    """
+    src_shards = shard_count(*src)
+    dst_shards = shard_count(*dst)
+    if dst_shards == 1 and src_shards > MAX_SOURCE_SHARDS:
+        raise ValueError(
+            f"parameter '{name}' would feed {src_shards} source shards into one "
+            f"replicated destination, over m2n's MAX_SOURCES={MAX_SOURCE_SHARDS}. "
+            "Shard the destination, replicate the source, or rebuild m2n with "
+            "larger arrays."
+        )
+    if src_shards == 1 and dst_shards > MAX_DEST_SHARDS:
+        raise ValueError(
+            f"parameter '{name}' would fan one source shard out to {dst_shards} "
+            f"destination shards, over m2n's MAX_TARGETS={MAX_DEST_SHARDS}. "
+            "Rebuild m2n with larger arrays to raise the bound."
+        )
+
+
+def validate_layout(
+    mesh: M2NMesh, placements: Placements, shape: Sequence[int], side: str
+) -> None:
+    """Check a resolved layout can describe `shape` before any collective runs.
+
+    Called on both sides at init so a bad plan surfaces as an error from the
+    init RPC rather than as a hang inside the first reshard.
+    """
+    for axis, code in enumerate(placements):
+        if code == REPLICATE:
+            continue
+        if code >= len(shape):
+            raise ValueError(
+                f"{side} layout shards tensor dim {code}, but the tensor "
+                f"has rank {len(shape)}"
+            )
+        factor = mesh.dims[axis]
+        if shape[code] % factor:
+            raise ValueError(
+                f"{side} layout shards dim {code} (size {shape[code]}) over "
+                f"{factor} ranks, which does not divide evenly"
+            )
+
+
+@dataclass(frozen=True)
+class M2NParamMeta(ParamMeta):
+    """`ParamMeta` extended with how the trainer places this tensor.
+
+    The base class carries only name / dtype / full shape, which is not enough
+    to plan a reshard. `placements` is relative to its side's `M2NMesh`, or
+    `REPLICATED` when every rank holds the whole tensor.
+    """
+
+    placements: Placements | None
+
+
+def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> None:
+    """Reject tensors m2n cannot move at all, with the parameter named."""
+    if dtype not in SUPPORTED_DTYPES:
+        raise ValueError(
+            f"parameter '{name}' has dtype {dtype}, which nccl_m2n does not "
+            f"support. Supported: {sorted(str(d) for d in SUPPORTED_DTYPES)}"
+        )
+    if not 1 <= len(shape) <= MAX_TENSOR_DIMS:
+        raise ValueError(
+            f"parameter '{name}' has rank {len(shape)}; nccl_m2n supports "
+            f"rank 1..{MAX_TENSOR_DIMS}"
+        )
+
+
+def to_mesh(m2n: Any, mesh: M2NMesh) -> Any:
+    return m2n.Mesh(mesh.dims, start_rank=mesh.start_rank)
+
+
+def to_placements(m2n: Any, placements: Placements) -> list[Any]:
+    return [
+        m2n.Replicate() if code == REPLICATE else m2n.Shard(code) for code in placements
+    ]
+
+
+def publish_destination_placements(
+    comm: "PyNcclCommunicator",
+    first_worker_rank: int,
+    placements: "Sequence[Placements | None] | None",
+    num_parameters: int,
+) -> list[Placements | None]:
+    """Share the worker-side destination plan with every rank in the group.
+
+    The trainer must issue each reshard with the same destination the workers
+    use, but once destinations are per-parameter they depend on the inference
+    model, which only the workers can see. The first worker publishes them here,
+    over the shared NCCL communicator; trainer ranks pass `None` and receive
+    them, and the other workers pass their own so a disagreement is caught
+    rather than deadlocking later.
+
+    Only placements travel: the destination *mesh* is still derived from the
+    rank split, identically on both sides. Use the NCCL communicator directly
+    because unique-id rendezvous deliberately has no bootstrap process group.
+    """
+    replicated_sentinel = REPLICATE - 1
+    encoded = torch.full(
+        (num_parameters, MESH_NDIMS),
+        replicated_sentinel,
+        dtype=torch.int8,
+        device=comm.device,
+    )
+    if comm.rank == first_worker_rank:
+        if placements is None or len(placements) != num_parameters:
+            raise ValueError(
+                f"publishing rank needs {num_parameters} destination placements"
+            )
+        rows = [
+            [replicated_sentinel] * MESH_NDIMS if placement is REPLICATED else placement
+            for placement in placements
+        ]
+        if rows:
+            encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=comm.device))
+
+    comm.broadcast(encoded, src=first_worker_rank)
+    return [
+        REPLICATED
+        if all(code == replicated_sentinel for code in row)
+        else cast(Placements, tuple(row))
+        for row in encoded.tolist()
+    ]
+
+
+def comm_ptr(comm: "PyNcclCommunicator") -> int:
+    """Raw `ncclComm_t` behind vLLM's `PyNcclCommunicator`.
+
+    m2n links its own NCCL, so the handle only means anything if vLLM loaded
+    the same `libnccl.so` — set `VLLM_NCCL_SO_PATH` accordingly.
+    """
+    handle = comm.comm
+    ptr = getattr(handle, "value", handle)
+    if not ptr:
+        raise RuntimeError("PyNcclCommunicator has no live NCCL communicator")
+    return int(ptr)

vllm/distributed/weight_transfer/m2n_engine.py

@@ -0,0 +1,453 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-side weight transfer engine built on NCCL M2N (`nccl_m2n`).
+
+The trainer and the inference workers share one NCCL communicator: trainer ranks
+occupy `[0, T)`, workers `[T, T + N)`. Each parameter is moved with a single
+`nccl.m2n.reshard`, which redistributes it from the trainer's layout (FSDP / EP
+/ arbitrary DTensor sharding) to the inference layout — so the trainer sends its
+local shards and never all-gathers a full tensor, which is what the broadcast
+NCCL backend forces it to do.
+
+When an incoming checkpoint parameter maps directly to a live vLLM parameter,
+M2N reshards into that worker's local model storage. Parameters whose names or
+layouts cannot be resolved receive a full tensor and fall back to
+`load_weights`, preserving the existing backend's loading behavior.
+
+Both meshes participate in every reshard, in the same order, so this backend has
+the same concurrency shape as the broadcast NCCL backend: the worker must be
+inside `receive_weights` while the trainer is sending. Driving that from the
+trainer is the trainer engine's job.
+"""
+
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+
+from vllm.config.weight_transfer import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferEngine,
+    WeightTransferInitInfo,
+    WeightTransferUpdateInfo,
+)
+from vllm.distributed.weight_transfer.m2n_common import (
+    DESTINATION_SHARD_AXIS,
+    MESH_NDIMS,
+    REPLICATED,
+    M2NMesh,
+    M2NParamMeta,
+    Placements,
+    check_plan_limits,
+    check_transferable,
+    comm_ptr,
+    import_m2n,
+    publish_destination_placements,
+    resolve_layout,
+    to_mesh,
+    to_placements,
+    validate_layout,
+)
+from vllm.distributed.weight_transfer.m2n_layout import (
+    M2NDestination,
+    resolve_parameter_destinations,
+)
+from vllm.distributed.weight_transfer.nccl_common import (
+    NCCLWeightTransferInitInfo,
+    decode_nccl_unique_id,
+    worker_init_process_group,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+
+__all__ = [
+    "M2NWeightTransferInitInfo",
+    "M2NWeightTransferUpdateInfo",
+    "M2NWeightTransferEngine",
+]
+
+
+# ---------------------------------------------------------------------------
+# Wire types
+# ---------------------------------------------------------------------------
+
+
+@dataclass(kw_only=True)
+class M2NWeightTransferInitInfo(WeightTransferInitInfo):
+    """Worker-side init info: the rendezvous plus the full transfer plan.
+
+    Layouts are static for the whole run, so they ride the one-time init
+    handshake and the per-round update info stays a list of names. Everything is
+    plain JSON so the HTTP control plane carries it unchanged.
+    """
+
+    master_address: str | None = None
+    master_port: int | None = None
+    rank_offset: int
+    """First worker rank, i.e. the number of trainer ranks."""
+    world_size: int
+    """Trainer ranks + all inference workers."""
+    src_mesh_dims: list[int]
+    """The trainer's mesh, shared by every parameter (`start_rank` is 0: the
+    trainer occupies the front of the communicator)."""
+    dst_mesh_dims: list[int]
+    """The inference mesh, starting at `rank_offset`. Declared rather than
+    derived so both sides describe the destination identically."""
+    names: list[str]
+    dtype_names: list[str]
+    shapes: list[list[int]]
+    src_placements: list[list[int] | None]
+    """Per parameter, relative to `src_mesh_dims`; `None` means replicated."""
+    nccl_unique_id_b64: str | None = field(default=None, repr=False)
+    max_cta: int | None = None
+
+    def __post_init__(self) -> None:
+        _ = self.nccl_unique_id_bytes
+        num_params = len(self.names)
+        for label, values in (
+            ("dtype_names", self.dtype_names),
+            ("shapes", self.shapes),
+            ("src_placements", self.src_placements),
+        ):
+            if len(values) != num_params:
+                raise ValueError(
+                    f"`{label}` should be of the same size as `names`: "
+                    f"got {len(values)} and {num_params}"
+                )
+        if self.rank_offset < 1 or self.rank_offset >= self.world_size:
+            raise ValueError(
+                f"`rank_offset` ({self.rank_offset}) must leave at least one "
+                f"trainer rank and one worker in world_size {self.world_size}"
+            )
+        num_workers = self.world_size - self.rank_offset
+        dst = tuple(self.dst_mesh_dims)
+        if len(dst) != MESH_NDIMS or dst[0] * dst[1] != num_workers:
+            raise ValueError(
+                f"`dst_mesh_dims` {dst} must be {MESH_NDIMS} dims covering the "
+                f"{num_workers} inference workers"
+            )
+
+    @property
+    def nccl_unique_id_bytes(self) -> bytes | None:
+        return decode_nccl_unique_id(
+            master_address=self.master_address,
+            master_port=self.master_port,
+            nccl_unique_id_b64=self.nccl_unique_id_b64,
+            ctx="M2NWeightTransferInitInfo",
+        )
+
+
+@dataclass
+class M2NWeightTransferUpdateInfo(WeightTransferUpdateInfo):
+    """Per-round update info: which parameters this chunk carries, in order.
+
+    Shapes, dtypes and layouts were fixed at init, so a round only needs to say
+    what is coming and in what order — both sides must issue their reshards in
+    exactly that order.
+    """
+
+    names: list[str]
+
+
+# ---------------------------------------------------------------------------
+# Worker engine
+# ---------------------------------------------------------------------------
+
+
+class M2NWeightTransferEngine(
+    WeightTransferEngine[M2NWeightTransferInitInfo, M2NWeightTransferUpdateInfo]
+):
+    """Inference-side engine: receives each parameter with one reshard.
+
+    Resolvable parameters are resharded from the trainer layout directly into
+    each worker's live local model storage. Unresolvable parameters receive a
+    full tensor and use `load_weights`. In both cases, the trainer sends its
+    local shards without materializing a full tensor.
+    """
+
+    init_info_cls = M2NWeightTransferInitInfo
+    update_info_cls = M2NWeightTransferUpdateInfo
+    # The destination plan is resolved against the model this engine was built
+    # for, so it cannot be retargeted at a draft model mid-flight.
+    supports_draft_weight_update = False
+
+    def __init__(
+        self,
+        config: WeightTransferConfig,
+        vllm_config: "VllmConfig",
+        device: torch.device,
+        model: torch.nn.Module,
+    ) -> None:
+        super().__init__(config, vllm_config, device, model)
+        self.model_update_group: PyNcclCommunicator | None = None
+        self._m2n: Any = None
+        self._handle: Any = None
+        self._metas: list[M2NParamMeta] = []
+        self._src_mesh: M2NMesh | None = None
+        self._dst_mesh: M2NMesh | None = None
+        self._parameter_destinations: list[M2NDestination] = []
+        self._index: dict[str, int] = {}
+        self._uses_load_weights = False
+
+    def init_transfer_engine(self, init_info: M2NWeightTransferInitInfo) -> None:
+        """Join the trainer's communicator and rebuild the transfer plan.
+
+        Every precondition (dtype, tensor rank, divisibility) is checked here so
+        a bad plan fails the init RPC instead of hanging the first collective.
+        """
+        self._m2n = import_m2n()
+        self._metas = []
+
+        self._src_mesh = M2NMesh(
+            cast(tuple[int, int], tuple(init_info.src_mesh_dims)), 0
+        )
+        if self._src_mesh.size != init_info.rank_offset:
+            raise ValueError(
+                f"source mesh covers {self._src_mesh.size} ranks, but there are "
+                f"{init_info.rank_offset} trainer ranks"
+            )
+        # The inference topology, as declared by the trainer. This is the mesh
+        # a sharded destination is placed over; a replicated one is described
+        # by the alternate descriptor `resolve_layout` derives from it, which
+        # covers the same ranks but is not this mesh.
+        self._dst_mesh = M2NMesh(
+            cast(tuple[int, int], tuple(init_info.dst_mesh_dims)),
+            init_info.rank_offset,
+        )
+
+        for name, name_dtype, shape, placements in zip(
+            init_info.names,
+            init_info.dtype_names,
+            init_info.shapes,
+            init_info.src_placements,
+        ):
+            dtype = getattr(torch, name_dtype, None)
+            if not isinstance(dtype, torch.dtype):
+                raise ValueError(
+                    f"parameter '{name}' has invalid dtype name "
+                    f"'{name_dtype}'; expected the name of a torch.dtype"
+                )
+            check_transferable(name, dtype, shape)
+            codes = (
+                REPLICATED
+                if placements is None
+                else cast(Placements, tuple(placements))
+            )
+            src_mesh, src_placements = resolve_layout(
+                self._src_mesh, codes, f"parameter '{name}' source placements"
+            )
+            validate_layout(src_mesh, src_placements, shape, "source")
+            # Destination checks are deferred until loader-aware resolution;
+            # assuming replication here can reject a valid sharded plan.
+            self._metas.append(M2NParamMeta(name, dtype, tuple(shape), codes))
+
+        parallel_config = self.parallel_config
+        shard_axis_size = self._dst_mesh.dims[DESTINATION_SHARD_AXIS]
+        quantization = getattr(self.model_config, "quantization", None)
+        # A quantized weight loader may dequantize / transpose / requantize, and
+        # under PP a rank's local shape is not the checkpoint shape split over
+        # the shard axis. Both break shape-derived resolution, so take the
+        # fallback.
+        allow_direct = (
+            quantization is None and parallel_config.pipeline_parallel_size == 1
+        )
+        # Match every incoming checkpoint parameter to its rank-local target.
+        # Resolvable parameters can be resharded directly into live model
+        # storage; the rest receive a full tensor for `load_weights`.
+        self._parameter_destinations = resolve_parameter_destinations(
+            self.model,
+            [m.name for m in self._metas],
+            [m.dtype for m in self._metas],
+            [m.shape for m in self._metas],
+            num_workers=self._dst_mesh.size,
+            shard_axis_size=shard_axis_size,
+            allow_direct=allow_direct,
+        )
+        # Check the resolved layout rather than a provisional replicated one.
+        # For example, 32 source shards feeding 4 destination shards may need
+        # only 8 sources per destination, while a replicated plan needs all 32.
+        for meta, destination in zip(self._metas, self._parameter_destinations):
+            src_mesh, src_placements = resolve_layout(
+                self._src_mesh,
+                meta.placements,
+                f"parameter '{meta.name}' source placements",
+            )
+            dst_mesh, dst_placements = resolve_layout(
+                self._dst_mesh,
+                destination.placements,
+                f"parameter '{meta.name}' destination placements",
+            )
+            validate_layout(dst_mesh, dst_placements, meta.shape, "destination")
+            check_plan_limits(
+                (src_mesh, src_placements),
+                (dst_mesh, dst_placements),
+                meta.name,
+            )
+
+        # Update requests carry names only, so cache their plan indices. Track
+        # whether any fallback entry requires the `load_weights` lifecycle.
+        self._index = {meta.name: i for i, meta in enumerate(self._metas)}
+        self._uses_load_weights = any(
+            not destination.direct for destination in self._parameter_destinations
+        )
+
+        self.model_update_group = worker_init_process_group(
+            NCCLWeightTransferInitInfo(
+                master_address=init_info.master_address,
+                master_port=init_info.master_port,
+                nccl_unique_id_b64=init_info.nccl_unique_id_b64,
+                rank_offset=init_info.rank_offset,
+                world_size=init_info.world_size,
+            ),
+            parallel_config,
+        )
+        # Destinations are per-parameter and depend on the inference model, so
+        # the trainer cannot derive them. The first worker publishes the plan
+        # and every other rank checks its own against it — a silent
+        # disagreement between workers would mean mismatched collectives.
+        mine = [destination.placements for destination in self._parameter_destinations]
+        agreed = publish_destination_placements(
+            self.model_update_group, init_info.rank_offset, mine, len(self._metas)
+        )
+        if agreed != mine:
+            mismatched = next(
+                meta.name for meta, a, b in zip(self._metas, mine, agreed) if a != b
+            )
+            raise RuntimeError(
+                "inference workers resolved different nccl_m2n destinations "
+                f"(first mismatch: '{mismatched}'); all workers must run the "
+                "same model and parallel config"
+            )
+
+        self._handle = self._m2n.Handle.create(
+            self._m2n.Config(max_cta=init_info.max_cta)
+        )
+
+    def start_weight_update(self) -> None:
+        """Set up layerwise reloading, but only if some parameter needs it.
+
+        Directly-resharded parameters are written in place and never go through
+        `load_weights`, so a plan with no fallback entries has nothing to
+        reload.
+        """
+        if not self._uses_load_weights:
+            return
+        from vllm.model_executor.model_loader.reload import initialize_layerwise_reload
+
+        initialize_layerwise_reload(self.model)
+
+    def finish_weight_update(self) -> None:
+        """Finalize layerwise reloading when the plan uses fallback entries."""
+        if not self._uses_load_weights:
+            return
+        from vllm.model_executor.model_loader.reload import finalize_layerwise_reload
+
+        finalize_layerwise_reload(self.model, self.model_config)
+
+    def receive_weights(self, update_info: M2NWeightTransferUpdateInfo) -> None:
+        """Receive each requested parameter using its initialization-time plan."""
+        if self._handle is None or self.model_update_group is None:
+            raise RuntimeError(
+                "nccl_m2n weight transfer not initialized. "
+                "Call init_transfer_engine() first."
+            )
+
+        requested = []
+        for name in update_info.names:
+            index = self._index.get(name)
+            if index is None:
+                raise ValueError(
+                    f"parameter '{name}' was not declared at init; the "
+                    "trainer must send the same parameter set it announced"
+                )
+            requested.append((name, index))
+
+        from vllm.model_executor.model_loader.mtp_validation import (
+            disable_mtp_completeness_check,
+        )
+
+        comm = comm_ptr(self.model_update_group)
+        stream = torch.cuda.current_stream()
+
+        # Reshard lazily so one `load_weights` invocation sees the complete
+        # fallback sequence without staging every full tensor at once. Direct
+        # destinations still execute in request order as the iterator advances.
+        def reshard_requested_weights() -> Iterator[tuple[str, torch.Tensor]]:
+            for name, index in requested:
+                meta = self._metas[index]
+                destination = self._parameter_destinations[index]
+                buffer = destination.tensor
+                if buffer is None:
+                    buffer = torch.empty(
+                        meta.shape, dtype=meta.dtype, device=self.device
+                    )
+
+                self._reshard(comm, stream, meta, destination.placements, buffer)
+
+                if destination.tensor is None:
+                    # `load_weights` reads on the host stream, so the transfer
+                    # has to have landed before it runs.
+                    stream.synchronize()
+                    yield name, buffer
+
+        has_fallback = any(
+            self._parameter_destinations[index].tensor is None for _, index in requested
+        )
+        with disable_mtp_completeness_check():
+            received_weights = reshard_requested_weights()
+            if has_fallback:
+                # Some model loaders keep invocation-local state to combine
+                # related checkpoint tensors, so preserve one iterable scope.
+                self.model.load_weights(received_weights)
+            else:
+                # The iterator performs the reshards. With nothing to yield to
+                # `load_weights`, exhaust it here to execute direct transfers.
+                for _ in received_weights:
+                    pass
+
+    def _reshard(
+        self,
+        comm: int,
+        stream: torch.cuda.Stream,
+        src_meta: M2NParamMeta,
+        dst_placements: Placements | None,
+        dst_buffer: torch.Tensor,
+    ) -> None:
+        """Reshard one parameter from its trainer layout into a worker buffer.
+
+        `dst_placements` describes the buffer's logical placement over the
+        worker mesh; `dst_buffer` is this rank's physical storage.
+        """
+        assert self._src_mesh is not None
+        assert self._dst_mesh is not None
+        m2n = self._m2n
+        src_mesh, resolved_src_placements = resolve_layout(
+            self._src_mesh, src_meta.placements
+        )
+        dst_mesh, resolved_dst_placements = resolve_layout(
+            self._dst_mesh, dst_placements
+        )
+        m2n.reshard(
+            None,
+            dst_buffer,
+            comm,
+            stream,
+            src_mesh=to_mesh(m2n, src_mesh),
+            src_placements=to_placements(m2n, resolved_src_placements),
+            dst_mesh=to_mesh(m2n, dst_mesh),
+            dst_placements=to_placements(m2n, resolved_dst_placements),
+            handle=self._handle,
+        )
+
+    def shutdown(self) -> None:
+        """Finish pending GPU work and release M2N communication state."""
+        if self._handle is not None:
+            # M2N does not synchronize caller streams on finalize.
+            torch.accelerator.synchronize()
+            self._handle.destroy()
+            self._handle = None
+        self.model_update_group = None
+        self._parameter_destinations = []

vllm/distributed/weight_transfer/m2n_layout.py

@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Destination-layout resolution for the NCCL M2N weight transfer backend.
+
+For every incoming checkpoint parameter the worker needs a destination buffer
+plus its placement over the inference mesh. Two outcomes are possible:
+
+* **direct** — when the model is not quantized, the parameter maps 1:1 onto a
+  live vLLM parameter and its loader is a small, known-safe copy or tensor-
+  parallel loader. The loader's declared input/output dimension determines the
+  placement; shapes are validation, not an inference mechanism. The reshard
+  writes straight into the live parameter, so each rank receives only its own
+  shard and nothing is copied afterwards.
+* **fallback** — anything else. The reshard delivers the whole tensor to every
+  rank and `load_weights` does the sharding, exactly as the broadcast NCCL
+  backend does. Fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`)
+  take this path: the checkpoint name does not name a vLLM parameter, so there
+  is nothing to resolve against.
+
+Correctness never depends on a parameter resolving — the fallback is always
+available and is the same path the existing backend uses.
+"""
+
+import inspect
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from functools import partial
+from typing import cast
+
+import torch
+
+from vllm.distributed.weight_transfer.m2n_common import (
+    DESTINATION_REPLICA_AXIS,
+    DESTINATION_SHARD_AXIS,
+    MESH_NDIMS,
+    REPLICATE,
+    REPLICATED,
+    Placements,
+)
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_DIRECT_WEIGHT_LOADING_BLOCKLIST = {
+    "vllm.model_executor.models.gpt2.GPT2Model",
+}
+
+
+def _module_type_name(module: torch.nn.Module) -> str:
+    cls = type(module)
+    return f"{cls.__module__}.{cls.__name__}"
+
+
+def _destination_placements(shard_dim: int) -> Placements:
+    """Place one tensor shard on the configured destination mesh axis."""
+    placements = [REPLICATE] * MESH_NDIMS
+    placements[DESTINATION_REPLICA_AXIS] = REPLICATE
+    placements[DESTINATION_SHARD_AXIS] = shard_dim
+    return cast(Placements, tuple(placements))
+
+
+@dataclass
+class M2NDestination:
+    """Where one checkpoint parameter lands on this worker."""
+
+    name: str
+    placements: Placements | None
+    """Placement over the inference mesh, or `REPLICATED` for the fallback."""
+    tensor: torch.Tensor | None
+    """The live parameter view to reshard into, or None for the fallback path
+    (the engine allocates a full replica per round and calls `load_weights`)."""
+
+    @property
+    def direct(self) -> bool:
+        return self.tensor is not None
+
+
+def _base_loader_shard_dim(
+    param: torch.nn.Parameter,
+    shard_axis_size: int,
+) -> int | None:
+    """Return a known-safe loader's declared shard dim, or reject it.
+
+    Loader identity is deliberately fail-closed. A missing loader is not
+    treated as ``default_weight_loader`` because a model-level ``load_weights``
+    implementation may transform the tensor before it reaches that default.
+    Wrappers and partials are also rejected: recognizing their wrapped callable
+    would ignore behavior added by the wrapper.
+    """
+    loader = getattr(param, "weight_loader", None)
+    if loader is None or isinstance(loader, partial):
+        return None
+
+    owner = loader.__self__ if inspect.ismethod(loader) else None
+    loader_fn = loader.__func__ if inspect.ismethod(loader) else loader
+
+    # Lazy imports avoid making the weight-transfer registry import all model
+    # executor layers during normal vLLM startup.
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        ReplicatedLinear,
+        RowParallelLinear,
+    )
+    from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+    replicated_loaders = {
+        default_weight_loader,
+        ReplicatedLinear.weight_loader,
+    }
+    column_loaders = {
+        ColumnParallelLinear.weight_loader,
+        ColumnParallelLinear.weight_loader_v2,
+    }
+    row_loaders = {
+        RowParallelLinear.weight_loader,
+        RowParallelLinear.weight_loader_v2,
+    }
+
+    # These attributes signal packing, a transform, or pre-sharded checkpoint
+    # storage. Such semantics cannot be reproduced by a plain M2N placement.
+    # Check them before accepting replicated loaders too: model-level loading
+    # may transform a tensor before handing it to an otherwise plain copier.
+    metadata_attrs = ("packed_dim", "packed_factor", "marlin_tile_size")
+    flag_attrs = ("needs_scalar_to_array", "is_sharded_weight", "is_transposed")
+    if any(getattr(param, attr, None) is not None for attr in metadata_attrs):
+        return None
+    if any(bool(getattr(param, attr, False)) for attr in flag_attrs):
+        return None
+
+    if loader_fn in replicated_loaders:
+        return REPLICATE
+    if loader_fn not in column_loaders | row_loaders:
+        return None
+
+    if owner is None or getattr(owner, "tp_size", None) != shard_axis_size:
+        return None
+
+    attr = "output_dim" if loader_fn in column_loaders else "input_dim"
+    dim = getattr(param, attr, None)
+    return dim if isinstance(dim, int) and dim >= 0 else None
+
+
+def _validated_shard_dim(
+    global_shape: Sequence[int],
+    local_shape: Sequence[int],
+    declared_dim: int,
+    shard_axis_size: int,
+) -> int | None:
+    """Validate a loader-declared shard dimension against both shapes."""
+    if len(global_shape) != len(local_shape):
+        return None
+
+    if declared_dim == REPLICATE:
+        return REPLICATE if tuple(global_shape) == tuple(local_shape) else None
+    if declared_dim >= len(global_shape):
+        return None
+    for dim, (whole, local) in enumerate(zip(global_shape, local_shape)):
+        expected = local * shard_axis_size if dim == declared_dim else local
+        if whole != expected:
+            return None
+    return declared_dim
+
+
+def resolve_parameter_destinations(
+    model: torch.nn.Module,
+    names: Sequence[str],
+    dtypes: Sequence[torch.dtype],
+    shapes: Sequence[Sequence[int]],
+    *,
+    num_workers: int,
+    shard_axis_size: int,
+    allow_direct: bool,
+) -> list[M2NDestination]:
+    """Build this worker's destination plan, one entry per checkpoint parameter.
+
+    Placements are relative to the inference mesh: axis 0 replicates and axis 1
+    shards. `allow_direct=False` forces every parameter onto the fallback path;
+    the engine sets it for pipeline-parallel or quantized deployments, where a
+    parameter's local shape is not simply the checkpoint shape split across the
+    shard axis. Parameters below a model type with known model-level checkpoint
+    preprocessing also fall back, since bypassing `load_weights` would skip the
+    transformation.
+    """
+    params = dict(model.named_parameters()) if allow_direct else {}
+    blocked_param_ids = {
+        id(param)
+        for module in model.modules()
+        if _module_type_name(module) in _DIRECT_WEIGHT_LOADING_BLOCKLIST
+        for param in module.parameters()
+    }
+
+    destinations: list[M2NDestination] = []
+    for name, dtype, shape in zip(names, dtypes, shapes):
+        param = params.get(name)
+        dim = None
+        if (
+            param is not None
+            and id(param) not in blocked_param_ids
+            and param.dtype == dtype
+            and param.data.is_contiguous()
+            and num_workers % shard_axis_size == 0
+        ):
+            declared_dim = _base_loader_shard_dim(param, shard_axis_size)
+            if declared_dim is not None:
+                dim = _validated_shard_dim(
+                    shape, param.shape, declared_dim, shard_axis_size
+                )
+
+        if dim is None:
+            # A replicated parameter is identical on every rank, so it needs no
+            # placement of its own — REPLICATED lets resolve_layout spread it
+            # over all the workers regardless of how the mesh is factored.
+            destinations.append(M2NDestination(name, REPLICATED, None))
+        elif dim == REPLICATE:
+            assert param is not None
+            destinations.append(M2NDestination(name, REPLICATED, param.data))
+        else:
+            assert param is not None
+            destinations.append(
+                M2NDestination(name, _destination_placements(dim), param.data)
+            )
+
+    # Layerwise reload finalizes a module as a unit. Mixing direct and fallback
+    # parameters within one module can copy stale storage over the direct load.
+    fallback_modules = {
+        name.rpartition(".")[0]
+        for name, destination in zip(names, destinations)
+        if not destination.direct and name in params
+    }
+    destinations = [
+        M2NDestination(destination.name, REPLICATED, None)
+        if destination.direct
+        and destination.name.rpartition(".")[0] in fallback_modules
+        else destination
+        for destination in destinations
+    ]
+
+    num_direct = sum(d.direct for d in destinations)
+    parameter_bytes = [
+        math.prod(shape) * dtype.itemsize for dtype, shape in zip(dtypes, shapes)
+    ]
+    total_bytes = sum(parameter_bytes)
+    direct_bytes = sum(
+        nbytes
+        for destination, nbytes in zip(destinations, parameter_bytes)
+        if destination.direct
+    )
+    direct_byte_percentage = 100 * direct_bytes / total_bytes if total_bytes else 0
+    logger.info(
+        "nccl_m2n destination plan: %d/%d parameters resharded directly into "
+        "the model, %d via full-tensor fallback; direct byte coverage: "
+        "%d/%d bytes (%.1f%%)",
+        num_direct,
+        len(destinations),
+        len(destinations) - num_direct,
+        direct_bytes,
+        total_bytes,
+        direct_byte_percentage,
+    )
+    return destinations

vllm/distributed/weight_transfer/m2n_source.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Trainer-side weight sources for the NCCL M2N backend.
+
+m2n plans a transfer from both sides' layouts, so the trainer has to say how it
+holds each parameter — which the base `WeightSource` / `ParamMeta` pair does not
+express. This module supplies the m2n flavor of both, and the DTensor-backed
+implementation that covers the common trainer.
+
+A source declares its mesh once (`mesh()`) and a placement per parameter, since
+one topology describes every tensor on the side.
+"""
+
+from collections.abc import Iterator
+from typing import Any
+
+import torch
+
+from vllm.distributed.weight_transfer.base import ParamMeta, WeightSource
+from vllm.distributed.weight_transfer.m2n_common import (
+    REPLICATE,
+    REPLICATED,
+    M2NMesh,
+    M2NParamMeta,
+    Placements,
+)
+
+__all__ = [
+    "M2NWeightSource",
+    "DTensorModuleSource",
+    "mesh_from_tensor",
+    "placements_from_tensor",
+]
+
+
+class M2NWeightSource(WeightSource):
+    """A `WeightSource` that also describes how the trainer holds its weights.
+
+    Unlike `ModuleSource`, iteration yields each rank's **local shard**, not a
+    materialized full tensor: gathering is exactly the cost m2n removes.
+    """
+
+    def mesh(self) -> M2NMesh:
+        """The trainer's rank topology, shared by every parameter."""
+        raise NotImplementedError
+
+    def metadata(self) -> list[ParamMeta]:
+        """Name, dtype, full shape, and trainer placement for each parameter."""
+        raise NotImplementedError
+
+    def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
+        """Yield `(name, local shard)` pairs in the same order as `metadata()`."""
+        raise NotImplementedError
+
+
+def _placement_code(placement: Any) -> int:
+    """Map a `torch.distributed` placement onto an m2n placement code."""
+    name = type(placement).__name__
+    if name == "Replicate":
+        return REPLICATE
+    if name == "Shard":
+        return int(placement.dim)
+    raise ValueError(
+        f"nccl_m2n cannot express the {name} placement; only Replicate and "
+        "Shard are supported"
+    )
+
+
+def mesh_from_tensor(tensor: torch.Tensor, num_trainer_ranks: int) -> M2NMesh:
+    """The mesh a parameter lives on, as an `M2NMesh`.
+
+    A plain tensor is identical on every trainer rank, so it spans all of them.
+    """
+    device_mesh = getattr(tensor, "device_mesh", None)
+    if device_mesh is None:
+        return M2NMesh((num_trainer_ranks, 1), 0)
+
+    grid = device_mesh.mesh
+    ranks = grid.flatten().tolist()
+    if ranks != list(range(len(ranks))):
+        raise ValueError(
+            "nccl_m2n requires the trainer's device mesh to cover the "
+            f"contiguous rank interval [0, {len(ranks)}); got {ranks}"
+        )
+    if grid.ndim > 2:
+        raise ValueError(
+            f"nccl_m2n supports 1-D and 2-D device meshes, got {grid.ndim}-D"
+        )
+    dims = tuple(grid.shape)
+    return M2NMesh(dims if len(dims) == 2 else (dims[0], 1), 0)
+
+
+def placements_from_tensor(tensor: torch.Tensor) -> Placements | None:
+    """How a parameter is placed over its mesh, or `REPLICATED`.
+
+    `REPLICATED` is returned rather than a `(REPLICATE, REPLICATE)` pair: m2n
+    cannot express that, and the size-1-axis encoding it needs instead is
+    applied later by `resolve_layout`, which owns that workaround.
+    """
+    placements = getattr(tensor, "placements", None)
+    if placements is None:
+        return REPLICATED
+    codes = [_placement_code(p) for p in placements]
+    if all(code == REPLICATE for code in codes):
+        return REPLICATED
+    if len(codes) == 1:
+        return (codes[0], REPLICATE)
+    return (codes[0], codes[1])
+
+
+class DTensorModuleSource(M2NWeightSource):
+    """`M2NWeightSource` over `module.named_parameters()`.
+
+    Covers both the FSDP/DTensor trainer (placement read off each parameter,
+    local shard yielded via `to_local()`) and the plain replicated trainer, with
+    no special casing. Trainers with a custom producer (a Megatron export, MoE
+    re-fusing) subclass `M2NWeightSource` instead.
+    """
+
+    def __init__(self, module: torch.nn.Module, num_trainer_ranks: int) -> None:
+        """Wrap `module`; `num_trainer_ranks` sizes the mesh for plain tensors."""
+        self._module = module
+        self._num_trainer_ranks = num_trainer_ranks
+
+    def mesh(self) -> M2NMesh:
+        """The mesh every sharded parameter agrees on.
+
+        Replicated parameters span the trainer without a device mesh of their
+        own, so they never decide this; a model whose *sharded* parameters
+        disagree cannot be described by one side mesh and is rejected.
+        """
+        meshes = {
+            mesh_from_tensor(p, self._num_trainer_ranks)
+            for _, p in self._module.named_parameters()
+            if placements_from_tensor(p) is not REPLICATED
+        }
+        if len(meshes) > 1:
+            raise ValueError(
+                "nccl_m2n needs one mesh per side, but this module's sharded "
+                f"parameters span several: {sorted(m.dims for m in meshes)}"
+            )
+        return meshes.pop() if meshes else M2NMesh((self._num_trainer_ranks, 1), 0)
+
+    def metadata(self) -> list[ParamMeta]:
+        """Read global shape/dtype and placements without gathering shards."""
+        return [
+            M2NParamMeta(name, p.dtype, tuple(p.shape), placements_from_tensor(p))
+            for name, p in self._module.named_parameters()
+        ]
+
+    def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
+        """Yield each parameter's local shard (`to_local()`), or the tensor itself."""
+        for name, param in self._module.named_parameters():
+            to_local = getattr(param, "to_local", None)
+            yield name, (to_local() if callable(to_local) else param)

vllm/distributed/weight_transfer/m2n_trainer.py

@@ -0,0 +1,414 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Trainer-side weight transfer engine for the NCCL M2N backend.
+
+Symmetric to `M2NWeightTransferEngine` but in the training process. Every
+trainer rank joins the shared communicator and runs every reshard, sending its
+own local shard; only rank 0 touches the inference control plane.
+
+Workers join the communicator during init and run reshard from inside
+`update_weights`, so those two RPCs overlap with work on this side. Both run on
+a helper thread and are joined afterwards, which keeps that overlap inside the
+engine -- where the `TrainerWeightTransferEngine` contract puts it -- rather than
+in every caller.
+"""
+
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, ClassVar, cast
+
+import torch
+from typing_extensions import Self
+
+from vllm.distributed.weight_transfer.base import (
+    TrainerInitInfo,
+    TrainerWeightTransferEngine,
+    VLLMWeightSyncClient,
+    WeightSource,
+)
+from vllm.distributed.weight_transfer.m2n_common import (
+    M2NMesh,
+    M2NParamMeta,
+    Placements,
+    check_transferable,
+    comm_ptr,
+    import_m2n,
+    publish_destination_placements,
+    resolve_layout,
+    to_mesh,
+    to_placements,
+    validate_layout,
+)
+from vllm.distributed.weight_transfer.m2n_source import M2NWeightSource
+from vllm.distributed.weight_transfer.nccl_common import (
+    decode_nccl_unique_id,
+    trainer_init,
+    uid_init_process_group,
+)
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+
+logger = init_logger(__name__)
+
+__all__ = ["M2NTrainerInitInfo", "M2NTrainerWeightTransferEngine"]
+
+
+def _dtype_name(dtype: torch.dtype) -> str:
+    """Short name for a dtype, e.g. `bfloat16`, for the worker-side init dict."""
+    return str(dtype).split(".")[-1]
+
+
+@dataclass(kw_only=True)
+class M2NTrainerInitInfo(TrainerInitInfo):
+    """Trainer-side init info for nccl_m2n.
+
+    `rank` (from `TrainerInitInfo`) is this trainer process's rank; it is also
+    its rank in the shared communicator, since the trainer occupies
+    `[0, num_trainer_ranks)`. Rank 0 drives the control plane.
+    """
+
+    backend: ClassVar[str] = "nccl_m2n"
+
+    master_address: str | None = None
+    master_port: int | None = None
+    world_size: int
+    """Trainer ranks + all inference workers."""
+    num_trainer_ranks: int = 1
+    dst_mesh_dims: tuple[int, int] | None = None
+    """How the inference ranks are laid out: axis 0 replicates and axis 1
+    shards. The trainer declares it so both sides describe the destination
+    identically; defaults to a flat `(num_workers, 1)`, which is all a
+    replicated destination needs."""
+    max_cta: int | None = None
+    nccl_unique_id_b64: str | None = field(default=None, repr=False)
+
+    @property
+    def destination_mesh_dims(self) -> tuple[int, int]:
+        """`dst_mesh_dims`, or a flat mesh over every inference worker."""
+        if self.dst_mesh_dims is not None:
+            return self.dst_mesh_dims
+        return (self.world_size - self.num_trainer_ranks, 1)
+
+    def __post_init__(self) -> None:
+        """Reject rank counts or destination meshes that cannot form a group."""
+        _ = self.nccl_unique_id_bytes
+        if not 0 < self.num_trainer_ranks < self.world_size:
+            raise ValueError(
+                f"`num_trainer_ranks` ({self.num_trainer_ranks}) must leave at "
+                f"least one worker in world_size {self.world_size}"
+            )
+        num_workers = self.world_size - self.num_trainer_ranks
+        if self.dst_mesh_dims is not None:
+            dims = tuple(self.dst_mesh_dims)
+            if len(dims) != 2 or dims[0] * dims[1] != num_workers:
+                raise ValueError(
+                    f"`dst_mesh_dims` {dims} must be two dims covering the "
+                    f"{num_workers} inference workers"
+                )
+        if not 0 <= self.rank < self.num_trainer_ranks:
+            raise ValueError(
+                f"trainer `rank` ({self.rank}) must be below "
+                f"`num_trainer_ranks` ({self.num_trainer_ranks})"
+            )
+
+    @property
+    def nccl_unique_id_bytes(self) -> bytes | None:
+        return decode_nccl_unique_id(
+            master_address=self.master_address,
+            master_port=self.master_port,
+            nccl_unique_id_b64=self.nccl_unique_id_b64,
+            ctx="M2NTrainerInitInfo",
+        )
+
+
+class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitInfo]):
+    """Trainer-side engine: sends every rank's local shard, once per parameter.
+
+    Called on every trainer rank. All ranks join the communicator and run every
+    reshard; only rank 0 touches the control plane.
+    """
+
+    init_info_cls = M2NTrainerInitInfo
+
+    def __init__(
+        self,
+        *,
+        client: VLLMWeightSyncClient,
+        source: M2NWeightSource,
+        is_controller: bool,
+        num_trainer_ranks: int,
+    ) -> None:
+        """Hold the client and source; `trainer_init` fills in the group."""
+        super().__init__(client=client, source=source, is_sender=is_controller)
+        self.source: M2NWeightSource = source
+        self.is_controller = is_controller
+        self.num_trainer_ranks = num_trainer_ranks
+        self.group: PyNcclCommunicator | None = None
+        self._m2n: Any = None
+        self._handle: Any = None
+        self._metas: list[M2NParamMeta] = []
+        self._src_mesh: M2NMesh | None = None
+        self._dst_mesh: M2NMesh | None = None
+        self._dst_placements: list[Placements | None] = []
+        self._executor: ThreadPoolExecutor | None = None
+
+    @classmethod
+    def trainer_init(
+        cls,
+        init_info: M2NTrainerInitInfo,
+        *,
+        client: VLLMWeightSyncClient,
+        source: WeightSource | None = None,
+    ) -> Self:
+        """Build the engine and rendezvous with the inference side.
+
+        Runs on *every* trainer rank. Rank 0 additionally ships the transfer
+        plan to the workers and drives the control plane; the other ranks only
+        build local state and join the communicator. Every rank participates
+        in every reshard and sends its local shard.
+
+        The ordering here is not arbitrary -- see the comments inline. In
+        short: validate before anything blocks; start the init RPC without
+        waiting (workers join the communicator inside that RPC, so it cannot
+        return until this side has joined too); then join the NCCL
+        communicator; then create the M2N handle; and only then wait for the RPC
+        to finish.
+        """
+        # The base WeightSource carries name/dtype/shape but not layout, which
+        # is not enough to plan a reshard. Fail here rather than at the first
+        # missing attribute deep inside the send loop.
+        if not isinstance(source, M2NWeightSource):
+            raise TypeError(
+                "nccl_m2n needs per-parameter layouts, so its source must be an "
+                f"M2NWeightSource (e.g. DTensorModuleSource); got "
+                f"{type(source).__name__}"
+            )
+
+        # Trainer-side weight transfer engine for this rank. `__init__` only
+        # holds the client and source; the communicator, meshes, and M2N handle
+        # are filled in below.
+        engine = cls(
+            client=client,
+            source=source,
+            is_controller=init_info.rank == 0,
+            num_trainer_ranks=init_info.num_trainer_ranks,
+        )
+        engine._m2n = import_m2n()
+
+        engine._src_mesh = source.mesh()
+        if engine._src_mesh.size != init_info.num_trainer_ranks:
+            raise ValueError(
+                f"source mesh covers {engine._src_mesh.size} ranks, but there "
+                f"are {init_info.num_trainer_ranks} trainer ranks"
+            )
+
+        # Validate the whole plan before any rendezvous. A layout m2n cannot
+        # express is a mismatched collective at send time, which hangs rather
+        # than raises -- so every precondition is checked while an exception can
+        # still propagate to the caller.
+        engine._metas = cast(list[M2NParamMeta], source.metadata())
+        for meta in engine._metas:
+            check_transferable(meta.name, meta.dtype, meta.shape)
+            mesh, placements = resolve_layout(engine._src_mesh, meta.placements)
+            validate_layout(mesh, placements, meta.shape, "source")
+
+        rendezvous: Future | None = None
+        if engine.is_controller:
+            engine._executor = ThreadPoolExecutor(
+                max_workers=1, thread_name_prefix="m2n-weight-sync"
+            )
+            # Only trainer rank 0 starts the control-plane RPC. Workers join
+            # the communicator inside it, so submit it first, join below, and
+            # wait for it at the end.
+            rendezvous = engine._executor.submit(
+                client.init_weight_transfer_engine,
+                engine._worker_init_info(init_info),
+            )
+
+        # Every trainer rank joins -- reshard is collective over the whole
+        # communicator. Trainer ranks take
+        # [0, num_trainer_ranks) and the workers follow, which is the
+        # contiguous-interval layout m2n requires of both meshes.
+        unique_id_bytes = init_info.nccl_unique_id_bytes
+        if unique_id_bytes is not None:
+            engine.group = uid_init_process_group(
+                unique_id_bytes,
+                rank=init_info.rank,
+                world_size=init_info.world_size,
+                device=torch.accelerator.current_device_index(),
+            )
+        else:
+            assert init_info.master_address is not None
+            assert init_info.master_port is not None
+            engine.group = trainer_init(
+                {
+                    "master_address": init_info.master_address,
+                    "master_port": init_info.master_port,
+                    "world_size": init_info.world_size,
+                },
+                rank=init_info.rank,
+            )
+
+        # The trainer declares the inference topology; the worker uses exactly
+        # this, so neither side has to infer the other's factorization.
+        engine._dst_mesh = M2NMesh(
+            init_info.destination_mesh_dims, init_info.num_trainer_ranks
+        )
+        engine._dst_placements = engine._receive_destination_plan(
+            init_info.num_trainer_ranks
+        )
+        logger.info(
+            "nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), "
+            "inference ranks [%d, %d)",
+            len(engine._metas),
+            init_info.num_trainer_ranks,
+            init_info.num_trainer_ranks,
+            init_info.world_size,
+        )
+
+        # Created after the group so a failed rendezvous does not leave a live
+        # handle behind; shutdown() releases it once caller streams are synced.
+        engine._handle = engine._m2n.Handle.create(
+            engine._m2n.Config(max_cta=init_info.max_cta)
+        )
+
+        # Now safe to collect: both sides have joined, so the RPC has returned
+        # or raised. Doing this last means a worker-side init failure surfaces
+        # here, as an exception from trainer_init, rather than at the first send.
+        if rendezvous is not None:
+            rendezvous.result()
+        return engine
+
+    def _receive_destination_plan(
+        self, first_worker_rank: int
+    ) -> list[Placements | None]:
+        """Receive and validate the inference-side per-parameter placements."""
+        if self.group is None or self._dst_mesh is None:
+            raise RuntimeError(
+                "destination plan requires the communicator and destination mesh"
+            )
+
+        # The first inference worker publishes one placement per parameter.
+        # Every trainer rank receives the complete plan in this single,
+        # initialization-time broadcast.
+        placements = publish_destination_placements(
+            self.group, first_worker_rank, None, len(self._metas)
+        )
+        # The send loop matches placements to metadata by position, so reject
+        # an incomplete or oversized plan before entering any M2N collective.
+        if len(placements) != len(self._metas):
+            raise RuntimeError(
+                f"inference side planned {len(placements)} parameters, but "
+                f"this trainer announced {len(self._metas)}"
+            )
+        # Check every worker-selected destination against the corresponding
+        # global parameter shape before the first transfer.
+        for meta, placement in zip(self._metas, placements):
+            mesh, resolved = resolve_layout(self._dst_mesh, placement)
+            validate_layout(mesh, resolved, meta.shape, "destination")
+        return placements
+
+    def _worker_init_info(self, init_info: M2NTrainerInitInfo) -> dict[str, Any]:
+        """Handshake payload: rendezvous, both meshes, and the transfer plan."""
+        assert self._src_mesh is not None
+        payload: dict[str, Any] = {
+            "master_address": init_info.master_address,
+            "master_port": init_info.master_port,
+            "nccl_unique_id_b64": init_info.nccl_unique_id_b64,
+            "rank_offset": init_info.num_trainer_ranks,
+            "world_size": init_info.world_size,
+            "src_mesh_dims": list(self._src_mesh.dims),
+            "dst_mesh_dims": list(init_info.destination_mesh_dims),
+            "names": [m.name for m in self._metas],
+            "dtype_names": [_dtype_name(m.dtype) for m in self._metas],
+            "shapes": [list(m.shape) for m in self._metas],
+            "src_placements": [
+                None if m.placements is None else list(m.placements)
+                for m in self._metas
+            ],
+            "max_cta": init_info.max_cta,
+        }
+        return {key: value for key, value in payload.items() if value is not None}
+
+    def send_weights(self) -> None:
+        """Drive one update round: start, reshard concurrently, then finish."""
+        if self._handle is None or self.group is None:
+            raise RuntimeError("nccl_m2n trainer engine is not initialized")
+
+        update: Future | None = None
+        if self.is_controller:
+            self.client.start_weight_update()
+            # The workers reshard from inside `update_weights`, concurrently
+            # with the sends below.
+            assert self._executor is not None
+            update = self._executor.submit(
+                self.client.update_weights,
+                {"names": [m.name for m in self._metas]},
+            )
+
+        try:
+            self._send()
+        finally:
+            if update is not None:
+                update.result()
+
+        if self.is_controller:
+            self.client.finish_weight_update()
+
+    def _send(self) -> None:
+        """Reshard each local shard into its worker-planned destination."""
+        assert self.group is not None
+        assert self._src_mesh is not None
+        assert self._dst_mesh is not None
+        m2n = self._m2n
+        comm = comm_ptr(self.group)
+        stream = torch.cuda.current_stream()
+
+        for index, (name, tensor) in enumerate(self.source):
+            meta = self._metas[index]
+            if name != meta.name:
+                raise RuntimeError(
+                    f"weight source yielded '{name}' at position {index}, but "
+                    f"its metadata declared '{meta.name}'; the source must "
+                    "iterate in a stable order"
+                )
+            shard = tensor.detach().contiguous()
+            src_mesh, src_placements = resolve_layout(self._src_mesh, meta.placements)
+            dst_mesh, dst_placements = resolve_layout(
+                self._dst_mesh, self._dst_placements[index]
+            )
+            m2n.reshard(
+                shard,
+                None,
+                comm,
+                stream,
+                src_mesh=to_mesh(m2n, src_mesh),
+                src_placements=to_placements(m2n, src_placements),
+                dst_mesh=to_mesh(m2n, dst_mesh),
+                dst_placements=to_placements(m2n, dst_placements),
+                handle=self._handle,
+            )
+            # A custom source may yield temporary storage allocated on another
+            # stream. M2N retains only its device pointer for asynchronous work,
+            # so tell the allocator not to reuse that storage until this stream
+            # has finished with it. For storage allocated on this stream,
+            # record_stream is effectively a no-op.
+            shard.record_stream(stream)
+
+        # As in the other weight-transfer backends, wait before returning
+        # rather than relying on same-stream ordering. The caller may then
+        # safely mutate parameters or continue work on another stream.
+        stream.synchronize()
+
+    def shutdown(self) -> None:
+        """Destroy the m2n handle, join the helper thread, and drop the group."""
+        if self._handle is not None:
+            torch.accelerator.synchronize()
+            self._handle.destroy()
+            self._handle = None
+        if self._executor is not None:
+            self._executor.shutdown()
+            self._executor = None
+        self.group = None

vllm/distributed/weight_transfer/nccl_common.py

@@ -220,11 +220,11 @@ def worker_init_process_group(
 
 def trainer_init(
     init_info: NCCLRendezvous | dict,
+    rank: int = 0,
 ) -> "PyNcclCommunicator":
     """Initialize NCCL process group for trainer-side weight transfer.
 
-    The trainer is always rank 0 in the process group. Uses the current
-    CUDA device (torch.accelerator.current_device_index()).
+    Uses the current CUDA device (torch.accelerator.current_device_index()).
 
     Args:
         init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or
@@ -232,6 +232,9 @@ def trainer_init(
             - master_address: str
             - master_port: int
             - world_size: int
+        rank: This trainer process's rank in the group. The broadcast backends
+            have a single trainer at rank 0; multi-rank trainers (m2n) occupy
+            ranks `[0, num_trainer_ranks)` and the workers start after them.
 
     Returns:
         PyNcclCommunicator for weight transfer.
@@ -246,12 +249,11 @@ def trainer_init(
         master_port = init_info.master_port
         world_size = init_info.world_size
 
-    # Trainer is always rank 0
     device = torch.accelerator.current_device_index()
     return stateless_init_process_group(
         master_address,
         master_port,
-        0,
+        rank,
         world_size,
         device,
     )