| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-side weight transfer engine built on NCCL M2N (`nccl_m2n`). |
| 4 | |
| 5 | The trainer and the inference workers share one NCCL communicator: trainer ranks |
| 6 | occupy `[0, T)`, workers `[T, T + N)`. Each parameter is moved with a single |
| 7 | `nccl.m2n.reshard`, which redistributes it from the trainer's layout (FSDP / EP |
| 8 | / arbitrary DTensor sharding) to the inference layout — so the trainer sends its |
| 9 | local shards and never all-gathers a full tensor, which is what the broadcast |
| 10 | NCCL backend forces it to do. |
| 11 | |
| 12 | When an incoming checkpoint parameter maps directly to a live vLLM parameter, |
| 13 | M2N reshards into that worker's local model storage. Parameters whose names or |
| 14 | layouts cannot be resolved receive a full tensor and fall back to |
| 15 | `load_weights`, preserving the existing backend's loading behavior. |
| 16 | |
| 17 | Both meshes participate in every reshard, in the same order, so this backend has |
| 18 | the same concurrency shape as the broadcast NCCL backend: the worker must be |
| 19 | inside `receive_weights` while the trainer is sending. Driving that from the |
| 20 | trainer is the trainer engine's job. |
| 21 | """ |
| 22 | |
| 23 | from collections.abc import Iterator |
| 24 | from dataclasses import dataclass, field |
| 25 | from typing import TYPE_CHECKING, Any, cast |
| 26 | |
| 27 | import torch |
| 28 | |
| 29 | from vllm.config.weight_transfer import WeightTransferConfig |
| 30 | from vllm.distributed.weight_transfer.base import ( |
| 31 | WeightTransferEngine, |
| 32 | WeightTransferInitInfo, |
| 33 | WeightTransferUpdateInfo, |
| 34 | ) |
| 35 | from vllm.distributed.weight_transfer.m2n_common import ( |
| 36 | DESTINATION_SHARD_AXIS, |
| 37 | MESH_NDIMS, |
| 38 | REPLICATED, |
| 39 | M2NMesh, |
| 40 | M2NParamMeta, |
| 41 | Placements, |
| 42 | check_plan_limits, |
| 43 | check_transferable, |
| 44 | comm_ptr, |
| 45 | import_m2n, |
| 46 | publish_destination_placements, |
| 47 | resolve_layout, |
| 48 | to_mesh, |
| 49 | to_placements, |
| 50 | validate_layout, |
| 51 | ) |
| 52 | from vllm.distributed.weight_transfer.m2n_layout import ( |
| 53 | M2NDestination, |
| 54 | resolve_parameter_destinations, |
| 55 | ) |
| 56 | from vllm.distributed.weight_transfer.nccl_common import ( |
| 57 | NCCLWeightTransferInitInfo, |
| 58 | decode_nccl_unique_id, |
| 59 | worker_init_process_group, |
| 60 | ) |
| 61 | |
| 62 | if TYPE_CHECKING: |
| 63 | from vllm.config import VllmConfig |
| 64 | from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator |
| 65 | |
| 66 | __all__ = [ |
| 67 | "M2NWeightTransferInitInfo", |
| 68 | "M2NWeightTransferUpdateInfo", |
| 69 | "M2NWeightTransferEngine", |
| 70 | ] |
| 71 | |
| 72 | |
| 73 | # --------------------------------------------------------------------------- |
| 74 | # Wire types |
| 75 | # --------------------------------------------------------------------------- |
| 76 | |
| 77 | |
| 78 | @dataclass(kw_only=True) |
| 79 | class M2NWeightTransferInitInfo(WeightTransferInitInfo): |
| 80 | """Worker-side init info: the rendezvous plus the full transfer plan. |
| 81 | |
| 82 | Layouts are static for the whole run, so they ride the one-time init |
| 83 | handshake and the per-round update info stays a list of names. Everything is |
| 84 | plain JSON so the HTTP control plane carries it unchanged. |
| 85 | """ |
| 86 | |
| 87 | master_address: str | None = None |
| 88 | master_port: int | None = None |
| 89 | rank_offset: int |
| 90 | """First worker rank, i.e. the number of trainer ranks.""" |
| 91 | world_size: int |
| 92 | """Trainer ranks + all inference workers.""" |
| 93 | src_mesh_dims: list[int] |
| 94 | """The trainer's mesh, shared by every parameter (`start_rank` is 0: the |
| 95 | trainer occupies the front of the communicator).""" |
| 96 | dst_mesh_dims: list[int] |
| 97 | """The inference mesh, starting at `rank_offset`. Declared rather than |
| 98 | derived so both sides describe the destination identically.""" |
| 99 | names: list[str] |
| 100 | dtype_names: list[str] |
| 101 | shapes: list[list[int]] |
| 102 | src_placements: list[list[int] | None] |
| 103 | """Per parameter, relative to `src_mesh_dims`; `None` means replicated.""" |
| 104 | nccl_unique_id_b64: str | None = field(default=None, repr=False) |
| 105 | max_cta: int | None = None |
| 106 | |
| 107 | def __post_init__(self) -> None: |
| 108 | _ = self.nccl_unique_id_bytes |
| 109 | num_params = len(self.names) |
| 110 | for label, values in ( |
| 111 | ("dtype_names", self.dtype_names), |
| 112 | ("shapes", self.shapes), |
| 113 | ("src_placements", self.src_placements), |
| 114 | ): |
| 115 | if len(values) != num_params: |
| 116 | raise ValueError( |
| 117 | f"`{label}` should be of the same size as `names`: " |
| 118 | f"got {len(values)} and {num_params}" |
| 119 | ) |
| 120 | if self.rank_offset < 1 or self.rank_offset >= self.world_size: |
| 121 | raise ValueError( |
| 122 | f"`rank_offset` ({self.rank_offset}) must leave at least one " |
| 123 | f"trainer rank and one worker in world_size {self.world_size}" |
| 124 | ) |
| 125 | num_workers = self.world_size - self.rank_offset |
| 126 | dst = tuple(self.dst_mesh_dims) |
| 127 | if len(dst) != MESH_NDIMS or dst[0] * dst[1] != num_workers: |
| 128 | raise ValueError( |
| 129 | f"`dst_mesh_dims` {dst} must be {MESH_NDIMS} dims covering the " |
| 130 | f"{num_workers} inference workers" |
| 131 | ) |
| 132 | |
| 133 | @property |
| 134 | def nccl_unique_id_bytes(self) -> bytes | None: |
| 135 | return decode_nccl_unique_id( |
| 136 | master_address=self.master_address, |
| 137 | master_port=self.master_port, |
| 138 | nccl_unique_id_b64=self.nccl_unique_id_b64, |
| 139 | ctx="M2NWeightTransferInitInfo", |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | @dataclass |
| 144 | class M2NWeightTransferUpdateInfo(WeightTransferUpdateInfo): |
| 145 | """Per-round update info: which parameters this chunk carries, in order. |
| 146 | |
| 147 | Shapes, dtypes and layouts were fixed at init, so a round only needs to say |
| 148 | what is coming and in what order — both sides must issue their reshards in |
| 149 | exactly that order. |
| 150 | """ |
| 151 | |
| 152 | names: list[str] |
| 153 | |
| 154 | |
| 155 | # --------------------------------------------------------------------------- |
| 156 | # Worker engine |
| 157 | # --------------------------------------------------------------------------- |
| 158 | |
| 159 | |
| 160 | class M2NWeightTransferEngine( |
| 161 | WeightTransferEngine[M2NWeightTransferInitInfo, M2NWeightTransferUpdateInfo] |
| 162 | ): |
| 163 | """Inference-side engine: receives each parameter with one reshard. |
| 164 | |
| 165 | Resolvable parameters are resharded from the trainer layout directly into |
| 166 | each worker's live local model storage. Unresolvable parameters receive a |
| 167 | full tensor and use `load_weights`. In both cases, the trainer sends its |
| 168 | local shards without materializing a full tensor. |
| 169 | """ |
| 170 | |
| 171 | init_info_cls = M2NWeightTransferInitInfo |
| 172 | update_info_cls = M2NWeightTransferUpdateInfo |
| 173 | # The destination plan is resolved against the model this engine was built |
| 174 | # for, so it cannot be retargeted at a draft model mid-flight. |
| 175 | supports_draft_weight_update = False |
| 176 | |
| 177 | def __init__( |
| 178 | self, |
| 179 | config: WeightTransferConfig, |
| 180 | vllm_config: "VllmConfig", |
| 181 | device: torch.device, |
| 182 | model: torch.nn.Module, |
| 183 | ) -> None: |
| 184 | super().__init__(config, vllm_config, device, model) |
| 185 | self.model_update_group: PyNcclCommunicator | None = None |
| 186 | self._m2n: Any = None |
| 187 | self._handle: Any = None |
| 188 | self._metas: list[M2NParamMeta] = [] |
| 189 | self._src_mesh: M2NMesh | None = None |
| 190 | self._dst_mesh: M2NMesh | None = None |
| 191 | self._parameter_destinations: list[M2NDestination] = [] |
| 192 | self._index: dict[str, int] = {} |
| 193 | self._uses_load_weights = False |
| 194 | |
| 195 | def init_transfer_engine(self, init_info: M2NWeightTransferInitInfo) -> None: |
| 196 | """Join the trainer's communicator and rebuild the transfer plan. |
| 197 | |
| 198 | Every precondition (dtype, tensor rank, divisibility) is checked here so |
| 199 | a bad plan fails the init RPC instead of hanging the first collective. |
| 200 | """ |
| 201 | self._m2n = import_m2n() |
| 202 | self._metas = [] |
| 203 | |
| 204 | self._src_mesh = M2NMesh( |
| 205 | cast(tuple[int, int], tuple(init_info.src_mesh_dims)), 0 |
| 206 | ) |
| 207 | if self._src_mesh.size != init_info.rank_offset: |
| 208 | raise ValueError( |
| 209 | f"source mesh covers {self._src_mesh.size} ranks, but there are " |
| 210 | f"{init_info.rank_offset} trainer ranks" |
| 211 | ) |
| 212 | # The inference topology, as declared by the trainer. This is the mesh |
| 213 | # a sharded destination is placed over; a replicated one is described |
| 214 | # by the alternate descriptor `resolve_layout` derives from it, which |
| 215 | # covers the same ranks but is not this mesh. |
| 216 | self._dst_mesh = M2NMesh( |
| 217 | cast(tuple[int, int], tuple(init_info.dst_mesh_dims)), |
| 218 | init_info.rank_offset, |
| 219 | ) |
| 220 | |
| 221 | for name, name_dtype, shape, placements in zip( |
| 222 | init_info.names, |
| 223 | init_info.dtype_names, |
| 224 | init_info.shapes, |
| 225 | init_info.src_placements, |
| 226 | ): |
| 227 | dtype = getattr(torch, name_dtype, None) |
| 228 | if not isinstance(dtype, torch.dtype): |
| 229 | raise ValueError( |
| 230 | f"parameter '{name}' has invalid dtype name " |
| 231 | f"'{name_dtype}'; expected the name of a torch.dtype" |
| 232 | ) |
| 233 | check_transferable(name, dtype, shape) |
| 234 | codes = ( |
| 235 | REPLICATED |
| 236 | if placements is None |
| 237 | else cast(Placements, tuple(placements)) |
| 238 | ) |
| 239 | src_mesh, src_placements = resolve_layout( |
| 240 | self._src_mesh, codes, f"parameter '{name}' source placements" |
| 241 | ) |
| 242 | validate_layout(src_mesh, src_placements, shape, "source") |
| 243 | # Destination checks are deferred until loader-aware resolution; |
| 244 | # assuming replication here can reject a valid sharded plan. |
| 245 | self._metas.append(M2NParamMeta(name, dtype, tuple(shape), codes)) |
| 246 | |
| 247 | parallel_config = self.parallel_config |
| 248 | shard_axis_size = self._dst_mesh.dims[DESTINATION_SHARD_AXIS] |
| 249 | quantization = getattr(self.model_config, "quantization", None) |
| 250 | # A quantized weight loader may dequantize / transpose / requantize, and |
| 251 | # under PP a rank's local shape is not the checkpoint shape split over |
| 252 | # the shard axis. Both break shape-derived resolution, so take the |
| 253 | # fallback. |
| 254 | allow_direct = ( |
| 255 | quantization is None and parallel_config.pipeline_parallel_size == 1 |
| 256 | ) |
| 257 | # Match every incoming checkpoint parameter to its rank-local target. |
| 258 | # Resolvable parameters can be resharded directly into live model |
| 259 | # storage; the rest receive a full tensor for `load_weights`. |
| 260 | self._parameter_destinations = resolve_parameter_destinations( |
| 261 | self.model, |
| 262 | [m.name for m in self._metas], |
| 263 | [m.dtype for m in self._metas], |
| 264 | [m.shape for m in self._metas], |
| 265 | num_workers=self._dst_mesh.size, |
| 266 | shard_axis_size=shard_axis_size, |
| 267 | allow_direct=allow_direct, |
| 268 | ) |
| 269 | # Check the resolved layout rather than a provisional replicated one. |
| 270 | # For example, 32 source shards feeding 4 destination shards may need |
| 271 | # only 8 sources per destination, while a replicated plan needs all 32. |
| 272 | for meta, destination in zip(self._metas, self._parameter_destinations): |
| 273 | src_mesh, src_placements = resolve_layout( |
| 274 | self._src_mesh, |
| 275 | meta.placements, |
| 276 | f"parameter '{meta.name}' source placements", |
| 277 | ) |
| 278 | dst_mesh, dst_placements = resolve_layout( |
| 279 | self._dst_mesh, |
| 280 | destination.placements, |
| 281 | f"parameter '{meta.name}' destination placements", |
| 282 | ) |
| 283 | validate_layout(dst_mesh, dst_placements, meta.shape, "destination") |
| 284 | check_plan_limits( |
| 285 | (src_mesh, src_placements), |
| 286 | (dst_mesh, dst_placements), |
| 287 | meta.name, |
| 288 | ) |
| 289 | |
| 290 | # Update requests carry names only, so cache their plan indices. Track |
| 291 | # whether any fallback entry requires the `load_weights` lifecycle. |
| 292 | self._index = {meta.name: i for i, meta in enumerate(self._metas)} |
| 293 | self._uses_load_weights = any( |
| 294 | not destination.direct for destination in self._parameter_destinations |
| 295 | ) |
| 296 | |
| 297 | self.model_update_group = worker_init_process_group( |
| 298 | NCCLWeightTransferInitInfo( |
| 299 | master_address=init_info.master_address, |
| 300 | master_port=init_info.master_port, |
| 301 | nccl_unique_id_b64=init_info.nccl_unique_id_b64, |
| 302 | rank_offset=init_info.rank_offset, |
| 303 | world_size=init_info.world_size, |
| 304 | ), |
| 305 | parallel_config, |
| 306 | ) |
| 307 | # Destinations are per-parameter and depend on the inference model, so |
| 308 | # the trainer cannot derive them. The first worker publishes the plan |
| 309 | # and every other rank checks its own against it — a silent |
| 310 | # disagreement between workers would mean mismatched collectives. |
| 311 | mine = [destination.placements for destination in self._parameter_destinations] |
| 312 | agreed = publish_destination_placements( |
| 313 | self.model_update_group, init_info.rank_offset, mine, len(self._metas) |
| 314 | ) |
| 315 | if agreed != mine: |
| 316 | mismatched = next( |
| 317 | meta.name for meta, a, b in zip(self._metas, mine, agreed) if a != b |
| 318 | ) |
| 319 | raise RuntimeError( |
| 320 | "inference workers resolved different nccl_m2n destinations " |
| 321 | f"(first mismatch: '{mismatched}'); all workers must run the " |
| 322 | "same model and parallel config" |
| 323 | ) |
| 324 | |
| 325 | self._handle = self._m2n.Handle.create( |
| 326 | self._m2n.Config(max_cta=init_info.max_cta) |
| 327 | ) |
| 328 | |
| 329 | def start_weight_update(self) -> None: |
| 330 | """Set up layerwise reloading, but only if some parameter needs it. |
| 331 | |
| 332 | Directly-resharded parameters are written in place and never go through |
| 333 | `load_weights`, so a plan with no fallback entries has nothing to |
| 334 | reload. |
| 335 | """ |
| 336 | if not self._uses_load_weights: |
| 337 | return |
| 338 | from vllm.model_executor.model_loader.reload import initialize_layerwise_reload |
| 339 | |
| 340 | initialize_layerwise_reload(self.model) |
| 341 | |
| 342 | def finish_weight_update(self) -> None: |
| 343 | """Finalize layerwise reloading when the plan uses fallback entries.""" |
| 344 | if not self._uses_load_weights: |
| 345 | return |
| 346 | from vllm.model_executor.model_loader.reload import finalize_layerwise_reload |
| 347 | |
| 348 | finalize_layerwise_reload(self.model, self.model_config) |
| 349 | |
| 350 | def receive_weights(self, update_info: M2NWeightTransferUpdateInfo) -> None: |
| 351 | """Receive each requested parameter using its initialization-time plan.""" |
| 352 | if self._handle is None or self.model_update_group is None: |
| 353 | raise RuntimeError( |
| 354 | "nccl_m2n weight transfer not initialized. " |
| 355 | "Call init_transfer_engine() first." |
| 356 | ) |
| 357 | |
| 358 | requested = [] |
| 359 | for name in update_info.names: |
| 360 | index = self._index.get(name) |
| 361 | if index is None: |
| 362 | raise ValueError( |
| 363 | f"parameter '{name}' was not declared at init; the " |
| 364 | "trainer must send the same parameter set it announced" |
| 365 | ) |
| 366 | requested.append((name, index)) |
| 367 | |
| 368 | from vllm.model_executor.model_loader.mtp_validation import ( |
| 369 | disable_mtp_completeness_check, |
| 370 | ) |
| 371 | |
| 372 | comm = comm_ptr(self.model_update_group) |
| 373 | stream = torch.cuda.current_stream() |
| 374 | |
| 375 | # Reshard lazily so one `load_weights` invocation sees the complete |
| 376 | # fallback sequence without staging every full tensor at once. Direct |
| 377 | # destinations still execute in request order as the iterator advances. |
| 378 | def reshard_requested_weights() -> Iterator[tuple[str, torch.Tensor]]: |
| 379 | for name, index in requested: |
| 380 | meta = self._metas[index] |
| 381 | destination = self._parameter_destinations[index] |
| 382 | buffer = destination.tensor |
| 383 | if buffer is None: |
| 384 | buffer = torch.empty( |
| 385 | meta.shape, dtype=meta.dtype, device=self.device |
| 386 | ) |
| 387 | |
| 388 | self._reshard(comm, stream, meta, destination.placements, buffer) |
| 389 | |
| 390 | if destination.tensor is None: |
| 391 | # `load_weights` reads on the host stream, so the transfer |
| 392 | # has to have landed before it runs. |
| 393 | stream.synchronize() |
| 394 | yield name, buffer |
| 395 | |
| 396 | has_fallback = any( |
| 397 | self._parameter_destinations[index].tensor is None for _, index in requested |
| 398 | ) |
| 399 | with disable_mtp_completeness_check(): |
| 400 | received_weights = reshard_requested_weights() |
| 401 | if has_fallback: |
| 402 | # Some model loaders keep invocation-local state to combine |
| 403 | # related checkpoint tensors, so preserve one iterable scope. |
| 404 | self.model.load_weights(received_weights) |
| 405 | else: |
| 406 | # The iterator performs the reshards. With nothing to yield to |
| 407 | # `load_weights`, exhaust it here to execute direct transfers. |
| 408 | for _ in received_weights: |
| 409 | pass |
| 410 | |
| 411 | def _reshard( |
| 412 | self, |
| 413 | comm: int, |
| 414 | stream: torch.cuda.Stream, |
| 415 | src_meta: M2NParamMeta, |
| 416 | dst_placements: Placements | None, |
| 417 | dst_buffer: torch.Tensor, |
| 418 | ) -> None: |
| 419 | """Reshard one parameter from its trainer layout into a worker buffer. |
| 420 | |
| 421 | `dst_placements` describes the buffer's logical placement over the |
| 422 | worker mesh; `dst_buffer` is this rank's physical storage. |
| 423 | """ |
| 424 | assert self._src_mesh is not None |
| 425 | assert self._dst_mesh is not None |
| 426 | m2n = self._m2n |
| 427 | src_mesh, resolved_src_placements = resolve_layout( |
| 428 | self._src_mesh, src_meta.placements |
| 429 | ) |
| 430 | dst_mesh, resolved_dst_placements = resolve_layout( |
| 431 | self._dst_mesh, dst_placements |
| 432 | ) |
| 433 | m2n.reshard( |
| 434 | None, |
| 435 | dst_buffer, |
| 436 | comm, |
| 437 | stream, |
| 438 | src_mesh=to_mesh(m2n, src_mesh), |
| 439 | src_placements=to_placements(m2n, resolved_src_placements), |
| 440 | dst_mesh=to_mesh(m2n, dst_mesh), |
| 441 | dst_placements=to_placements(m2n, resolved_dst_placements), |
| 442 | handle=self._handle, |
| 443 | ) |
| 444 | |
| 445 | def shutdown(self) -> None: |
| 446 | """Finish pending GPU work and release M2N communication state.""" |
| 447 | if self._handle is not None: |
| 448 | # M2N does not synchronize caller streams on finalize. |
| 449 | torch.accelerator.synchronize() |
| 450 | self._handle.destroy() |
| 451 | self._handle = None |
| 452 | self.model_update_group = None |
| 453 | self._parameter_destinations = [] |