| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Apply dense or sparse weight updates in checkpoint coordinates. |
| 4 | |
| 5 | The model's ``load_weights`` method still handles checkpoint-name mapping, |
| 6 | TP slicing, and packed parameters. Sparse patches use NaN to mark unchanged |
| 7 | checkpoint elements. The supported sparse-loader contract permits only final |
| 8 | same-shaped floating-point ``Tensor.copy_`` writes. Composed, multi-stage, and |
| 9 | custom-write loaders are unsupported. |
| 10 | """ |
| 11 | |
| 12 | from __future__ import annotations |
| 13 | |
| 14 | import math |
| 15 | from collections.abc import Iterable |
| 16 | from typing import NamedTuple |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | _DEFAULT_PATCH_CHUNK_BYTES = 512 << 20 |
| 21 | |
| 22 | __all__ = [ |
| 23 | "CheckpointWeightPatch", |
| 24 | "load_checkpoint_weight_patches", |
| 25 | ] |
| 26 | |
| 27 | |
| 28 | class CheckpointWeightPatch(NamedTuple): |
| 29 | """Describe one dense or sparse checkpoint-coordinate update. |
| 30 | |
| 31 | Attributes: |
| 32 | name: Checkpoint weight name passed to the model loader. |
| 33 | shape: Full checkpoint tensor shape. |
| 34 | dtype: Checkpoint tensor dtype. |
| 35 | values: The flattened full tensor for a dense patch, or the values at |
| 36 | ``indices`` for a sparse patch. |
| 37 | indices: Flat indices into the full checkpoint tensor described by |
| 38 | ``shape``. ``None`` makes ``values`` a dense replacement. |
| 39 | |
| 40 | """ |
| 41 | |
| 42 | name: str |
| 43 | shape: tuple[int, ...] |
| 44 | dtype: torch.dtype |
| 45 | values: torch.Tensor |
| 46 | indices: torch.Tensor | None = None |
| 47 | |
| 48 | |
| 49 | def _load_nan_masked_weights( |
| 50 | model: torch.nn.Module, |
| 51 | weights: list[tuple[str, torch.Tensor]], |
| 52 | ) -> Iterable[str] | None: |
| 53 | """Load NaN-masked checkpoint tensors while preserving runtime values at NaNs.""" |
| 54 | # Save the original copy_ so it can be restored after loading. |
| 55 | original_copy = torch.Tensor.copy_ |
| 56 | |
| 57 | # Enforce the sparse-loader contract and preserve destination values where |
| 58 | # the corresponding source is NaN. |
| 59 | def copy_non_nan_( |
| 60 | destination: torch.Tensor, |
| 61 | source: torch.Tensor, |
| 62 | non_blocking: bool = False, |
| 63 | ) -> torch.Tensor: |
| 64 | if not ( |
| 65 | destination.is_floating_point() |
| 66 | and source.is_floating_point() |
| 67 | and destination.shape == source.shape |
| 68 | ): |
| 69 | raise NotImplementedError( |
| 70 | "Sparse checkpoint patches require a same-shaped " |
| 71 | "floating-point final copy" |
| 72 | ) |
| 73 | source = source.to(dtype=destination.dtype, device=destination.device) |
| 74 | return torch.where( |
| 75 | torch.isnan(source), |
| 76 | destination, |
| 77 | source, |
| 78 | out=destination, |
| 79 | ) |
| 80 | |
| 81 | torch.Tensor.copy_ = copy_non_nan_ |
| 82 | try: |
| 83 | return model.load_weights(weights) |
| 84 | finally: |
| 85 | torch.Tensor.copy_ = original_copy |
| 86 | |
| 87 | |
| 88 | def _validate_patch_structure(patch: CheckpointWeightPatch) -> int: |
| 89 | if not patch.name: |
| 90 | raise ValueError("Checkpoint weight name must be non-empty") |
| 91 | if not isinstance(patch.dtype, torch.dtype): |
| 92 | raise TypeError(f"{patch.name}: dtype must be a torch.dtype") |
| 93 | if any(dim < 0 for dim in patch.shape): |
| 94 | raise ValueError(f"{patch.name}: shape dimensions must be non-negative") |
| 95 | if patch.values.ndim != 1: |
| 96 | raise ValueError(f"{patch.name}: patch values must be a 1D tensor") |
| 97 | |
| 98 | numel = math.prod(patch.shape) |
| 99 | if patch.indices is None: |
| 100 | if patch.values.numel() != numel: |
| 101 | raise ValueError( |
| 102 | f"{patch.name}: dense patch has {patch.values.numel()} values " |
| 103 | f"for a {numel}-element tensor" |
| 104 | ) |
| 105 | return numel |
| 106 | |
| 107 | if not patch.dtype.is_floating_point: |
| 108 | raise TypeError( |
| 109 | f"{patch.name}: sparse checkpoint patches require a floating dtype " |
| 110 | "because NaN is the unchanged-value sentinel" |
| 111 | ) |
| 112 | if patch.indices.ndim != 1: |
| 113 | raise ValueError(f"{patch.name}: sparse patch indices must be a 1D tensor") |
| 114 | if patch.indices.dtype not in (torch.int32, torch.int64): |
| 115 | raise TypeError(f"{patch.name}: sparse patch indices must be int32 or int64") |
| 116 | if patch.indices.numel() != patch.values.numel(): |
| 117 | raise ValueError( |
| 118 | f"{patch.name}: sparse indices and values must have matching lengths" |
| 119 | ) |
| 120 | |
| 121 | return numel |
| 122 | |
| 123 | |
| 124 | def _load_weight_chunk( |
| 125 | model: torch.nn.Module, |
| 126 | weights: list[tuple[str, torch.Tensor]], |
| 127 | *, |
| 128 | sparse: bool, |
| 129 | ) -> set[str]: |
| 130 | if not weights: |
| 131 | return set() |
| 132 | if sparse: |
| 133 | loaded_names = _load_nan_masked_weights(model, weights) |
| 134 | else: |
| 135 | loaded_names = model.load_weights(weights) |
| 136 | return set() if loaded_names is None else set(loaded_names) |
| 137 | |
| 138 | |
| 139 | @torch.no_grad() |
| 140 | def load_checkpoint_weight_patches( |
| 141 | model: torch.nn.Module, |
| 142 | patches: Iterable[CheckpointWeightPatch], |
| 143 | *, |
| 144 | max_chunk_bytes: int = _DEFAULT_PATCH_CHUNK_BYTES, |
| 145 | validate_unique_indices: bool = True, |
| 146 | ) -> set[str]: |
| 147 | """Load ordered patches through the model's checkpoint loader. |
| 148 | |
| 149 | A call may contain either dense or sparse patches, but not both. A repeated |
| 150 | name starts a new loader call, so patches for that weight are applied in input |
| 151 | order. A later patch may update positions changed by an earlier patch. Each |
| 152 | sparse patch creates a full checkpoint-shaped tensor whose NaNs mark unchanged |
| 153 | elements. Each locally applied sparse patch must use one final same-shaped |
| 154 | floating-point ``Tensor.copy_``. Intermediate or unrelated ``copy_`` calls, |
| 155 | composed loaders, and custom write paths are unsupported. ``max_chunk_bytes`` |
| 156 | is a batching target; one tensor may exceed it. |
| 157 | |
| 158 | Online checkpoint-format dense updates require the caller to manage vLLM's |
| 159 | layerwise reload lifecycle. Sparse online updates modify initialized model |
| 160 | tensors in place and must not use that lifecycle. Patch shapes, dtypes, value |
| 161 | lengths, and sparse indices are validated before loading begins. |
| 162 | |
| 163 | This function does not roll back a partial update. If a loader fails after a |
| 164 | write, some destinations may already be changed. Do not serve from the affected |
| 165 | worker until a known baseline has been restored or the worker has been |
| 166 | restarted. |
| 167 | |
| 168 | Args: |
| 169 | model: Model whose native ``load_weights`` method applies the patches. |
| 170 | patches: Dense replacements or sparse updates in checkpoint coordinates. |
| 171 | max_chunk_bytes: Target checkpoint tensor bytes per |
| 172 | ``model.load_weights`` call. A single tensor may exceed this target. |
| 173 | validate_unique_indices: Whether to reject duplicate sparse indices within |
| 174 | each patch. Repeated patches may update the same positions. Disable |
| 175 | only for a trusted producer that already guarantees unique positions |
| 176 | within each patch. |
| 177 | |
| 178 | Returns: |
| 179 | Union of the weight names reported by all ``model.load_weights`` calls. |
| 180 | |
| 181 | """ |
| 182 | if max_chunk_bytes <= 0: |
| 183 | raise ValueError("max_chunk_bytes must be positive") |
| 184 | |
| 185 | patch_list = list(patches) |
| 186 | patch_numels = [_validate_patch_structure(patch) for patch in patch_list] |
| 187 | # All patches in one call must use the same dense or sparse representation. |
| 188 | sparse_flags = {patch.indices is not None for patch in patch_list} |
| 189 | if len(sparse_flags) > 1: |
| 190 | raise ValueError("Dense and sparse checkpoint patches cannot be mixed") |
| 191 | sparse = sparse_flags == {True} |
| 192 | |
| 193 | # Keep validation results on their devices and read one combined result per |
| 194 | # device, instead of synchronizing the CPU once for every patch. |
| 195 | checks_by_device: dict[torch.device, list[torch.Tensor]] = {} |
| 196 | if sparse: |
| 197 | for patch, numel in zip(patch_list, patch_numels, strict=True): |
| 198 | assert patch.indices is not None |
| 199 | |
| 200 | # NaN is reserved to mean "leave this checkpoint position unchanged," |
| 201 | # so it cannot also be supplied as a new sparse value. |
| 202 | patch_checks = [torch.isnan(patch.values).any()] |
| 203 | indices = patch.indices.to(device=patch.values.device) |
| 204 | if indices.numel(): |
| 205 | # Sparse indices address the flattened full checkpoint tensor. |
| 206 | patch_checks.append( |
| 207 | torch.logical_or(indices < 0, indices >= numel).any() |
| 208 | ) |
| 209 | if validate_unique_indices: |
| 210 | # Equal neighboring indices after sorting indicate duplicates. |
| 211 | sorted_indices = torch.sort(indices).values |
| 212 | patch_checks.append( |
| 213 | (sorted_indices[1:] == sorted_indices[:-1]).any() |
| 214 | ) |
| 215 | |
| 216 | checks_by_device.setdefault(patch.values.device, []).extend(patch_checks) |
| 217 | |
| 218 | for device_checks in checks_by_device.values(): |
| 219 | if torch.stack(device_checks).any().item(): |
| 220 | invalid_reasons = ( |
| 221 | "NaN values, out-of-range indices, or duplicate indices" |
| 222 | if validate_unique_indices |
| 223 | else "NaN values or out-of-range indices" |
| 224 | ) |
| 225 | raise ValueError(f"Sparse checkpoint patches contain {invalid_reasons}") |
| 226 | |
| 227 | loaded_names: set[str] = set() |
| 228 | # Collect checkpoint tensors for one model.load_weights() call. |
| 229 | weight_chunk: list[tuple[str, torch.Tensor]] = [] |
| 230 | chunk_names: set[str] = set() |
| 231 | chunk_bytes = 0 |
| 232 | |
| 233 | for patch, numel in zip(patch_list, patch_numels, strict=True): |
| 234 | # Count the full checkpoint tensor, including sparse NaN staging. |
| 235 | checkpoint_tensor_bytes = ( |
| 236 | numel * torch.empty((), dtype=patch.dtype).element_size() |
| 237 | ) |
| 238 | # Flush before exceeding the batching target. Repeated names also start a |
| 239 | # new loader call so their patches are applied in input order. |
| 240 | if weight_chunk and ( |
| 241 | chunk_bytes + checkpoint_tensor_bytes > max_chunk_bytes |
| 242 | or patch.name in chunk_names |
| 243 | ): |
| 244 | loaded_names.update(_load_weight_chunk(model, weight_chunk, sparse=sparse)) |
| 245 | weight_chunk = [] |
| 246 | chunk_names = set() |
| 247 | chunk_bytes = 0 |
| 248 | |
| 249 | if patch.indices is None: |
| 250 | checkpoint_tensor = patch.values.to(dtype=patch.dtype).view(patch.shape) |
| 251 | else: |
| 252 | # NaNs mark unchanged checkpoint positions; index_copy_ fills only |
| 253 | # the positions supplied by this sparse patch. |
| 254 | flat = torch.full( |
| 255 | (numel,), |
| 256 | float("nan"), |
| 257 | dtype=patch.dtype, |
| 258 | device=patch.values.device, |
| 259 | ) |
| 260 | if patch.values.numel(): |
| 261 | flat.index_copy_( |
| 262 | 0, |
| 263 | patch.indices.to(device=flat.device, dtype=torch.long), |
| 264 | patch.values.to(dtype=patch.dtype, device=flat.device), |
| 265 | ) |
| 266 | checkpoint_tensor = flat.view(patch.shape) |
| 267 | |
| 268 | weight_chunk.append((patch.name, checkpoint_tensor)) |
| 269 | chunk_names.add(patch.name) |
| 270 | chunk_bytes += checkpoint_tensor_bytes |
| 271 | |
| 272 | # Load the final partial chunk; the helper also accepts an empty chunk. |
| 273 | loaded_names.update(_load_weight_chunk(model, weight_chunk, sparse=sparse)) |
| 274 | return loaded_names |