| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from inspect import BoundArguments |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from .types import LayerReloadingInfo, LayerTensors |
| 9 | |
| 10 | __all__ = [ |
| 11 | "get_layer_tensors", |
| 12 | "get_layer_params_buffers", |
| 13 | "get_layer_size", |
| 14 | "get_tensor_load_numel", |
| 15 | "has_device_tensors", |
| 16 | "get_info_size", |
| 17 | ] |
| 18 | |
| 19 | |
| 20 | def get_layer_tensors(layer: torch.nn.Module) -> dict[str, torch.Tensor]: |
| 21 | """Get all parameters and buffers from a module as a dict.""" |
| 22 | params, buffers = get_layer_params_buffers(layer) |
| 23 | return params | buffers |
| 24 | |
| 25 | |
| 26 | def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors: |
| 27 | """Get all parameters and buffers of a module as a tuple of dicts.""" |
| 28 | return ( |
| 29 | {name: param for name, param in layer._parameters.items() if param is not None}, |
| 30 | {name: buffer for name, buffer in layer._buffers.items() if buffer is not None}, |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | def get_tensor_load_numel(tensor: torch.Tensor) -> int: |
| 35 | """Count checkpoint elements, excluding padding declared by the weight creator.""" |
| 36 | return getattr(tensor, "weight_loader_numel", tensor.numel()) |
| 37 | |
| 38 | |
| 39 | def get_layer_size(layer: torch.nn.Module) -> int: |
| 40 | """Calculate total number of elements across loadable tensors in a layer. |
| 41 | |
| 42 | Excludes SKIP_LOAD_TENSORS (e.g. _expert_map) which are never loaded via |
| 43 | weight_loader during layerwise reload. |
| 44 | """ |
| 45 | from .meta import SKIP_LOAD_TENSORS |
| 46 | |
| 47 | return sum( |
| 48 | get_tensor_load_numel(tensor) |
| 49 | for name, tensor in get_layer_tensors(layer).items() |
| 50 | if name not in SKIP_LOAD_TENSORS |
| 51 | ) |
| 52 | |
| 53 | |
| 54 | def has_device_tensors(bound_args: BoundArguments) -> bool: |
| 55 | """Return True if the loaded weights exist on an accelerator device. |
| 56 | |
| 57 | Args: |
| 58 | bound_args: args to load weights |
| 59 | |
| 60 | Returns: |
| 61 | True if weights are on accelerator device |
| 62 | |
| 63 | """ |
| 64 | return any( |
| 65 | isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu") |
| 66 | for value in bound_args.arguments.values() |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | def get_info_size(info: LayerReloadingInfo) -> int: |
| 71 | """Calculate the number of bytes used by loaded weights for a given layer. |
| 72 | |
| 73 | Args: |
| 74 | info: layerwise info to get size of |
| 75 | |
| 76 | Returns: |
| 77 | number of bytes used by loaded weights |
| 78 | |
| 79 | """ |
| 80 | return sum( |
| 81 | value.nbytes |
| 82 | for _, args in info.loaded_weights |
| 83 | for value in args.arguments.values() |
| 84 | if isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu") |
| 85 | ) |