| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import inspect |
| 4 | from collections.abc import Callable |
| 5 | |
| 6 | import torch |
| 7 | from torch.nn.parameter import UninitializedParameter |
| 8 | from torch.utils._python_dispatch import TorchDispatchMode |
| 9 | |
| 10 | from .sanitize import restore_layer_refs, sanitize_layer_refs |
| 11 | from .types import LayerReloadingInfo, LayerTensors |
| 12 | from .utils import get_layer_params_buffers, get_layer_tensors, get_tensor_load_numel |
| 13 | |
| 14 | __all__ = [ |
| 15 | "to_meta_tensor", |
| 16 | "materialize_meta_tensor", |
| 17 | "capture_layer_to_meta", |
| 18 | "restore_layer_on_meta", |
| 19 | "materialize_layer", |
| 20 | "get_numel_loaded", |
| 21 | ] |
| 22 | |
| 23 | # Modules whose tensors are never moved to, or materialized from, the meta device. |
| 24 | SKIP_MODULES: set[str] = {"HadamardTransform"} |
| 25 | |
| 26 | # Tensors never loaded by a weight loader, so the layerwise trigger ignores them. |
| 27 | SKIP_LOAD_TENSORS: set[str] = { |
| 28 | "_expert_map", |
| 29 | "expert_mask", |
| 30 | "expert_global_to_physical", |
| 31 | "expert_physical_to_global", |
| 32 | "expert_local_to_global", |
| 33 | "e_score_correction_bias", |
| 34 | } |
| 35 | |
| 36 | # Tensors which are never moved to, or materialized from, the meta device. |
| 37 | # `bias` is built after create_weights(), so it is never on meta to begin with. |
| 38 | SKIP_TENSORS: set[str] = SKIP_LOAD_TENSORS | {"bias"} |
| 39 | |
| 40 | |
| 41 | def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor: |
| 42 | """Convert a tensor to a meta tensor while preserving class and attributes.""" |
| 43 | meta_tensor = tensor.data.to("meta") |
| 44 | meta_tensor.__class__ = tensor.__class__ |
| 45 | meta_tensor.__dict__ = tensor.__dict__.copy() |
| 46 | return meta_tensor |
| 47 | |
| 48 | |
| 49 | def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor: |
| 50 | """Materialize a meta tensor into an actual tensor on the current device. |
| 51 | Should be called within the torch device context for the given rank. |
| 52 | """ |
| 53 | tensor = torch.empty_strided( |
| 54 | size=tuple(meta_tensor.size()), |
| 55 | stride=tuple(meta_tensor.stride()), |
| 56 | dtype=meta_tensor.dtype, |
| 57 | requires_grad=False, |
| 58 | ) |
| 59 | tensor.__class__ = meta_tensor.__class__ |
| 60 | tensor.__dict__ = meta_tensor.__dict__.copy() |
| 61 | return tensor |
| 62 | |
| 63 | |
| 64 | def _is_non_persistent_parameter_alias_buffer( |
| 65 | layer: torch.nn.Module, |
| 66 | name: str, |
| 67 | buffer: torch.Tensor, |
| 68 | parameter_storage_ptrs: set[int], |
| 69 | ) -> bool: |
| 70 | if name not in layer._non_persistent_buffers_set: |
| 71 | return False |
| 72 | |
| 73 | buffer_storage_ptr = _tensor_storage_ptr(buffer) |
| 74 | return ( |
| 75 | buffer_storage_ptr is not None and buffer_storage_ptr in parameter_storage_ptrs |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def _tensor_storage_ptr(tensor: torch.Tensor) -> int | None: |
| 80 | if isinstance(tensor, UninitializedParameter): |
| 81 | return None |
| 82 | |
| 83 | try: |
| 84 | return tensor.untyped_storage().data_ptr() |
| 85 | except (RuntimeError, ValueError): |
| 86 | return None |
| 87 | |
| 88 | |
| 89 | def _parameter_storage_ptrs(layer: torch.nn.Module) -> set[int]: |
| 90 | return { |
| 91 | storage_ptr |
| 92 | for param in layer.parameters(recurse=True) |
| 93 | if (storage_ptr := _tensor_storage_ptr(param)) is not None |
| 94 | } |
| 95 | |
| 96 | |
| 97 | def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors: |
| 98 | if layer.__class__.__name__ in SKIP_MODULES: |
| 99 | return ({}, {}) |
| 100 | |
| 101 | params, buffers = get_layer_params_buffers(layer) |
| 102 | parameter_storage_ptrs = _parameter_storage_ptrs(layer) |
| 103 | return ( |
| 104 | { |
| 105 | name: sanitize_layer_refs(to_meta_tensor(param), layer) |
| 106 | for name, param in params.items() |
| 107 | if name not in SKIP_TENSORS |
| 108 | }, |
| 109 | { |
| 110 | name: sanitize_layer_refs(to_meta_tensor(buffer), layer) |
| 111 | for name, buffer in buffers.items() |
| 112 | if name not in SKIP_TENSORS |
| 113 | and not _is_non_persistent_parameter_alias_buffer( |
| 114 | layer, name, buffer, parameter_storage_ptrs |
| 115 | ) |
| 116 | }, |
| 117 | ) |
| 118 | |
| 119 | |
| 120 | def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo): |
| 121 | """Restore a layer to model format with tensors on the meta device.""" |
| 122 | if layer.__class__.__name__ in SKIP_MODULES: |
| 123 | return |
| 124 | |
| 125 | non_persistent = set(layer._non_persistent_buffers_set) |
| 126 | restore_params, restore_buffers = info.restore_metadata |
| 127 | tensor_names = ( |
| 128 | get_layer_tensors(layer).keys() | restore_params.keys() | restore_buffers.keys() |
| 129 | ) |
| 130 | for name in tensor_names: |
| 131 | if name not in SKIP_TENSORS and hasattr(layer, name): |
| 132 | delattr(layer, name) |
| 133 | |
| 134 | for name, param in restore_params.items(): |
| 135 | if name not in SKIP_TENSORS: |
| 136 | param = restore_layer_refs(param, layer) |
| 137 | layer.register_parameter(name, param) |
| 138 | |
| 139 | for name, buffer in restore_buffers.items(): |
| 140 | if name not in SKIP_TENSORS: |
| 141 | buffer = restore_layer_refs(buffer, layer) |
| 142 | layer.register_buffer(name, buffer, persistent=name not in non_persistent) |
| 143 | |
| 144 | |
| 145 | def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo): |
| 146 | """Materialize all meta tensors in a layer to actual tensors.""" |
| 147 | if layer.__class__.__name__ in SKIP_MODULES: |
| 148 | return |
| 149 | |
| 150 | with info.restore_device: |
| 151 | for name, tensor in get_layer_tensors(layer).items(): |
| 152 | if name not in SKIP_TENSORS and tensor.is_meta: |
| 153 | setattr(layer, name, materialize_meta_tensor(tensor)) |
| 154 | |
| 155 | |
| 156 | class CopyCounter(TorchDispatchMode): |
| 157 | """Tracks total number of elements modified with `copy_`. |
| 158 | |
| 159 | Useful for keeping track of weight loading where underlying weights can be |
| 160 | arbitrarily transformed (such as with `narrow`) before calling copy. |
| 161 | |
| 162 | Note: Assumes that copy kwargs are not used. |
| 163 | """ |
| 164 | |
| 165 | def __init__(self): |
| 166 | super().__init__() |
| 167 | self.copied_numel = 0 |
| 168 | |
| 169 | def __torch_dispatch__(self, func, types, args=(), kwargs=None): |
| 170 | if kwargs is None: |
| 171 | kwargs = {} |
| 172 | |
| 173 | if func is torch.ops.aten.copy_.default: |
| 174 | assert args[0].numel() == args[1].numel() |
| 175 | self.copied_numel += args[0].numel() |
| 176 | |
| 177 | return func(*args, **kwargs) |
| 178 | |
| 179 | |
| 180 | def get_numel_loaded( |
| 181 | weight_loader: Callable, args: inspect.BoundArguments |
| 182 | ) -> tuple[int, object]: |
| 183 | """Determine how many elements would be loaded by a weight loader call. |
| 184 | |
| 185 | Args: |
| 186 | weight_loader: used to load weights |
| 187 | args: bound arguments to weight loader |
| 188 | |
| 189 | Returns: |
| 190 | number of elements loaded by the weight loader, the return value of the |
| 191 | weight loader |
| 192 | |
| 193 | """ |
| 194 | with CopyCounter() as counter: |
| 195 | return_value = weight_loader(*args.args, **args.kwargs) |
| 196 | |
| 197 | # A weight loader fills a single destination parameter, so the number of |
| 198 | # loaded elements is at most that parameter's size. Some loaders copy into |
| 199 | # the parameter more than once -- e.g. ``composed_weight_loader`` runs an |
| 200 | # in-place post-load transform (``param.copy_(fn(param))``) on top of the |
| 201 | # initial copy -- which would make CopyCounter report twice the parameter |
| 202 | # size. Over-counting inflates the layer's loaded-element total and can |
| 203 | # finalize the layer before every parameter is loaded, silently dropping |
| 204 | # the trailing parameter(s) (e.g. Mamba ``mixer.D``). Cap the count at the |
| 205 | # destination size to keep the per-layer accounting correct. |
| 206 | numel = counter.copied_numel |
| 207 | param = args.arguments.get("param", None) |
| 208 | if isinstance(param, torch.Tensor): |
| 209 | numel = min(numel, get_tensor_load_numel(param)) |
| 210 | return numel, return_value |