| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Trainer-side weight sources for the NCCL M2N backend. |
| 4 | |
| 5 | m2n plans a transfer from both sides' layouts, so the trainer has to say how it |
| 6 | holds each parameter — which the base `WeightSource` / `ParamMeta` pair does not |
| 7 | express. This module supplies the m2n flavor of both, and the DTensor-backed |
| 8 | implementation that covers the common trainer. |
| 9 | |
| 10 | A source declares its mesh once (`mesh()`) and a placement per parameter, since |
| 11 | one topology describes every tensor on the side. |
| 12 | """ |
| 13 | |
| 14 | from collections.abc import Iterator |
| 15 | from typing import Any |
| 16 | |
| 17 | import torch |
| 18 | |
| 19 | from vllm.distributed.weight_transfer.base import ParamMeta, WeightSource |
| 20 | from vllm.distributed.weight_transfer.m2n_common import ( |
| 21 | REPLICATE, |
| 22 | REPLICATED, |
| 23 | M2NMesh, |
| 24 | M2NParamMeta, |
| 25 | Placements, |
| 26 | ) |
| 27 | |
| 28 | __all__ = [ |
| 29 | "M2NWeightSource", |
| 30 | "DTensorModuleSource", |
| 31 | "mesh_from_tensor", |
| 32 | "placements_from_tensor", |
| 33 | ] |
| 34 | |
| 35 | |
| 36 | class M2NWeightSource(WeightSource): |
| 37 | """A `WeightSource` that also describes how the trainer holds its weights. |
| 38 | |
| 39 | Unlike `ModuleSource`, iteration yields each rank's **local shard**, not a |
| 40 | materialized full tensor: gathering is exactly the cost m2n removes. |
| 41 | """ |
| 42 | |
| 43 | def mesh(self) -> M2NMesh: |
| 44 | """The trainer's rank topology, shared by every parameter.""" |
| 45 | raise NotImplementedError |
| 46 | |
| 47 | def metadata(self) -> list[ParamMeta]: |
| 48 | """Name, dtype, full shape, and trainer placement for each parameter.""" |
| 49 | raise NotImplementedError |
| 50 | |
| 51 | def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: |
| 52 | """Yield `(name, local shard)` pairs in the same order as `metadata()`.""" |
| 53 | raise NotImplementedError |
| 54 | |
| 55 | |
| 56 | def _placement_code(placement: Any) -> int: |
| 57 | """Map a `torch.distributed` placement onto an m2n placement code.""" |
| 58 | name = type(placement).__name__ |
| 59 | if name == "Replicate": |
| 60 | return REPLICATE |
| 61 | if name == "Shard": |
| 62 | return int(placement.dim) |
| 63 | raise ValueError( |
| 64 | f"nccl_m2n cannot express the {name} placement; only Replicate and " |
| 65 | "Shard are supported" |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | def mesh_from_tensor(tensor: torch.Tensor, num_trainer_ranks: int) -> M2NMesh: |
| 70 | """The mesh a parameter lives on, as an `M2NMesh`. |
| 71 | |
| 72 | A plain tensor is identical on every trainer rank, so it spans all of them. |
| 73 | """ |
| 74 | device_mesh = getattr(tensor, "device_mesh", None) |
| 75 | if device_mesh is None: |
| 76 | return M2NMesh((num_trainer_ranks, 1), 0) |
| 77 | |
| 78 | grid = device_mesh.mesh |
| 79 | ranks = grid.flatten().tolist() |
| 80 | if ranks != list(range(len(ranks))): |
| 81 | raise ValueError( |
| 82 | "nccl_m2n requires the trainer's device mesh to cover the " |
| 83 | f"contiguous rank interval [0, {len(ranks)}); got {ranks}" |
| 84 | ) |
| 85 | if grid.ndim > 2: |
| 86 | raise ValueError( |
| 87 | f"nccl_m2n supports 1-D and 2-D device meshes, got {grid.ndim}-D" |
| 88 | ) |
| 89 | dims = tuple(grid.shape) |
| 90 | return M2NMesh(dims if len(dims) == 2 else (dims[0], 1), 0) |
| 91 | |
| 92 | |
| 93 | def placements_from_tensor(tensor: torch.Tensor) -> Placements | None: |
| 94 | """How a parameter is placed over its mesh, or `REPLICATED`. |
| 95 | |
| 96 | `REPLICATED` is returned rather than a `(REPLICATE, REPLICATE)` pair: m2n |
| 97 | cannot express that, and the size-1-axis encoding it needs instead is |
| 98 | applied later by `resolve_layout`, which owns that workaround. |
| 99 | """ |
| 100 | placements = getattr(tensor, "placements", None) |
| 101 | if placements is None: |
| 102 | return REPLICATED |
| 103 | codes = [_placement_code(p) for p in placements] |
| 104 | if all(code == REPLICATE for code in codes): |
| 105 | return REPLICATED |
| 106 | if len(codes) == 1: |
| 107 | return (codes[0], REPLICATE) |
| 108 | return (codes[0], codes[1]) |
| 109 | |
| 110 | |
| 111 | class DTensorModuleSource(M2NWeightSource): |
| 112 | """`M2NWeightSource` over `module.named_parameters()`. |
| 113 | |
| 114 | Covers both the FSDP/DTensor trainer (placement read off each parameter, |
| 115 | local shard yielded via `to_local()`) and the plain replicated trainer, with |
| 116 | no special casing. Trainers with a custom producer (a Megatron export, MoE |
| 117 | re-fusing) subclass `M2NWeightSource` instead. |
| 118 | """ |
| 119 | |
| 120 | def __init__(self, module: torch.nn.Module, num_trainer_ranks: int) -> None: |
| 121 | """Wrap `module`; `num_trainer_ranks` sizes the mesh for plain tensors.""" |
| 122 | self._module = module |
| 123 | self._num_trainer_ranks = num_trainer_ranks |
| 124 | |
| 125 | def mesh(self) -> M2NMesh: |
| 126 | """The mesh every sharded parameter agrees on. |
| 127 | |
| 128 | Replicated parameters span the trainer without a device mesh of their |
| 129 | own, so they never decide this; a model whose *sharded* parameters |
| 130 | disagree cannot be described by one side mesh and is rejected. |
| 131 | """ |
| 132 | meshes = { |
| 133 | mesh_from_tensor(p, self._num_trainer_ranks) |
| 134 | for _, p in self._module.named_parameters() |
| 135 | if placements_from_tensor(p) is not REPLICATED |
| 136 | } |
| 137 | if len(meshes) > 1: |
| 138 | raise ValueError( |
| 139 | "nccl_m2n needs one mesh per side, but this module's sharded " |
| 140 | f"parameters span several: {sorted(m.dims for m in meshes)}" |
| 141 | ) |
| 142 | return meshes.pop() if meshes else M2NMesh((self._num_trainer_ranks, 1), 0) |
| 143 | |
| 144 | def metadata(self) -> list[ParamMeta]: |
| 145 | """Read global shape/dtype and placements without gathering shards.""" |
| 146 | return [ |
| 147 | M2NParamMeta(name, p.dtype, tuple(p.shape), placements_from_tensor(p)) |
| 148 | for name, p in self._module.named_parameters() |
| 149 | ] |
| 150 | |
| 151 | def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: |
| 152 | """Yield each parameter's local shard (`to_local()`), or the tensor itself.""" |
| 153 | for name, param in self._module.named_parameters(): |
| 154 | to_local = getattr(param, "to_local", None) |
| 155 | yield name, (to_local() if callable(to_local) else param) |