chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,38 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Layerwise weight reloading utilities for vLLM.
|
||||
|
||||
This module provides functionality to reload model weights layer-by-layer,
|
||||
which is useful for weight updates without full model reconstruction.
|
||||
|
||||
Limitations:
|
||||
1. Composition with CPU offloading has not been implemented
|
||||
2. Tied parameters will only reflect processing from one of the parent layers (for
|
||||
example, only processing from embed_tokens will have an effect)
|
||||
3. This design assumes that the number of weights loaded from disk is the same as the
|
||||
number of weights created at model init time. This is not true for quant methods
|
||||
which (1) pad weights or (2) load qkv weights into the same parameter. Both of these
|
||||
cases are non-issues for today's quant methods, but future quantizations may cause
|
||||
reloading to fail
|
||||
"""
|
||||
|
||||
__all__ = [
|
||||
"record_metadata_for_reloading",
|
||||
"initialize_layerwise_reload",
|
||||
"finalize_layerwise_processing",
|
||||
"finalize_layerwise_reload",
|
||||
"set_torchao_reload_attrs",
|
||||
"support_quantized_model_reload_from_hp_weights",
|
||||
]
|
||||
|
||||
from .layerwise import (
|
||||
finalize_layerwise_processing,
|
||||
finalize_layerwise_reload,
|
||||
initialize_layerwise_reload,
|
||||
record_metadata_for_reloading,
|
||||
)
|
||||
from .torchao_decorator import (
|
||||
set_torchao_reload_attrs,
|
||||
support_quantized_model_reload_from_hp_weights,
|
||||
)
|
||||
@@ -0,0 +1,410 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from weakref import WeakKeyDictionary, WeakSet
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.attention import Attention, MLAAttention
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from .meta import (
|
||||
SKIP_TENSORS,
|
||||
capture_layer_to_meta,
|
||||
get_numel_loaded,
|
||||
materialize_layer,
|
||||
restore_layer_on_meta,
|
||||
)
|
||||
from .types import LayerReloadingInfo
|
||||
from .utils import (
|
||||
get_info_size,
|
||||
get_layer_params_buffers,
|
||||
get_layer_size,
|
||||
get_layer_tensors,
|
||||
has_device_tensors,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"get_layerwise_info",
|
||||
"record_metadata_for_reloading",
|
||||
"initialize_layerwise_reload",
|
||||
"finalize_layerwise_processing",
|
||||
"finalize_layerwise_reload",
|
||||
]
|
||||
|
||||
|
||||
# Global dict storing information used for layerwise restoring, loading, and processing.
|
||||
# For more information regarding what info is stored when, see `LayerReloadingInfo`
|
||||
#
|
||||
# Use a weak ref dictionary so that modules can be freed when the model is freed.
|
||||
# Values are sanitized from references to the layer key in order to avoid circular refs
|
||||
LAYERWISE_INFO: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (
|
||||
WeakKeyDictionary()
|
||||
)
|
||||
|
||||
# Global set used to track loading for logging purposes only
|
||||
LOADING_LAYERS: WeakSet[torch.nn.Module] = WeakSet()
|
||||
|
||||
|
||||
def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo:
|
||||
"""
|
||||
Get information related to restoring and layerwise processing. If no previous
|
||||
information existed, a new entry is constructed
|
||||
"""
|
||||
if layer not in LAYERWISE_INFO:
|
||||
LAYERWISE_INFO[layer] = LayerReloadingInfo(
|
||||
restore_metadata=({}, {}),
|
||||
restore_device=torch.get_default_device(),
|
||||
)
|
||||
|
||||
return LAYERWISE_INFO[layer]
|
||||
|
||||
|
||||
def record_metadata_for_reloading(model: torch.nn.Module):
|
||||
"""
|
||||
Record layer metadata needed for later reloading.
|
||||
|
||||
Stores parameter and buffer metadata as meta tensors for restoration.
|
||||
Must be called before `initialize_layerwise_reload`.
|
||||
"""
|
||||
for layer in model.modules():
|
||||
info = get_layerwise_info(layer)
|
||||
info.restore_metadata = capture_layer_to_meta(layer)
|
||||
info.restore_device = torch.get_default_device()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def initialize_layerwise_reload(model: torch.nn.Module):
|
||||
"""
|
||||
Set up layerwise weight loading with deferred processing.
|
||||
|
||||
Must be called after `record_metadata_for_reloading`. This function:
|
||||
1. Saves current kernel tensors for later copying
|
||||
2. Restores layer parameters/buffers from metadata (on meta device)
|
||||
3. Wraps weight loaders to defer processing until all weights are loaded
|
||||
|
||||
When all weights for a layer are loaded, the wrapped loaders will:
|
||||
1. Materialize the layer onto the target device
|
||||
2. Load all cached weights
|
||||
3. Run quantization processing if applicable
|
||||
4. Copy processed values back to original tensor storage
|
||||
"""
|
||||
# disable torchao reloading to avoid infinite recursion
|
||||
model._original_do_torchao_reload = getattr(model, "_do_torchao_reload", False)
|
||||
model._do_torchao_reload = False
|
||||
|
||||
for layer in model.modules():
|
||||
info = get_layerwise_info(layer)
|
||||
|
||||
# Skip if the layer has already been initialized
|
||||
if info.can_load():
|
||||
continue
|
||||
|
||||
# Save current tensors for later copying
|
||||
info.kernel_tensors = get_layer_params_buffers(layer)
|
||||
|
||||
# Restore layer parameters/buffers onto meta device
|
||||
restore_layer_on_meta(layer, info)
|
||||
|
||||
# Wrap weight loaders to buffer loading
|
||||
initialize_online_processing(layer)
|
||||
|
||||
|
||||
def initialize_online_processing(layer: torch.nn.Module):
|
||||
"""
|
||||
Wrap a layer's weight loaders with online processing loaders.
|
||||
Called by either `initialize_layerwise_reload` or an online quantization scheme,
|
||||
prevents double wrapping in the case of online quantization + reloading
|
||||
|
||||
Args:
|
||||
layer: layer whose parameter weight loaders will be wrapped
|
||||
"""
|
||||
info = get_layerwise_info(layer)
|
||||
|
||||
# Track loading progress to determine when to process/copy
|
||||
info.load_numel = 0
|
||||
info.load_numel_total = get_layer_size(layer)
|
||||
_wrap_parameters_weight_loader(layer)
|
||||
|
||||
|
||||
def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None:
|
||||
"""Wrap each parameter's weight loader."""
|
||||
# Note that nested wrapping will occur for shared tensors
|
||||
for name, tensor in get_layer_tensors(layer).items():
|
||||
if name in SKIP_TENSORS:
|
||||
continue
|
||||
if _get_weight_loader(tensor).__name__ != "online_process_loader":
|
||||
tensor.weight_loader = make_online_process_loader(layer, name)
|
||||
|
||||
|
||||
def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Callable:
|
||||
"""Create a wrapped weight loader that defers processing."""
|
||||
info = get_layerwise_info(layer)
|
||||
param = getattr(layer, param_name)
|
||||
original_loader = _get_original_loader(param)
|
||||
loader_signature = inspect.signature(original_loader)
|
||||
|
||||
@wraps(original_loader, assigned=("__doc__", "__annotations__"))
|
||||
def online_process_loader(*args, **kwargs):
|
||||
if not info.can_load():
|
||||
# Unfortunately, some qconfigs are set up to load the same weight
|
||||
# multiple times. For example, CT_WNA16 loads `weight_shape` for
|
||||
# each of the qkv partitions. This results in layers loading extra
|
||||
# weights (beyond load_numel_total) after it's already processed.
|
||||
#
|
||||
# Best solution is to ensure that `load_numel_total` reflects the
|
||||
# actual number of weights loaded, either by modifying qconfigs to
|
||||
# create as many weights as loaded (see padding issue as well)
|
||||
# or maybe capturing how many weights are loaded on first pass
|
||||
#
|
||||
# For now, `load_numel_total` is still safe to use as long as
|
||||
# there's no way to reach `load_numel_total` without loading all
|
||||
# necessary weights. `weight_shape` is very small, so this is safe.
|
||||
# see Limitations(4)
|
||||
logger.debug("%s: Excessive loading", layer.__class__.__name__)
|
||||
return
|
||||
|
||||
# Re-run on each load: layers may register parameters later (e.g., `bias`).
|
||||
# Wrap late parameters and refresh `load_numel_total` so processing waits
|
||||
# until all parameters are loaded.
|
||||
info.load_numel_total = get_layer_size(layer)
|
||||
_wrap_parameters_weight_loader(layer)
|
||||
|
||||
# Bind and normalize arguments
|
||||
bound_args = loader_signature.bind(*args, **kwargs)
|
||||
bound_args.apply_defaults()
|
||||
|
||||
# Buffer loaded weights, track loading progress
|
||||
info.loaded_weights.append((param_name, bound_args))
|
||||
num_loaded, ret = get_numel_loaded(original_loader, bound_args)
|
||||
info.load_numel += num_loaded
|
||||
|
||||
logger.debug(
|
||||
"%s: %d / %d",
|
||||
layer.__class__.__name__,
|
||||
info.load_numel,
|
||||
info.load_numel_total,
|
||||
)
|
||||
|
||||
# Do not online process attention layers, must wait until finalize
|
||||
if isinstance(layer, (Attention, MLAAttention)):
|
||||
return ret
|
||||
|
||||
# Log warnings allocating excessive buffers on device
|
||||
if has_device_tensors(bound_args):
|
||||
LOADING_LAYERS.add(layer)
|
||||
if len(LOADING_LAYERS) >= 2:
|
||||
names = sorted([layer.__class__.__name__ for layer in LOADING_LAYERS])
|
||||
mem_used = sum(
|
||||
get_info_size(LAYERWISE_INFO[layer]) for layer in LOADING_LAYERS
|
||||
)
|
||||
logger.warning_once(
|
||||
"Allocating %.1f MB of device memory to buffers to load %s layers. "
|
||||
"This extra memory usage can be avoided by ordering weights "
|
||||
"by their parent layer when reloading.",
|
||||
mem_used / 1e6,
|
||||
str(list(names)),
|
||||
)
|
||||
|
||||
# Process and copy when all weights are loaded
|
||||
if info.load_numel >= info.load_numel_total: # type: ignore[operator]
|
||||
_layerwise_process(layer, info)
|
||||
LOADING_LAYERS.discard(layer)
|
||||
|
||||
return ret
|
||||
|
||||
return online_process_loader
|
||||
|
||||
|
||||
def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelConfig):
|
||||
"""
|
||||
Apply processing to any layers which were not layerwise processed during loading.
|
||||
This includes attention layers and layers which have weight elements which are not
|
||||
loaded (due to padding).
|
||||
|
||||
This function should be applied after `initialize_layerwise_reload` is applied
|
||||
unwrap the layerwise weight loaders.
|
||||
|
||||
Args:
|
||||
model: model to finalize processing for
|
||||
model_config: config needed for applying processing to attention layers
|
||||
"""
|
||||
if hasattr(model, "_original_do_torchao_reload"):
|
||||
model._do_torchao_reload = model._original_do_torchao_reload
|
||||
|
||||
deferred_attn: list[tuple[torch.nn.Module, LayerReloadingInfo]] = []
|
||||
|
||||
for layer in model.modules():
|
||||
info = get_layerwise_info(layer)
|
||||
if not info.can_load():
|
||||
info.reset()
|
||||
continue
|
||||
|
||||
# Attention/MLA layers are processed after all other layers
|
||||
if isinstance(layer, (Attention, MLAAttention)):
|
||||
deferred_attn.append((layer, info))
|
||||
continue
|
||||
|
||||
# No weights were loaded
|
||||
if info.load_numel <= 0:
|
||||
# first load: checkpoint did not contain weights for this layer
|
||||
if info.kernel_tensors is None:
|
||||
_layerwise_process(layer, info)
|
||||
continue
|
||||
|
||||
# reloading: place kernel tensors back as a fallback
|
||||
elif info.load_numel_total > 0: # type: ignore[operator]
|
||||
logger.warning("%s: Failed to load weights", layer.__class__.__name__)
|
||||
_place_kernel_tensors(layer, info)
|
||||
|
||||
# Process non-attention layers which did not load all elements. This can happen
|
||||
# if the created weight has extra padding elements which are not loaded
|
||||
# Having too many of these delayed layers can lead to excess memory usage
|
||||
# see Limitations(4)
|
||||
elif info.load_numel > 0 and info.load_numel < info.load_numel_total: # type: ignore[operator]
|
||||
logger.debug("%s: Delayed processing", layer.__class__.__name__)
|
||||
_layerwise_process(layer, info)
|
||||
|
||||
info.reset()
|
||||
|
||||
# Process attention layers after all other layers are done
|
||||
for layer, info in deferred_attn:
|
||||
_finalize_attention_layer(layer, info, model_config)
|
||||
info.reset()
|
||||
|
||||
LOADING_LAYERS.clear()
|
||||
|
||||
|
||||
def finalize_layerwise_reload(*args, **kwargs):
|
||||
finalize_layerwise_processing(*args, **kwargs)
|
||||
|
||||
|
||||
def _finalize_attention_layer(
|
||||
layer: torch.nn.Module, info: LayerReloadingInfo, model_config: ModelConfig
|
||||
) -> None:
|
||||
if info.load_numel > 0 and info.kernel_tensors is not None:
|
||||
# Reload with new scale weights from checkpoint
|
||||
_place_kernel_tensors(layer, info)
|
||||
_reload_attention_scales(layer, info)
|
||||
elif info.load_numel > 0 or info.kernel_tensors is None:
|
||||
raise ValueError(
|
||||
"Layerwise loading of attention layers is not supported. "
|
||||
"Attention must always process after linears."
|
||||
)
|
||||
else:
|
||||
_place_kernel_tensors(layer, info)
|
||||
layer.process_weights_after_loading(model_config.dtype)
|
||||
|
||||
|
||||
def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -> None:
|
||||
"""Load and process attention scale weights (k_scale, v_scale, etc.)
|
||||
during reload.
|
||||
|
||||
Assumes dtype/shapes of attention tensors do not change during
|
||||
processing, since we use .data.copy_() to preserve kernel tensor
|
||||
references."""
|
||||
quant_method = getattr(layer, "quant_method", None)
|
||||
if quant_method is None:
|
||||
return
|
||||
|
||||
# Re-create scale Parameters with sentinel values so unloaded scales
|
||||
# are correctly detected by process_weights_after_loading
|
||||
quant_method.create_weights(layer)
|
||||
|
||||
for name, args in info.loaded_weights:
|
||||
param = getattr(layer, name)
|
||||
args.arguments["param"] = param
|
||||
_get_weight_loader(param)(*args.args, **args.kwargs)
|
||||
|
||||
quant_method.process_weights_after_loading(layer)
|
||||
|
||||
_copy_and_restore_kernel_tensors(layer, info)
|
||||
|
||||
|
||||
def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
|
||||
"""
|
||||
Finalize layer loading after all weights have been buffered.
|
||||
|
||||
This function:
|
||||
1. Materializes the layer onto the target device
|
||||
2. Loads all buffered weights
|
||||
3. Runs quantization processing if applicable
|
||||
4. Copies processed values back to original tensor storage
|
||||
"""
|
||||
# Materialize layer tensors onto device
|
||||
materialize_layer(layer, info)
|
||||
|
||||
# Reset online quantization flag so process_weights_after_loading
|
||||
# will run again during reload
|
||||
if hasattr(layer, "_already_called_process_weights_after_loading"):
|
||||
delattr(layer, "_already_called_process_weights_after_loading")
|
||||
|
||||
# Unwrap layerwise loading wrappers
|
||||
for param in get_layer_tensors(layer).values():
|
||||
param.weight_loader = _get_original_loader(param)
|
||||
|
||||
# Load all buffered weights into materialized layer (using original loaders)
|
||||
for name, args in info.loaded_weights:
|
||||
param = getattr(layer, name)
|
||||
args.arguments["param"] = param
|
||||
param.weight_loader(*args.args, **args.kwargs)
|
||||
|
||||
# Process weights (quantization, repacking, etc.)
|
||||
quant_method = getattr(layer, "quant_method", None)
|
||||
if isinstance(quant_method, QuantizeMethodBase):
|
||||
quant_method.process_weights_after_loading(layer)
|
||||
|
||||
# Copy processed values into original tensor storage (preserves cudagraph refs)
|
||||
# this code is a no-op if not reloading (because kernel tensors is empty)
|
||||
if info.kernel_tensors is not None:
|
||||
_copy_and_restore_kernel_tensors(layer, info)
|
||||
|
||||
info.reset()
|
||||
logger.debug("%s: Processed", layer.__class__.__name__)
|
||||
|
||||
|
||||
def _get_original_loader(tensor: torch.Tensor) -> Callable:
|
||||
"""Return the weight loader with any layerwise wrappers removed"""
|
||||
loader = _get_weight_loader(tensor)
|
||||
while loader.__name__ == "online_process_loader":
|
||||
loader = loader.__wrapped__ # type: ignore[union-attr]
|
||||
|
||||
return loader
|
||||
|
||||
|
||||
def _get_weight_loader(tensor: torch.Tensor):
|
||||
return getattr(tensor, "weight_loader", default_weight_loader)
|
||||
|
||||
|
||||
def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
|
||||
"""Copy processed values into original kernel tensor storage and restore
|
||||
kernel tensor references on the layer. Preserves cudagraph references."""
|
||||
assert info.kernel_tensors is not None
|
||||
parameters, buffers = info.kernel_tensors
|
||||
for name, param in parameters.items():
|
||||
param.data.copy_(getattr(layer, name))
|
||||
for name, buffer in buffers.items():
|
||||
if name not in layer._buffers:
|
||||
continue
|
||||
buffer.data.copy_(getattr(layer, name))
|
||||
|
||||
_place_kernel_tensors(layer, info)
|
||||
|
||||
|
||||
def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
|
||||
for name in get_layer_tensors(layer):
|
||||
delattr(layer, name)
|
||||
|
||||
assert info.kernel_tensors is not None
|
||||
parameters, buffers = info.kernel_tensors
|
||||
for name, param in parameters.items():
|
||||
layer.register_parameter(name, param)
|
||||
for name, buffer in buffers.items():
|
||||
layer.register_buffer(name, buffer)
|
||||
@@ -0,0 +1,202 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import UninitializedParameter
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
from .sanitize import restore_layer_refs, sanitize_layer_refs
|
||||
from .types import LayerReloadingInfo, LayerTensors
|
||||
from .utils import get_layer_params_buffers, get_layer_tensors
|
||||
|
||||
__all__ = [
|
||||
"to_meta_tensor",
|
||||
"materialize_meta_tensor",
|
||||
"capture_layer_to_meta",
|
||||
"restore_layer_on_meta",
|
||||
"materialize_layer",
|
||||
"get_numel_loaded",
|
||||
]
|
||||
|
||||
SKIP_MODULES: set[str] = {"HadamardTransform"}
|
||||
|
||||
SKIP_TENSORS: set[str] = {
|
||||
"_expert_map",
|
||||
"expert_mask",
|
||||
"expert_global_to_physical",
|
||||
"expert_physical_to_global",
|
||||
"expert_local_to_global",
|
||||
"e_score_correction_bias",
|
||||
}
|
||||
|
||||
|
||||
def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert a tensor to a meta tensor while preserving class and attributes."""
|
||||
meta_tensor = tensor.data.to("meta")
|
||||
meta_tensor.__class__ = tensor.__class__
|
||||
meta_tensor.__dict__ = tensor.__dict__.copy()
|
||||
return meta_tensor
|
||||
|
||||
|
||||
def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Materialize a meta tensor into an actual tensor on the current device.
|
||||
Should be called within the torch device context for the given rank.
|
||||
"""
|
||||
tensor = torch.empty_strided(
|
||||
size=tuple(meta_tensor.size()),
|
||||
stride=tuple(meta_tensor.stride()),
|
||||
dtype=meta_tensor.dtype,
|
||||
requires_grad=False,
|
||||
)
|
||||
tensor.__class__ = meta_tensor.__class__
|
||||
tensor.__dict__ = meta_tensor.__dict__.copy()
|
||||
return tensor
|
||||
|
||||
|
||||
def _is_non_persistent_parameter_alias_buffer(
|
||||
layer: torch.nn.Module,
|
||||
name: str,
|
||||
buffer: torch.Tensor,
|
||||
parameter_storage_ptrs: set[int],
|
||||
) -> bool:
|
||||
if name not in layer._non_persistent_buffers_set:
|
||||
return False
|
||||
|
||||
buffer_storage_ptr = _tensor_storage_ptr(buffer)
|
||||
return (
|
||||
buffer_storage_ptr is not None and buffer_storage_ptr in parameter_storage_ptrs
|
||||
)
|
||||
|
||||
|
||||
def _tensor_storage_ptr(tensor: torch.Tensor) -> int | None:
|
||||
if isinstance(tensor, UninitializedParameter):
|
||||
return None
|
||||
|
||||
try:
|
||||
return tensor.untyped_storage().data_ptr()
|
||||
except (RuntimeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _parameter_storage_ptrs(layer: torch.nn.Module) -> set[int]:
|
||||
return {
|
||||
storage_ptr
|
||||
for param in layer.parameters(recurse=True)
|
||||
if (storage_ptr := _tensor_storage_ptr(param)) is not None
|
||||
}
|
||||
|
||||
|
||||
def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors:
|
||||
if layer.__class__.__name__ in SKIP_MODULES:
|
||||
return ({}, {})
|
||||
|
||||
params, buffers = get_layer_params_buffers(layer)
|
||||
parameter_storage_ptrs = _parameter_storage_ptrs(layer)
|
||||
return (
|
||||
{
|
||||
name: sanitize_layer_refs(to_meta_tensor(param), layer)
|
||||
for name, param in params.items()
|
||||
if name not in SKIP_TENSORS
|
||||
},
|
||||
{
|
||||
name: sanitize_layer_refs(to_meta_tensor(buffer), layer)
|
||||
for name, buffer in buffers.items()
|
||||
if name not in SKIP_TENSORS
|
||||
and not _is_non_persistent_parameter_alias_buffer(
|
||||
layer, name, buffer, parameter_storage_ptrs
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo):
|
||||
"""Restore a layer to model format with tensors on the meta device"""
|
||||
if layer.__class__.__name__ in SKIP_MODULES:
|
||||
return
|
||||
|
||||
for name in get_layer_tensors(layer):
|
||||
if name not in SKIP_TENSORS:
|
||||
delattr(layer, name)
|
||||
|
||||
restore_params, restore_buffers = info.restore_metadata
|
||||
for name, param in restore_params.items():
|
||||
if name not in SKIP_TENSORS:
|
||||
param = restore_layer_refs(param, layer)
|
||||
layer.register_parameter(name, param)
|
||||
|
||||
for name, buffer in restore_buffers.items():
|
||||
if name not in SKIP_TENSORS:
|
||||
buffer = restore_layer_refs(buffer, layer)
|
||||
layer.register_buffer(name, buffer)
|
||||
|
||||
|
||||
def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo):
|
||||
"""Materialize all meta tensors in a layer to actual tensors."""
|
||||
if layer.__class__.__name__ in SKIP_MODULES:
|
||||
return
|
||||
|
||||
with info.restore_device:
|
||||
for name, tensor in get_layer_tensors(layer).items():
|
||||
if name not in SKIP_TENSORS and tensor.is_meta:
|
||||
setattr(layer, name, materialize_meta_tensor(tensor))
|
||||
|
||||
|
||||
class CopyCounter(TorchDispatchMode):
|
||||
"""
|
||||
Tracks total number of elements modified with `copy_`.
|
||||
|
||||
Useful for keeping track of weight loading where underlying weights can be
|
||||
arbitrarily transformed (such as with `narrow`) before calling copy.
|
||||
|
||||
Note: Assumes that copy kwargs are not used.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.copied_numel = 0
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
if func is torch.ops.aten.copy_.default:
|
||||
assert args[0].numel() == args[1].numel()
|
||||
self.copied_numel += args[0].numel()
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
|
||||
def get_numel_loaded(
|
||||
weight_loader: Callable, args: inspect.BoundArguments
|
||||
) -> tuple[int, object]:
|
||||
"""
|
||||
Determine how many elements would be loaded by a weight loader call.
|
||||
|
||||
Args:
|
||||
weight_loader: used to load weights
|
||||
args: bound arguments to weight loader
|
||||
|
||||
Returns:
|
||||
number of elements loaded by the weight loader, the return value of the
|
||||
weight loader
|
||||
"""
|
||||
with CopyCounter() as counter:
|
||||
return_value = weight_loader(*args.args, **args.kwargs)
|
||||
|
||||
# A weight loader fills a single destination parameter, so the number of
|
||||
# loaded elements is at most that parameter's size. Some loaders copy into
|
||||
# the parameter more than once -- e.g. ``composed_weight_loader`` runs an
|
||||
# in-place post-load transform (``param.copy_(fn(param))``) on top of the
|
||||
# initial copy -- which would make CopyCounter report twice the parameter
|
||||
# size. Over-counting inflates the layer's loaded-element total and can
|
||||
# finalize the layer before every parameter is loaded, silently dropping
|
||||
# the trailing parameter(s) (e.g. Mamba ``mixer.D``). Cap the count at the
|
||||
# destination size to keep the per-layer accounting correct.
|
||||
numel = counter.copied_numel
|
||||
param = args.arguments.get("param", None)
|
||||
if isinstance(param, torch.Tensor):
|
||||
numel = min(numel, param.numel())
|
||||
return numel, return_value
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from types import MethodType
|
||||
|
||||
import torch
|
||||
|
||||
__all__ = ["sanitize_layer_refs", "restore_layer_refs"]
|
||||
|
||||
|
||||
layer_ref_sentinel = object()
|
||||
|
||||
|
||||
def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor:
|
||||
"""
|
||||
Removes references to layer held by tensor attributes. Specifically, removes the
|
||||
`__self__` attribute of weight loader methods attached to the tensor.
|
||||
|
||||
Used by `capture_layer_to_meta` to avoid circular references to layers in
|
||||
`LAYERWISE_INFO`, leading to modules never being cleaned up. Without sanitation,
|
||||
tensors will reference layers, and the WeakKeyDictionary will never evict entries,
|
||||
even when the model is deleted.
|
||||
|
||||
Args:
|
||||
tensor: tensor to be sanitized
|
||||
layer: layer whose references should be removed
|
||||
|
||||
Returns:
|
||||
sanitized tensor
|
||||
"""
|
||||
for key, value in tensor.__dict__.items():
|
||||
if isinstance(value, MethodType) and value.__self__ is layer:
|
||||
tensor.__dict__[key] = value.__func__.__get__(layer_ref_sentinel)
|
||||
|
||||
return tensor
|
||||
|
||||
|
||||
def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Tensor:
|
||||
"""
|
||||
Restores references to layer held by tensor attributes.
|
||||
|
||||
Used by `restore_layer_on_meta` to add back layer references, allowing for proper
|
||||
weight loading.
|
||||
|
||||
Args:
|
||||
tensor: tensor to be sanitized
|
||||
layer: layer whose references should be removed
|
||||
|
||||
Returns:
|
||||
sanitized tensor
|
||||
"""
|
||||
for key, value in tensor.__dict__.items():
|
||||
if isinstance(value, MethodType) and value.__self__ is layer_ref_sentinel:
|
||||
tensor.__dict__[key] = value.__func__.__get__(layer)
|
||||
|
||||
return tensor
|
||||
@@ -0,0 +1,58 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Iterable
|
||||
from functools import wraps
|
||||
from types import FunctionType
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
|
||||
from .layerwise import (
|
||||
finalize_layerwise_reload,
|
||||
initialize_layerwise_reload,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
__all__ = ["set_torchao_reload_attrs", "support_quantized_model_reload_from_hp_weights"]
|
||||
|
||||
|
||||
def set_torchao_reload_attrs(model: torch.nn.Module, model_config: ModelConfig):
|
||||
model._do_torchao_reload = True
|
||||
model._model_config = model_config
|
||||
|
||||
|
||||
def support_quantized_model_reload_from_hp_weights(original_load_weights: FunctionType):
|
||||
"""
|
||||
Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
|
||||
reloading high precision (bfloat16/float16/float32) weight for an already quantized
|
||||
model, this involves restoring the weights to a high precision weights and
|
||||
then online quantize the weights.
|
||||
|
||||
Only applies to torchao quantized models. Assumes that all model weights are
|
||||
loaded within a single weights iterator (cannot perform batched updates)
|
||||
"""
|
||||
|
||||
@wraps(original_load_weights)
|
||||
def patched_model_load_weights(
|
||||
self: "AutoWeightsLoader",
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
model = self.module
|
||||
|
||||
if not getattr(model, "_do_torchao_reload", False):
|
||||
return original_load_weights(self, weights, *args, **kwargs)
|
||||
|
||||
initialize_layerwise_reload(model)
|
||||
loaded_weights = original_load_weights(self, weights, *args, **kwargs)
|
||||
finalize_layerwise_reload(model, model._model_config)
|
||||
|
||||
return loaded_weights
|
||||
|
||||
return patched_model_load_weights
|
||||
@@ -0,0 +1,38 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass, field
|
||||
from inspect import BoundArguments
|
||||
|
||||
import torch
|
||||
|
||||
__all__ = ["LayerTensors", "LayerReloadingInfo"]
|
||||
|
||||
# encodes both parameters and buffers separately
|
||||
LayerTensors = tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class LayerReloadingInfo:
|
||||
# model format metadata, recorded by `record_metadata_for_reloading`
|
||||
restore_metadata: LayerTensors
|
||||
|
||||
# device to materialize layers with, recorded by `record_metadata_for_reloading`
|
||||
restore_device: torch.device
|
||||
|
||||
# track how many elements are ready for loading, used by `online_process_loader`
|
||||
load_numel: int = 0
|
||||
load_numel_total: int | None = None
|
||||
|
||||
# used by `online_process_loader` to buffer args and tensors until ready to load
|
||||
loaded_weights: list[tuple[str, BoundArguments]] = field(default_factory=list)
|
||||
|
||||
# kernel formatted tensors, copied into by `_layerwise_process` when reloading
|
||||
kernel_tensors: LayerTensors | None = None
|
||||
|
||||
def reset(self):
|
||||
self.__init__( # type: ignore[misc]
|
||||
restore_metadata=self.restore_metadata, restore_device=self.restore_device
|
||||
)
|
||||
|
||||
def can_load(self) -> bool:
|
||||
return self.load_numel_total is not None
|
||||
@@ -0,0 +1,79 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from inspect import BoundArguments
|
||||
|
||||
import torch
|
||||
|
||||
from .types import LayerReloadingInfo, LayerTensors
|
||||
|
||||
__all__ = [
|
||||
"get_layer_tensors",
|
||||
"get_layer_params_buffers",
|
||||
"get_layer_size",
|
||||
"has_device_tensors",
|
||||
"get_info_size",
|
||||
]
|
||||
|
||||
|
||||
def get_layer_tensors(layer: torch.nn.Module) -> dict[str, torch.Tensor]:
|
||||
"""Get all parameters and buffers from a module as a dict."""
|
||||
params, buffers = get_layer_params_buffers(layer)
|
||||
return params | buffers
|
||||
|
||||
|
||||
def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
|
||||
"""Get all parameters and buffers of a module as a tuple of dicts."""
|
||||
return (
|
||||
{name: param for name, param in layer._parameters.items() if param is not None},
|
||||
{name: buffer for name, buffer in layer._buffers.items() if buffer is not None},
|
||||
)
|
||||
|
||||
|
||||
def get_layer_size(layer: torch.nn.Module) -> int:
|
||||
"""Calculate total number of elements across loadable tensors in a layer.
|
||||
|
||||
Excludes SKIP_TENSORS (e.g. _expert_map) which are never moved to meta
|
||||
device and never loaded via weight_loader during layerwise reload.
|
||||
"""
|
||||
from .meta import SKIP_TENSORS
|
||||
|
||||
return sum(
|
||||
tensor.numel()
|
||||
for name, tensor in get_layer_tensors(layer).items()
|
||||
if name not in SKIP_TENSORS
|
||||
)
|
||||
|
||||
|
||||
def has_device_tensors(bound_args: BoundArguments) -> bool:
|
||||
"""
|
||||
Return True if the loaded weights exist on an accelerator device
|
||||
|
||||
Args:
|
||||
bound_args: args to load weights
|
||||
|
||||
Returns:
|
||||
True if weights are on accelerator device
|
||||
"""
|
||||
return any(
|
||||
isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu")
|
||||
for value in bound_args.arguments.values()
|
||||
)
|
||||
|
||||
|
||||
def get_info_size(info: LayerReloadingInfo) -> int:
|
||||
"""
|
||||
Calculate the number of bytes used by loaded weights for a given layer
|
||||
|
||||
Args:
|
||||
info: layerwise info to get size of
|
||||
|
||||
Returns:
|
||||
number of bytes used by loaded weights
|
||||
"""
|
||||
return sum(
|
||||
value.nbytes
|
||||
for _, args in info.loaded_weights
|
||||
for value in args.arguments.values()
|
||||
if isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu")
|
||||
)
|
||||
Reference in New Issue
Block a user