| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared helpers for the NCCL M2N (`nccl_m2n`) weight transfer backend. |
| 4 | |
| 5 | M2N reshards a tensor between two disjoint meshes of ranks that live in one |
| 6 | communicator: the trainer occupies ranks `[0, T)` and the inference workers |
| 7 | `[T, T + N)`. This module holds everything both sides need — the optional |
| 8 | runtime import, layout descriptors, and the conversion to `nccl.m2n` types — |
| 9 | so the engine module stays about the transfer itself. |
| 10 | """ |
| 11 | |
| 12 | from collections.abc import Sequence |
| 13 | from dataclasses import dataclass |
| 14 | from typing import TYPE_CHECKING, Any, cast |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | from vllm.distributed.weight_transfer.base import ParamMeta |
| 19 | |
| 20 | if TYPE_CHECKING: |
| 21 | from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator |
| 22 | |
| 23 | # Placement code for a mesh axis that replicates. Any other (non-negative) code |
| 24 | # is the tensor dim that axis shards. Ints rather than `nccl.m2n` objects so |
| 25 | # layouts survive the JSON init handshake without a custom encoder. |
| 26 | REPLICATE = -1 |
| 27 | |
| 28 | # `NCCL_RESHARD_MAX_TENSOR_DIMS`; 4-D and higher are rejected by the library. |
| 29 | MAX_TENSOR_DIMS = 3 |
| 30 | |
| 31 | MESH_NDIMS = 2 |
| 32 | |
| 33 | # Destination mesh convention. Keeping these roles named avoids scattering |
| 34 | # positional axis assumptions through the worker-side planner. |
| 35 | DESTINATION_REPLICA_AXIS = 0 |
| 36 | DESTINATION_SHARD_AXIS = 1 |
| 37 | |
| 38 | # m2n bounds how many source shards may feed one destination shard, and how many |
| 39 | # destination shards one source shard may feed, with compile-time arrays in |
| 40 | # `reshard_limits.h`. The bindings do not expose them, so they are mirrored |
| 41 | # here; a build with larger arrays makes these conservative, never wrong. |
| 42 | MAX_SOURCE_SHARDS = 16 |
| 43 | MAX_DEST_SHARDS = 64 |
| 44 | |
| 45 | # Wire dtypes `ncclReshard` accepts. Notably excludes fp4 and any packed |
| 46 | # sub-byte type, so quantized checkpoints are out of scope for now. |
| 47 | SUPPORTED_DTYPES: frozenset[torch.dtype] = frozenset( |
| 48 | dtype |
| 49 | for dtype in ( |
| 50 | getattr(torch, name, None) |
| 51 | for name in ( |
| 52 | "int8", |
| 53 | "uint8", |
| 54 | "float8_e4m3fn", |
| 55 | "float8_e5m2", |
| 56 | "float16", |
| 57 | "bfloat16", |
| 58 | "int32", |
| 59 | "uint32", |
| 60 | "float32", |
| 61 | "int64", |
| 62 | "uint64", |
| 63 | "float64", |
| 64 | ) |
| 65 | ) |
| 66 | if dtype is not None |
| 67 | ) |
| 68 | |
| 69 | _IMPORT_HINT = ( |
| 70 | "The nccl_m2n weight transfer backend requires the `nccl-extensions` " |
| 71 | "package (and its `nccl4py` dependency), which is not installed by vLLM. " |
| 72 | "See https://github.com/NVIDIA/nccl-extensions for build " |
| 73 | "and install instructions. It needs NCCL 2.30.5 or newer, and " |
| 74 | "VLLM_NCCL_SO_PATH must point at the same libnccl.so that libnccl_m2n.so " |
| 75 | "was linked against." |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def import_m2n() -> Any: |
| 80 | """Import `nccl.m2n` lazily, with an actionable error when it is missing. |
| 81 | |
| 82 | Deferred so that importing vLLM — or any other weight transfer backend — |
| 83 | never requires the m2n runtime to be present. |
| 84 | """ |
| 85 | try: |
| 86 | import nccl.m2n as m2n |
| 87 | except ImportError as e: |
| 88 | raise ImportError(f"{_IMPORT_HINT} (import failed: {e})") from e |
| 89 | return m2n |
| 90 | |
| 91 | |
| 92 | @dataclass(frozen=True) |
| 93 | class M2NMesh: |
| 94 | """One side's rank topology — pure topology, no tensor placement. |
| 95 | |
| 96 | Mirrors `ncclMesh_t`: a 2-axis mesh owning the contiguous rank interval |
| 97 | `[start_rank, start_rank + dims[0] * dims[1])`. There is no 1-D mesh; a |
| 98 | single-axis topology is spelled with a second axis of size 1. |
| 99 | |
| 100 | One mesh describes every tensor on its side, so it is exchanged once at the |
| 101 | init handshake rather than per parameter. |
| 102 | """ |
| 103 | |
| 104 | dims: tuple[int, int] |
| 105 | start_rank: int |
| 106 | |
| 107 | def __post_init__(self) -> None: |
| 108 | if len(self.dims) != MESH_NDIMS or any(d <= 0 for d in self.dims): |
| 109 | raise ValueError(f"mesh dims must be {MESH_NDIMS} positive ints") |
| 110 | if self.start_rank < 0: |
| 111 | raise ValueError( |
| 112 | f"mesh start_rank must be non-negative, got {self.start_rank}" |
| 113 | ) |
| 114 | |
| 115 | @property |
| 116 | def size(self) -> int: |
| 117 | return self.dims[0] * self.dims[1] |
| 118 | |
| 119 | |
| 120 | # A tensor's placement over its side's mesh: one code per mesh axis |
| 121 | # (`REPLICATE`, or the tensor dim that axis shards), or `REPLICATED` for a |
| 122 | # tensor every rank holds in full. |
| 123 | Placements = tuple[int, int] |
| 124 | REPLICATED: Placements | None = None |
| 125 | |
| 126 | |
| 127 | def check_placements(placements: Placements, context: str = "placements") -> None: |
| 128 | """Reject placement pairs m2n cannot express. |
| 129 | |
| 130 | The header requires exactly one SHARD axis and one REPLICATE axis. |
| 131 | `{REPLICATE, REPLICATE}` hits a degenerate prepare branch, which is why |
| 132 | full replication is carried as `REPLICATED` and resolved separately; |
| 133 | sharding both axes is not expressible at all. |
| 134 | """ |
| 135 | if len(placements) != MESH_NDIMS: |
| 136 | raise ValueError(f"{context} must have {MESH_NDIMS} entries") |
| 137 | invalid = [code for code in placements if code < REPLICATE] |
| 138 | if invalid: |
| 139 | raise ValueError( |
| 140 | f"{context} contains invalid placement code {invalid[0]}; valid " |
| 141 | f"codes are {REPLICATE} (Replicate) and non-negative tensor dimensions" |
| 142 | ) |
| 143 | num_sharded = sum(code != REPLICATE for code in placements) |
| 144 | if num_sharded == 0: |
| 145 | raise ValueError( |
| 146 | f"{context}: a fully replicated tensor is carried as REPLICATED, " |
| 147 | "not as two " |
| 148 | "REPLICATE axes" |
| 149 | ) |
| 150 | if num_sharded == MESH_NDIMS: |
| 151 | raise ValueError( |
| 152 | f"{context}: nccl_m2n needs one REPLICATE mesh axis, but " |
| 153 | f"{placements} shards both" |
| 154 | ) |
| 155 | |
| 156 | |
| 157 | def resolve_layout( |
| 158 | mesh: M2NMesh, |
| 159 | placements: Placements | None, |
| 160 | context: str = "placements", |
| 161 | ) -> tuple[M2NMesh, Placements]: |
| 162 | """Pair one tensor's placement with the mesh m2n should see for it. |
| 163 | |
| 164 | A replicated tensor needs a size-1 mesh axis to carry a no-op shard, since |
| 165 | m2n has no `{REPLICATE, REPLICATE}`. It is therefore described over the |
| 166 | *same rank interval* re-factored as `(size, 1)`. That re-factoring is sound |
| 167 | precisely because replication is order-independent — every rank holds the |
| 168 | whole tensor, so it does not matter that `(a, b)` and `(size, 1)` walk the |
| 169 | interval in a different order. A sharded tensor keeps its side's own |
| 170 | factorization, where rank order decides who owns which shard. |
| 171 | """ |
| 172 | if placements is None: |
| 173 | return M2NMesh((mesh.size, 1), mesh.start_rank), (REPLICATE, 0) |
| 174 | check_placements(placements, context) |
| 175 | return mesh, placements |
| 176 | |
| 177 | |
| 178 | def shard_count(mesh: M2NMesh, placements: Placements) -> int: |
| 179 | """How many pieces this layout splits the tensor into.""" |
| 180 | return next( |
| 181 | (mesh.dims[axis] for axis, code in enumerate(placements) if code != REPLICATE), |
| 182 | 1, |
| 183 | ) |
| 184 | |
| 185 | |
| 186 | def check_plan_limits( |
| 187 | src: tuple[M2NMesh, Placements], |
| 188 | dst: tuple[M2NMesh, Placements], |
| 189 | name: str, |
| 190 | ) -> None: |
| 191 | """Reject plans that exceed m2n's static per-shard fan-in / fan-out arrays. |
| 192 | |
| 193 | Only the unambiguous cases are checked: when one side is a single shard it |
| 194 | is fed by (or feeds) every shard on the other side, so the count is exactly |
| 195 | the other side's shard count. In the general sharded-to-sharded case the |
| 196 | overlap depends on the library's chunking, and reproducing that arithmetic |
| 197 | here would duplicate internals that can change under us. |
| 198 | |
| 199 | m2n re-checks authoritatively and, because the plan is derived from the |
| 200 | shared descriptors, fails identically on every rank -- so this is about |
| 201 | reporting at init with a message that names the parameter, not about |
| 202 | avoiding a hang. |
| 203 | """ |
| 204 | src_shards = shard_count(*src) |
| 205 | dst_shards = shard_count(*dst) |
| 206 | if dst_shards == 1 and src_shards > MAX_SOURCE_SHARDS: |
| 207 | raise ValueError( |
| 208 | f"parameter '{name}' would feed {src_shards} source shards into one " |
| 209 | f"replicated destination, over m2n's MAX_SOURCES={MAX_SOURCE_SHARDS}. " |
| 210 | "Shard the destination, replicate the source, or rebuild m2n with " |
| 211 | "larger arrays." |
| 212 | ) |
| 213 | if src_shards == 1 and dst_shards > MAX_DEST_SHARDS: |
| 214 | raise ValueError( |
| 215 | f"parameter '{name}' would fan one source shard out to {dst_shards} " |
| 216 | f"destination shards, over m2n's MAX_TARGETS={MAX_DEST_SHARDS}. " |
| 217 | "Rebuild m2n with larger arrays to raise the bound." |
| 218 | ) |
| 219 | |
| 220 | |
| 221 | def validate_layout( |
| 222 | mesh: M2NMesh, placements: Placements, shape: Sequence[int], side: str |
| 223 | ) -> None: |
| 224 | """Check a resolved layout can describe `shape` before any collective runs. |
| 225 | |
| 226 | Called on both sides at init so a bad plan surfaces as an error from the |
| 227 | init RPC rather than as a hang inside the first reshard. |
| 228 | """ |
| 229 | for axis, code in enumerate(placements): |
| 230 | if code == REPLICATE: |
| 231 | continue |
| 232 | if code >= len(shape): |
| 233 | raise ValueError( |
| 234 | f"{side} layout shards tensor dim {code}, but the tensor " |
| 235 | f"has rank {len(shape)}" |
| 236 | ) |
| 237 | factor = mesh.dims[axis] |
| 238 | if shape[code] % factor: |
| 239 | raise ValueError( |
| 240 | f"{side} layout shards dim {code} (size {shape[code]}) over " |
| 241 | f"{factor} ranks, which does not divide evenly" |
| 242 | ) |
| 243 | |
| 244 | |
| 245 | @dataclass(frozen=True) |
| 246 | class M2NParamMeta(ParamMeta): |
| 247 | """`ParamMeta` extended with how the trainer places this tensor. |
| 248 | |
| 249 | The base class carries only name / dtype / full shape, which is not enough |
| 250 | to plan a reshard. `placements` is relative to its side's `M2NMesh`, or |
| 251 | `REPLICATED` when every rank holds the whole tensor. |
| 252 | """ |
| 253 | |
| 254 | placements: Placements | None |
| 255 | |
| 256 | |
| 257 | def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> None: |
| 258 | """Reject tensors m2n cannot move at all, with the parameter named.""" |
| 259 | if dtype not in SUPPORTED_DTYPES: |
| 260 | raise ValueError( |
| 261 | f"parameter '{name}' has dtype {dtype}, which nccl_m2n does not " |
| 262 | f"support. Supported: {sorted(str(d) for d in SUPPORTED_DTYPES)}" |
| 263 | ) |
| 264 | if not 1 <= len(shape) <= MAX_TENSOR_DIMS: |
| 265 | raise ValueError( |
| 266 | f"parameter '{name}' has rank {len(shape)}; nccl_m2n supports " |
| 267 | f"rank 1..{MAX_TENSOR_DIMS}" |
| 268 | ) |
| 269 | |
| 270 | |
| 271 | def to_mesh(m2n: Any, mesh: M2NMesh) -> Any: |
| 272 | return m2n.Mesh(mesh.dims, start_rank=mesh.start_rank) |
| 273 | |
| 274 | |
| 275 | def to_placements(m2n: Any, placements: Placements) -> list[Any]: |
| 276 | return [ |
| 277 | m2n.Replicate() if code == REPLICATE else m2n.Shard(code) for code in placements |
| 278 | ] |
| 279 | |
| 280 | |
| 281 | def publish_destination_placements( |
| 282 | comm: "PyNcclCommunicator", |
| 283 | first_worker_rank: int, |
| 284 | placements: "Sequence[Placements | None] | None", |
| 285 | num_parameters: int, |
| 286 | ) -> list[Placements | None]: |
| 287 | """Share the worker-side destination plan with every rank in the group. |
| 288 | |
| 289 | The trainer must issue each reshard with the same destination the workers |
| 290 | use, but once destinations are per-parameter they depend on the inference |
| 291 | model, which only the workers can see. The first worker publishes them here, |
| 292 | over the shared NCCL communicator; trainer ranks pass `None` and receive |
| 293 | them, and the other workers pass their own so a disagreement is caught |
| 294 | rather than deadlocking later. |
| 295 | |
| 296 | Only placements travel: the destination *mesh* is still derived from the |
| 297 | rank split, identically on both sides. Use the NCCL communicator directly |
| 298 | because unique-id rendezvous deliberately has no bootstrap process group. |
| 299 | """ |
| 300 | replicated_sentinel = REPLICATE - 1 |
| 301 | encoded = torch.full( |
| 302 | (num_parameters, MESH_NDIMS), |
| 303 | replicated_sentinel, |
| 304 | dtype=torch.int8, |
| 305 | device=comm.device, |
| 306 | ) |
| 307 | if comm.rank == first_worker_rank: |
| 308 | if placements is None or len(placements) != num_parameters: |
| 309 | raise ValueError( |
| 310 | f"publishing rank needs {num_parameters} destination placements" |
| 311 | ) |
| 312 | rows = [ |
| 313 | [replicated_sentinel] * MESH_NDIMS if placement is REPLICATED else placement |
| 314 | for placement in placements |
| 315 | ] |
| 316 | if rows: |
| 317 | encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=comm.device)) |
| 318 | |
| 319 | comm.broadcast(encoded, src=first_worker_rank) |
| 320 | return [ |
| 321 | REPLICATED |
| 322 | if all(code == replicated_sentinel for code in row) |
| 323 | else cast(Placements, tuple(row)) |
| 324 | for row in encoded.tolist() |
| 325 | ] |
| 326 | |
| 327 | |
| 328 | def comm_ptr(comm: "PyNcclCommunicator") -> int: |
| 329 | """Raw `ncclComm_t` behind vLLM's `PyNcclCommunicator`. |
| 330 | |
| 331 | m2n links its own NCCL, so the handle only means anything if vLLM loaded |
| 332 | the same `libnccl.so` — set `VLLM_NCCL_SO_PATH` accordingly. |
| 333 | """ |
| 334 | handle = comm.comm |
| 335 | ptr = getattr(handle, "value", handle) |
| 336 | if not ptr: |
| 337 | raise RuntimeError("PyNcclCommunicator has no live NCCL communicator") |
| 338 | return int(ptr) |