Skip to content

vllm.distributed.weight_transfer.m2n_source

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. Each parameter owns its source mesh and placements because dense and expert weights can use different factorizations over the same trainer ranks.

Classes:

Functions:

DTensorModuleSource

Bases: 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.

Methods:

  • __init__ –

    Wrap module; num_trainer_ranks sizes the mesh for plain tensors.

  • __iter__ –

    Yield each parameter's local shard (to_local()), or the tensor itself.

  • metadata –

    Read each global shape and source layout without gathering shards.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
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 metadata(self) -> list[M2NParamMeta]:  # type: ignore[override]
        """Read each global shape and source layout without gathering shards."""
        return [
            M2NParamMeta(
                name,
                p.dtype,
                tuple(p.shape),
                M2NLayout(
                    mesh_from_tensor(p, self._num_trainer_ranks),
                    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)

__init__(module, num_trainer_ranks)

Wrap module; num_trainer_ranks sizes the mesh for plain tensors.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
183
184
185
186
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

__iter__()

Yield each parameter's local shard (to_local()), or the tensor itself.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
203
204
205
206
207
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)

metadata()

Read each global shape and source layout without gathering shards.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def metadata(self) -> list[M2NParamMeta]:  # type: ignore[override]
    """Read each global shape and source layout without gathering shards."""
    return [
        M2NParamMeta(
            name,
            p.dtype,
            tuple(p.shape),
            M2NLayout(
                mesh_from_tensor(p, self._num_trainer_ranks),
                placements_from_tensor(p),
            ),
        )
        for name, p in self._module.named_parameters()
    ]

M2NManifestEntry dataclass

One immutable source descriptor with a lazy rank-local tensor view.

Attributes:

Source code in vllm/distributed/weight_transfer/m2n_source.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
@dataclass(frozen=True)
class M2NManifestEntry:
    """One immutable source descriptor with a lazy rank-local tensor view."""

    key: str
    dtype: torch.dtype
    global_shape: tuple[int, ...]
    source_layout: M2NLayout
    local_tensor: Callable[[], Any]

    def __post_init__(self) -> None:
        object.__setattr__(self, "global_shape", tuple(self.global_shape))
        if not self.key:
            raise ValueError("M2N manifest key must not be empty")
        if not isinstance(self.dtype, torch.dtype):
            raise TypeError(
                f"manifest entry '{self.key}' has invalid dtype {self.dtype!r}"
            )
        if not isinstance(self.source_layout, M2NLayout):
            raise TypeError(
                f"manifest entry '{self.key}' source_layout must be an M2NLayout"
            )
        if not callable(self.local_tensor):
            raise TypeError(
                f"manifest entry '{self.key}' local_tensor must be callable"
            )

    @property
    def metadata(self) -> M2NParamMeta:
        """Return the serializable portion consumed during initialization."""
        return M2NParamMeta(
            self.key,
            self.dtype,
            self.global_shape,
            self.source_layout,
        )

metadata property

Return the serializable portion consumed during initialization.

M2NWeightSource

Bases: 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.

Methods:

  • __iter__ –

    Yield (name, local buffer) pairs in metadata order.

  • metadata –

    Name, dtype, full shape, and source layout for each parameter.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
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 metadata(self) -> list[M2NParamMeta]:  # type: ignore[override]
        """Name, dtype, full shape, and source layout for each parameter."""
        raise NotImplementedError

    def __iter__(self) -> Iterator[tuple[str, Any]]:  # type: ignore[override]
        """Yield `(name, local buffer)` pairs in metadata order.

        Providers must enqueue materialization on the caller's current CUDA
        stream, or return a buffer whose writes are already visible to it.
        Buffers may be torch tensors, CUDA Array Interface objects, or raw
        device pointers accepted by nccl4py. Providers own correct
        communicator-device placement. Raw-pointer providers additionally own
        allocation lifetime, dtype, and shape correctness.
        """
        raise NotImplementedError

__iter__()

Yield (name, local buffer) pairs in metadata order.

Providers must enqueue materialization on the caller's current CUDA stream, or return a buffer whose writes are already visible to it. Buffers may be torch tensors, CUDA Array Interface objects, or raw device pointers accepted by nccl4py. Providers own correct communicator-device placement. Raw-pointer providers additionally own allocation lifetime, dtype, and shape correctness.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
50
51
52
53
54
55
56
57
58
59
60
def __iter__(self) -> Iterator[tuple[str, Any]]:  # type: ignore[override]
    """Yield `(name, local buffer)` pairs in metadata order.

    Providers must enqueue materialization on the caller's current CUDA
    stream, or return a buffer whose writes are already visible to it.
    Buffers may be torch tensors, CUDA Array Interface objects, or raw
    device pointers accepted by nccl4py. Providers own correct
    communicator-device placement. Raw-pointer providers additionally own
    allocation lifetime, dtype, and shape correctness.
    """
    raise NotImplementedError

metadata()

Name, dtype, full shape, and source layout for each parameter.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
46
47
48
def metadata(self) -> list[M2NParamMeta]:  # type: ignore[override]
    """Name, dtype, full shape, and source layout for each parameter."""
    raise NotImplementedError

ManifestM2NWeightSource

Bases: M2NWeightSource

An ordered M2N source backed by lazy, rank-local manifest entries.

Methods:

  • __iter__ –

    Materialize one rank-local buffer at a time.

  • metadata –

    Return a fresh list while preserving the immutable manifest order.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
class ManifestM2NWeightSource(M2NWeightSource):
    """An ordered M2N source backed by lazy, rank-local manifest entries."""

    def __init__(self, entries: Sequence[M2NManifestEntry]) -> None:
        self._entries = tuple(entries)
        self._metadata = tuple(entry.metadata for entry in self._entries)

    def metadata(self) -> list[M2NParamMeta]:  # type: ignore[override]
        """Return a fresh list while preserving the immutable manifest order."""
        return list(self._metadata)

    def __iter__(self) -> Iterator[tuple[str, Any]]:  # type: ignore[override]
        """Materialize one rank-local buffer at a time."""
        for entry in self._entries:
            yield entry.key, entry.local_tensor()

__iter__()

Materialize one rank-local buffer at a time.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
112
113
114
115
def __iter__(self) -> Iterator[tuple[str, Any]]:  # type: ignore[override]
    """Materialize one rank-local buffer at a time."""
    for entry in self._entries:
        yield entry.key, entry.local_tensor()

metadata()

Return a fresh list while preserving the immutable manifest order.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
108
109
110
def metadata(self) -> list[M2NParamMeta]:  # type: ignore[override]
    """Return a fresh list while preserving the immutable manifest order."""
    return list(self._metadata)

_placement_code(placement)

Map a torch.distributed placement onto an m2n placement code.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
118
119
120
121
122
123
124
125
126
127
128
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"
    )

mesh_from_tensor(tensor, num_trainer_ranks)

The mesh a parameter lives on, as an M2NMesh.

Tensors without DeviceMesh metadata are treated as replicated. Implicitly sharded tensors require a custom source with an explicit layout.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def mesh_from_tensor(tensor: torch.Tensor, num_trainer_ranks: int) -> M2NMesh:
    """The mesh a parameter lives on, as an `M2NMesh`.

    Tensors without DeviceMesh metadata are treated as replicated. Implicitly
    sharded tensors require a custom source with an explicit layout.
    """
    device_mesh = getattr(tensor, "device_mesh", None)
    if device_mesh is None:
        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)

placements_from_tensor(tensor)

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.

Source code in vllm/distributed/weight_transfer/m2n_source.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
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])