| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Destination-layout resolution for the NCCL M2N weight transfer backend. |
| 4 | |
| 5 | For every incoming checkpoint parameter the worker needs a destination buffer |
| 6 | plus its placement over the inference mesh. Two outcomes are possible: |
| 7 | |
| 8 | * **direct** — when the model is not quantized, the parameter maps 1:1 onto a |
| 9 | live vLLM parameter and its loader is a small, known-safe copy or tensor- |
| 10 | parallel loader. The loader's declared input/output dimension determines the |
| 11 | placement; shapes are validation, not an inference mechanism. The reshard |
| 12 | writes straight into the live parameter, so each rank receives only its own |
| 13 | shard and nothing is copied afterwards. |
| 14 | * **fallback** — anything else. The reshard delivers the whole tensor to every |
| 15 | rank and `load_weights` does the sharding, exactly as the broadcast NCCL |
| 16 | backend does. Fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`) |
| 17 | take this path: the checkpoint name does not name a vLLM parameter, so there |
| 18 | is nothing to resolve against. |
| 19 | |
| 20 | Correctness never depends on a parameter resolving — the fallback is always |
| 21 | available and is the same path the existing backend uses. |
| 22 | """ |
| 23 | |
| 24 | import inspect |
| 25 | import math |
| 26 | from collections.abc import Sequence |
| 27 | from dataclasses import dataclass |
| 28 | from functools import partial |
| 29 | from typing import cast |
| 30 | |
| 31 | import torch |
| 32 | |
| 33 | from vllm.distributed.weight_transfer.m2n_common import ( |
| 34 | DESTINATION_REPLICA_AXIS, |
| 35 | DESTINATION_SHARD_AXIS, |
| 36 | MESH_NDIMS, |
| 37 | REPLICATE, |
| 38 | REPLICATED, |
| 39 | Placements, |
| 40 | ) |
| 41 | from vllm.logger import init_logger |
| 42 | |
| 43 | logger = init_logger(__name__) |
| 44 | |
| 45 | _DIRECT_WEIGHT_LOADING_BLOCKLIST = { |
| 46 | "vllm.model_executor.models.gpt2.GPT2Model", |
| 47 | } |
| 48 | |
| 49 | |
| 50 | def _module_type_name(module: torch.nn.Module) -> str: |
| 51 | cls = type(module) |
| 52 | return f"{cls.__module__}.{cls.__name__}" |
| 53 | |
| 54 | |
| 55 | def _destination_placements(shard_dim: int) -> Placements: |
| 56 | """Place one tensor shard on the configured destination mesh axis.""" |
| 57 | placements = [REPLICATE] * MESH_NDIMS |
| 58 | placements[DESTINATION_REPLICA_AXIS] = REPLICATE |
| 59 | placements[DESTINATION_SHARD_AXIS] = shard_dim |
| 60 | return cast(Placements, tuple(placements)) |
| 61 | |
| 62 | |
| 63 | @dataclass |
| 64 | class M2NDestination: |
| 65 | """Where one checkpoint parameter lands on this worker.""" |
| 66 | |
| 67 | name: str |
| 68 | placements: Placements | None |
| 69 | """Placement over the inference mesh, or `REPLICATED` for the fallback.""" |
| 70 | tensor: torch.Tensor | None |
| 71 | """The live parameter view to reshard into, or None for the fallback path |
| 72 | (the engine allocates a full replica per round and calls `load_weights`).""" |
| 73 | |
| 74 | @property |
| 75 | def direct(self) -> bool: |
| 76 | return self.tensor is not None |
| 77 | |
| 78 | |
| 79 | def _base_loader_shard_dim( |
| 80 | param: torch.nn.Parameter, |
| 81 | shard_axis_size: int, |
| 82 | ) -> int | None: |
| 83 | """Return a known-safe loader's declared shard dim, or reject it. |
| 84 | |
| 85 | Loader identity is deliberately fail-closed. A missing loader is not |
| 86 | treated as ``default_weight_loader`` because a model-level ``load_weights`` |
| 87 | implementation may transform the tensor before it reaches that default. |
| 88 | Wrappers and partials are also rejected: recognizing their wrapped callable |
| 89 | would ignore behavior added by the wrapper. |
| 90 | """ |
| 91 | loader = getattr(param, "weight_loader", None) |
| 92 | if loader is None or isinstance(loader, partial): |
| 93 | return None |
| 94 | |
| 95 | owner = loader.__self__ if inspect.ismethod(loader) else None |
| 96 | loader_fn = loader.__func__ if inspect.ismethod(loader) else loader |
| 97 | |
| 98 | # Lazy imports avoid making the weight-transfer registry import all model |
| 99 | # executor layers during normal vLLM startup. |
| 100 | from vllm.model_executor.layers.linear import ( |
| 101 | ColumnParallelLinear, |
| 102 | ReplicatedLinear, |
| 103 | RowParallelLinear, |
| 104 | ) |
| 105 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 106 | |
| 107 | replicated_loaders = { |
| 108 | default_weight_loader, |
| 109 | ReplicatedLinear.weight_loader, |
| 110 | } |
| 111 | column_loaders = { |
| 112 | ColumnParallelLinear.weight_loader, |
| 113 | ColumnParallelLinear.weight_loader_v2, |
| 114 | } |
| 115 | row_loaders = { |
| 116 | RowParallelLinear.weight_loader, |
| 117 | RowParallelLinear.weight_loader_v2, |
| 118 | } |
| 119 | |
| 120 | # These attributes signal packing, a transform, or pre-sharded checkpoint |
| 121 | # storage. Such semantics cannot be reproduced by a plain M2N placement. |
| 122 | # Check them before accepting replicated loaders too: model-level loading |
| 123 | # may transform a tensor before handing it to an otherwise plain copier. |
| 124 | metadata_attrs = ("packed_dim", "packed_factor", "marlin_tile_size") |
| 125 | flag_attrs = ("needs_scalar_to_array", "is_sharded_weight", "is_transposed") |
| 126 | if any(getattr(param, attr, None) is not None for attr in metadata_attrs): |
| 127 | return None |
| 128 | if any(bool(getattr(param, attr, False)) for attr in flag_attrs): |
| 129 | return None |
| 130 | |
| 131 | if loader_fn in replicated_loaders: |
| 132 | return REPLICATE |
| 133 | if loader_fn not in column_loaders | row_loaders: |
| 134 | return None |
| 135 | |
| 136 | if owner is None or getattr(owner, "tp_size", None) != shard_axis_size: |
| 137 | return None |
| 138 | |
| 139 | attr = "output_dim" if loader_fn in column_loaders else "input_dim" |
| 140 | dim = getattr(param, attr, None) |
| 141 | return dim if isinstance(dim, int) and dim >= 0 else None |
| 142 | |
| 143 | |
| 144 | def _validated_shard_dim( |
| 145 | global_shape: Sequence[int], |
| 146 | local_shape: Sequence[int], |
| 147 | declared_dim: int, |
| 148 | shard_axis_size: int, |
| 149 | ) -> int | None: |
| 150 | """Validate a loader-declared shard dimension against both shapes.""" |
| 151 | if len(global_shape) != len(local_shape): |
| 152 | return None |
| 153 | |
| 154 | if declared_dim == REPLICATE: |
| 155 | return REPLICATE if tuple(global_shape) == tuple(local_shape) else None |
| 156 | if declared_dim >= len(global_shape): |
| 157 | return None |
| 158 | for dim, (whole, local) in enumerate(zip(global_shape, local_shape)): |
| 159 | expected = local * shard_axis_size if dim == declared_dim else local |
| 160 | if whole != expected: |
| 161 | return None |
| 162 | return declared_dim |
| 163 | |
| 164 | |
| 165 | def resolve_parameter_destinations( |
| 166 | model: torch.nn.Module, |
| 167 | names: Sequence[str], |
| 168 | dtypes: Sequence[torch.dtype], |
| 169 | shapes: Sequence[Sequence[int]], |
| 170 | *, |
| 171 | num_workers: int, |
| 172 | shard_axis_size: int, |
| 173 | allow_direct: bool, |
| 174 | ) -> list[M2NDestination]: |
| 175 | """Build this worker's destination plan, one entry per checkpoint parameter. |
| 176 | |
| 177 | Placements are relative to the inference mesh: axis 0 replicates and axis 1 |
| 178 | shards. `allow_direct=False` forces every parameter onto the fallback path; |
| 179 | the engine sets it for pipeline-parallel or quantized deployments, where a |
| 180 | parameter's local shape is not simply the checkpoint shape split across the |
| 181 | shard axis. Parameters below a model type with known model-level checkpoint |
| 182 | preprocessing also fall back, since bypassing `load_weights` would skip the |
| 183 | transformation. |
| 184 | """ |
| 185 | params = dict(model.named_parameters()) if allow_direct else {} |
| 186 | blocked_param_ids = { |
| 187 | id(param) |
| 188 | for module in model.modules() |
| 189 | if _module_type_name(module) in _DIRECT_WEIGHT_LOADING_BLOCKLIST |
| 190 | for param in module.parameters() |
| 191 | } |
| 192 | |
| 193 | destinations: list[M2NDestination] = [] |
| 194 | for name, dtype, shape in zip(names, dtypes, shapes): |
| 195 | param = params.get(name) |
| 196 | dim = None |
| 197 | if ( |
| 198 | param is not None |
| 199 | and id(param) not in blocked_param_ids |
| 200 | and param.dtype == dtype |
| 201 | and param.data.is_contiguous() |
| 202 | and num_workers % shard_axis_size == 0 |
| 203 | ): |
| 204 | declared_dim = _base_loader_shard_dim(param, shard_axis_size) |
| 205 | if declared_dim is not None: |
| 206 | dim = _validated_shard_dim( |
| 207 | shape, param.shape, declared_dim, shard_axis_size |
| 208 | ) |
| 209 | |
| 210 | if dim is None: |
| 211 | # A replicated parameter is identical on every rank, so it needs no |
| 212 | # placement of its own — REPLICATED lets resolve_layout spread it |
| 213 | # over all the workers regardless of how the mesh is factored. |
| 214 | destinations.append(M2NDestination(name, REPLICATED, None)) |
| 215 | elif dim == REPLICATE: |
| 216 | assert param is not None |
| 217 | destinations.append(M2NDestination(name, REPLICATED, param.data)) |
| 218 | else: |
| 219 | assert param is not None |
| 220 | destinations.append( |
| 221 | M2NDestination(name, _destination_placements(dim), param.data) |
| 222 | ) |
| 223 | |
| 224 | # Layerwise reload finalizes a module as a unit. Mixing direct and fallback |
| 225 | # parameters within one module can copy stale storage over the direct load. |
| 226 | fallback_modules = { |
| 227 | name.rpartition(".")[0] |
| 228 | for name, destination in zip(names, destinations) |
| 229 | if not destination.direct and name in params |
| 230 | } |
| 231 | destinations = [ |
| 232 | M2NDestination(destination.name, REPLICATED, None) |
| 233 | if destination.direct |
| 234 | and destination.name.rpartition(".")[0] in fallback_modules |
| 235 | else destination |
| 236 | for destination in destinations |
| 237 | ] |
| 238 | |
| 239 | num_direct = sum(d.direct for d in destinations) |
| 240 | parameter_bytes = [ |
| 241 | math.prod(shape) * dtype.itemsize for dtype, shape in zip(dtypes, shapes) |
| 242 | ] |
| 243 | total_bytes = sum(parameter_bytes) |
| 244 | direct_bytes = sum( |
| 245 | nbytes |
| 246 | for destination, nbytes in zip(destinations, parameter_bytes) |
| 247 | if destination.direct |
| 248 | ) |
| 249 | direct_byte_percentage = 100 * direct_bytes / total_bytes if total_bytes else 0 |
| 250 | logger.info( |
| 251 | "nccl_m2n destination plan: %d/%d parameters resharded directly into " |
| 252 | "the model, %d via full-tensor fallback; direct byte coverage: " |
| 253 | "%d/%d bytes (%.1f%%)", |
| 254 | num_direct, |
| 255 | len(destinations), |
| 256 | len(destinations) - num_direct, |
| 257 | direct_bytes, |
| 258 | total_bytes, |
| 259 | direct_byte_percentage, |
| 260 | ) |
| 261 | return destinations |