| 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 | from functools import wraps |
| 6 | from weakref import WeakKeyDictionary, WeakSet |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import ModelConfig |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.layers.attention import is_deferred_attention_layer |
| 13 | from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase |
| 14 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 15 | |
| 16 | from .meta import ( |
| 17 | SKIP_LOAD_TENSORS, |
| 18 | capture_layer_to_meta, |
| 19 | get_numel_loaded, |
| 20 | materialize_layer, |
| 21 | restore_layer_on_meta, |
| 22 | ) |
| 23 | from .types import LayerReloadingInfo |
| 24 | from .utils import ( |
| 25 | get_info_size, |
| 26 | get_layer_params_buffers, |
| 27 | get_layer_size, |
| 28 | get_layer_tensors, |
| 29 | has_device_tensors, |
| 30 | ) |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | __all__ = [ |
| 35 | "get_layerwise_info", |
| 36 | "record_metadata_for_reloading", |
| 37 | "initialize_layerwise_reload", |
| 38 | "finalize_layerwise_processing", |
| 39 | "finalize_layerwise_reload", |
| 40 | ] |
| 41 | |
| 42 | |
| 43 | # Global dict storing information used for layerwise restoring, loading, and processing. |
| 44 | # For more information regarding what info is stored when, see `LayerReloadingInfo` |
| 45 | # |
| 46 | # Use a weak ref dictionary so that modules can be freed when the model is freed. |
| 47 | # Values are sanitized from references to the layer key in order to avoid circular refs |
| 48 | LAYERWISE_INFO: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = ( |
| 49 | WeakKeyDictionary() |
| 50 | ) |
| 51 | |
| 52 | # Global set used to track loading for logging purposes only |
| 53 | LOADING_LAYERS: WeakSet[torch.nn.Module] = WeakSet() |
| 54 | |
| 55 | |
| 56 | def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo: |
| 57 | """Get information related to restoring and layerwise processing. If no previous |
| 58 | information existed, a new entry is constructed |
| 59 | """ |
| 60 | if layer not in LAYERWISE_INFO: |
| 61 | LAYERWISE_INFO[layer] = LayerReloadingInfo( |
| 62 | restore_metadata=({}, {}), |
| 63 | restore_device=torch.get_default_device(), |
| 64 | ) |
| 65 | |
| 66 | return LAYERWISE_INFO[layer] |
| 67 | |
| 68 | |
| 69 | def record_metadata_for_reloading(model: torch.nn.Module): |
| 70 | """Record layer metadata needed for later reloading. |
| 71 | |
| 72 | Stores parameter and buffer metadata as meta tensors for restoration. |
| 73 | Must be called before `initialize_layerwise_reload`. |
| 74 | """ |
| 75 | for layer in model.modules(): |
| 76 | info = get_layerwise_info(layer) |
| 77 | info.restore_metadata = capture_layer_to_meta(layer) |
| 78 | info.restore_device = torch.get_default_device() |
| 79 | |
| 80 | |
| 81 | @torch.no_grad() |
| 82 | def initialize_layerwise_reload(model: torch.nn.Module): |
| 83 | """Set up layerwise weight loading with deferred processing. |
| 84 | |
| 85 | Must be called after `record_metadata_for_reloading`. This function: |
| 86 | 1. Saves current kernel tensors for later copying |
| 87 | 2. Restores layer parameters/buffers from metadata (on meta device) |
| 88 | 3. Wraps weight loaders to defer processing until all weights are loaded |
| 89 | |
| 90 | When all weights for a layer are loaded, the wrapped loaders will: |
| 91 | 1. Materialize the layer onto the target device |
| 92 | 2. Load all cached weights |
| 93 | 3. Run quantization processing if applicable |
| 94 | 4. Copy processed values back to original tensor storage |
| 95 | """ |
| 96 | # disable torchao reloading to avoid infinite recursion |
| 97 | model._original_do_torchao_reload = getattr(model, "_do_torchao_reload", False) |
| 98 | model._do_torchao_reload = False |
| 99 | |
| 100 | for layer in model.modules(): |
| 101 | info = get_layerwise_info(layer) |
| 102 | |
| 103 | # Skip if the layer has already been initialized |
| 104 | if info.can_load(): |
| 105 | continue |
| 106 | |
| 107 | # Save current tensors for later copying |
| 108 | info.kernel_tensors = get_layer_params_buffers(layer) |
| 109 | # snapshot now: restore_layer_on_meta drops alias buffers from the live set |
| 110 | info.kernel_non_persistent_buffers = set(layer._non_persistent_buffers_set) |
| 111 | |
| 112 | # Restore layer parameters/buffers onto meta device |
| 113 | restore_layer_on_meta(layer, info) |
| 114 | |
| 115 | # Wrap weight loaders to buffer loading |
| 116 | initialize_online_processing(layer) |
| 117 | |
| 118 | |
| 119 | def initialize_online_processing(layer: torch.nn.Module): |
| 120 | """Wrap a layer's weight loaders with online processing loaders. |
| 121 | Called by either `initialize_layerwise_reload` or an online quantization scheme, |
| 122 | prevents double wrapping in the case of online quantization + reloading |
| 123 | |
| 124 | Args: |
| 125 | layer: layer whose parameter weight loaders will be wrapped |
| 126 | |
| 127 | """ |
| 128 | info = get_layerwise_info(layer) |
| 129 | |
| 130 | # Track loading progress to determine when to process/copy |
| 131 | info.load_numel = 0 |
| 132 | info.load_numel_total = get_layer_size(layer) |
| 133 | _wrap_parameters_weight_loader(layer) |
| 134 | |
| 135 | |
| 136 | def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None: |
| 137 | """Wrap each parameter's weight loader.""" |
| 138 | # Note that nested wrapping will occur for shared tensors |
| 139 | for name, tensor in get_layer_tensors(layer).items(): |
| 140 | if name in SKIP_LOAD_TENSORS: |
| 141 | continue |
| 142 | if _get_weight_loader(tensor).__name__ != "online_process_loader": |
| 143 | tensor.weight_loader = make_online_process_loader(layer, name) |
| 144 | |
| 145 | |
| 146 | def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Callable: |
| 147 | """Create a wrapped weight loader that defers processing.""" |
| 148 | info = get_layerwise_info(layer) |
| 149 | param = getattr(layer, param_name) |
| 150 | original_loader = _get_original_loader(param) |
| 151 | loader_signature = inspect.signature(original_loader) |
| 152 | |
| 153 | @wraps(original_loader, assigned=("__doc__", "__annotations__")) |
| 154 | def online_process_loader(*args, **kwargs): |
| 155 | if not info.can_load(): |
| 156 | # Unfortunately, some qconfigs are set up to load the same weight |
| 157 | # multiple times. For example, CT_WNA16 loads `weight_shape` for |
| 158 | # each of the qkv partitions. This results in layers loading extra |
| 159 | # weights (beyond load_numel_total) after it's already processed. |
| 160 | # |
| 161 | # Best solution is to ensure that `load_numel_total` reflects the |
| 162 | # actual number of weights loaded, either by modifying qconfigs to |
| 163 | # create as many weights as loaded (see padding issue as well) |
| 164 | # or maybe capturing how many weights are loaded on first pass |
| 165 | # |
| 166 | # For now, `load_numel_total` is still safe to use as long as |
| 167 | # there's no way to reach `load_numel_total` without loading all |
| 168 | # necessary weights. `weight_shape` is very small, so this is safe. |
| 169 | # see Limitations(4) |
| 170 | logger.debug("%s: Excessive loading", layer.__class__.__name__) |
| 171 | return |
| 172 | |
| 173 | # Re-run on each load: layers may register parameters later (e.g., `bias`). |
| 174 | # Wrap late parameters and refresh `load_numel_total` so processing waits |
| 175 | # until all parameters are loaded. |
| 176 | info.load_numel_total = get_layer_size(layer) |
| 177 | _wrap_parameters_weight_loader(layer) |
| 178 | |
| 179 | # Bind and normalize arguments |
| 180 | bound_args = loader_signature.bind(*args, **kwargs) |
| 181 | bound_args.apply_defaults() |
| 182 | |
| 183 | # Buffer loaded weights, track loading progress |
| 184 | info.loaded_weights.append((param_name, bound_args)) |
| 185 | num_loaded, ret = get_numel_loaded(original_loader, bound_args) |
| 186 | info.load_numel += num_loaded |
| 187 | |
| 188 | logger.debug( |
| 189 | "%s: %d / %d", |
| 190 | layer.__class__.__name__, |
| 191 | info.load_numel, |
| 192 | info.load_numel_total, |
| 193 | ) |
| 194 | |
| 195 | # Do not online process attention layers, must wait until finalize |
| 196 | if is_deferred_attention_layer(layer): |
| 197 | return ret |
| 198 | |
| 199 | # Log warnings allocating excessive buffers on device |
| 200 | if has_device_tensors(bound_args) and layer not in LOADING_LAYERS: |
| 201 | LOADING_LAYERS.add(layer) |
| 202 | if len(LOADING_LAYERS) == 2: |
| 203 | names = sorted([layer.__class__.__name__ for layer in LOADING_LAYERS]) |
| 204 | mem_used = sum( |
| 205 | get_info_size(LAYERWISE_INFO[layer]) for layer in LOADING_LAYERS |
| 206 | ) |
| 207 | logger.warning_once( |
| 208 | "Allocating %.1f MB of device memory to buffers to load %s layers. " |
| 209 | "This extra memory usage can be avoided by ordering weights " |
| 210 | "by their parent layer when reloading.", |
| 211 | mem_used / 1e6, |
| 212 | str(list(names)), |
| 213 | ) |
| 214 | |
| 215 | # Process and copy when all weights are loaded |
| 216 | if info.load_numel >= info.load_numel_total: # type: ignore[operator] |
| 217 | _layerwise_process(layer, info) |
| 218 | LOADING_LAYERS.discard(layer) |
| 219 | |
| 220 | return ret |
| 221 | |
| 222 | return online_process_loader |
| 223 | |
| 224 | |
| 225 | def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelConfig): |
| 226 | """Apply processing to any layers which were not layerwise processed during loading. |
| 227 | This includes attention layers and layers which have weight elements which are not |
| 228 | loaded (due to padding). |
| 229 | |
| 230 | This function should be applied after `initialize_layerwise_reload` is applied |
| 231 | unwrap the layerwise weight loaders. |
| 232 | |
| 233 | Args: |
| 234 | model: model to finalize processing for |
| 235 | model_config: config needed for applying processing to attention layers |
| 236 | |
| 237 | """ |
| 238 | if hasattr(model, "_original_do_torchao_reload"): |
| 239 | model._do_torchao_reload = model._original_do_torchao_reload |
| 240 | |
| 241 | deferred_attn: list[tuple[torch.nn.Module, LayerReloadingInfo]] = [] |
| 242 | |
| 243 | for layer in model.modules(): |
| 244 | info = get_layerwise_info(layer) |
| 245 | if not info.can_load(): |
| 246 | info.reset() |
| 247 | continue |
| 248 | |
| 249 | # Deferred attention-like layers are processed after all other layers |
| 250 | if is_deferred_attention_layer(layer): |
| 251 | deferred_attn.append((layer, info)) |
| 252 | continue |
| 253 | |
| 254 | # No weights were loaded |
| 255 | if info.load_numel <= 0: |
| 256 | # first load: checkpoint did not contain weights for this layer |
| 257 | if info.kernel_tensors is None: |
| 258 | _layerwise_process(layer, info) |
| 259 | continue |
| 260 | |
| 261 | # reloading: place kernel tensors back as a fallback. Always place, even |
| 262 | # when nothing is loadable (load_numel_total == 0), so parameter-alias |
| 263 | # buffers on such layers are restored rather than left deleted. |
| 264 | if info.load_numel_total > 0: # type: ignore[operator] |
| 265 | logger.warning("%s: Failed to load weights", layer.__class__.__name__) |
| 266 | _place_kernel_tensors(layer, info) |
| 267 | |
| 268 | # Process non-attention layers which did not load all elements. This can happen |
| 269 | # if the created weight has extra padding elements which are not loaded |
| 270 | # Having too many of these delayed layers can lead to excess memory usage |
| 271 | # see Limitations(4) |
| 272 | elif info.load_numel > 0 and info.load_numel < info.load_numel_total: # type: ignore[operator] |
| 273 | logger.debug("%s: Delayed processing", layer.__class__.__name__) |
| 274 | _layerwise_process(layer, info) |
| 275 | |
| 276 | info.reset() |
| 277 | |
| 278 | # Process attention layers after all other layers are done |
| 279 | for layer, info in deferred_attn: |
| 280 | _finalize_attention_layer(layer, info, model_config) |
| 281 | info.reset() |
| 282 | |
| 283 | LOADING_LAYERS.clear() |
| 284 | |
| 285 | |
| 286 | def finalize_layerwise_reload(*args, **kwargs): |
| 287 | finalize_layerwise_processing(*args, **kwargs) |
| 288 | |
| 289 | |
| 290 | def _finalize_attention_layer( |
| 291 | layer: torch.nn.Module, info: LayerReloadingInfo, model_config: ModelConfig |
| 292 | ) -> None: |
| 293 | if info.kernel_tensors is None: |
| 294 | if info.load_numel > 0: |
| 295 | _layerwise_process(layer, info) |
| 296 | elif info.load_numel > 0: |
| 297 | # Reload with new scale weights from checkpoint |
| 298 | _place_kernel_tensors(layer, info) |
| 299 | _reload_attention_scales(layer, info) |
| 300 | else: |
| 301 | _place_kernel_tensors(layer, info) |
| 302 | layer.process_weights_after_loading(model_config.dtype) |
| 303 | |
| 304 | |
| 305 | def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -> None: |
| 306 | """Load and process attention scale weights (k_scale, v_scale, etc.) |
| 307 | during reload. |
| 308 | |
| 309 | Assumes dtype/shapes of attention tensors do not change during |
| 310 | processing, since we use .data.copy_() to preserve kernel tensor |
| 311 | references.""" |
| 312 | quant_method = getattr(layer, "quant_method", None) |
| 313 | if quant_method is not None: |
| 314 | # Re-create scale Parameters with sentinel values so unloaded scales |
| 315 | # are correctly detected by process_weights_after_loading |
| 316 | quant_method.create_weights(layer) |
| 317 | |
| 318 | for name, args in info.loaded_weights: |
| 319 | param = getattr(layer, name) |
| 320 | args.arguments["param"] = param |
| 321 | _get_weight_loader(param)(*args.args, **args.kwargs) |
| 322 | |
| 323 | if quant_method is not None: |
| 324 | quant_method.process_weights_after_loading(layer) |
| 325 | |
| 326 | _copy_and_restore_kernel_tensors(layer, info) |
| 327 | |
| 328 | |
| 329 | def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): |
| 330 | """Finalize layer loading after all weights have been buffered. |
| 331 | |
| 332 | This function: |
| 333 | 1. Materializes the layer onto the target device |
| 334 | 2. Loads all buffered weights |
| 335 | 3. Runs quantization processing if applicable |
| 336 | 4. Copies processed values back to original tensor storage |
| 337 | """ |
| 338 | # Materialize layer tensors onto device |
| 339 | materialize_layer(layer, info) |
| 340 | |
| 341 | # Reset online quantization flag so process_weights_after_loading |
| 342 | # will run again during reload |
| 343 | if hasattr(layer, "_already_called_process_weights_after_loading"): |
| 344 | delattr(layer, "_already_called_process_weights_after_loading") |
| 345 | |
| 346 | # Unwrap layerwise loading wrappers |
| 347 | for param in get_layer_tensors(layer).values(): |
| 348 | param.weight_loader = _get_original_loader(param) |
| 349 | |
| 350 | # Load all buffered weights into materialized layer (using original loaders) |
| 351 | for name, args in info.loaded_weights: |
| 352 | param = getattr(layer, name) |
| 353 | args.arguments["param"] = param |
| 354 | param.weight_loader(*args.args, **args.kwargs) |
| 355 | |
| 356 | # Process weights (quantization, repacking, etc.) |
| 357 | quant_method = getattr(layer, "quant_method", None) |
| 358 | if isinstance(quant_method, QuantizeMethodBase): |
| 359 | quant_method.process_weights_after_loading(layer) |
| 360 | # Re-reconcile parameter TP state: process_weights_after_loading may |
| 361 | # have re-created Parameters (stamped with the global rank), which would |
| 362 | # otherwise break replicated (disable_tp) weights on a subsequent reload. |
| 363 | if hasattr(layer, "update_param_tp_status"): |
| 364 | layer.update_param_tp_status() |
| 365 | |
| 366 | # Copy processed values into original tensor storage (preserves cudagraph refs) |
| 367 | # this code is a no-op if not reloading (because kernel tensors is empty) |
| 368 | if info.kernel_tensors is not None: |
| 369 | _copy_and_restore_kernel_tensors(layer, info) |
| 370 | |
| 371 | info.reset() |
| 372 | logger.debug("%s: Processed", layer.__class__.__name__) |
| 373 | |
| 374 | |
| 375 | def _get_original_loader(tensor: torch.Tensor) -> Callable: |
| 376 | """Return the weight loader with any layerwise wrappers removed.""" |
| 377 | loader = _get_weight_loader(tensor) |
| 378 | while loader.__name__ == "online_process_loader": |
| 379 | loader = loader.__wrapped__ # type: ignore[union-attr] |
| 380 | |
| 381 | return loader |
| 382 | |
| 383 | |
| 384 | def _get_weight_loader(tensor: torch.Tensor): |
| 385 | return getattr(tensor, "weight_loader", default_weight_loader) |
| 386 | |
| 387 | |
| 388 | def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo): |
| 389 | """Copy processed values into original kernel tensor storage and restore |
| 390 | kernel tensor references on the layer. Preserves cudagraph references.""" |
| 391 | assert info.kernel_tensors is not None |
| 392 | parameters, buffers = info.kernel_tensors |
| 393 | non_persistent = info.kernel_non_persistent_buffers |
| 394 | loaded_tensor_names = {name for name, _ in info.loaded_weights} |
| 395 | for name, param in parameters.items(): |
| 396 | param.data.copy_(getattr(layer, name)) |
| 397 | for name, buffer in buffers.items(): |
| 398 | if name not in layer._buffers: |
| 399 | continue |
| 400 | if name in non_persistent and name not in loaded_tensor_names: |
| 401 | continue |
| 402 | buffer.data.copy_(getattr(layer, name)) |
| 403 | |
| 404 | _place_kernel_tensors(layer, info) |
| 405 | |
| 406 | |
| 407 | def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo): |
| 408 | for name in get_layer_tensors(layer): |
| 409 | delattr(layer, name) |
| 410 | |
| 411 | assert info.kernel_tensors is not None |
| 412 | parameters, buffers = info.kernel_tensors |
| 413 | non_persistent = info.kernel_non_persistent_buffers |
| 414 | for name, param in parameters.items(): |
| 415 | layer.register_parameter(name, param) |
| 416 | for name, buffer in buffers.items(): |
| 417 | layer.register_buffer(name, buffer, persistent=name not in non_persistent) |