vllm.model_executor.model_loader.reload.meta

Functions:

• get_numel_loaded –

  Determine how many elements would be loaded by a weight loader call.

• materialize_layer –

  Materialize all meta tensors in a layer to actual tensors.

• materialize_meta_tensor –

  Materialize a meta tensor into an actual tensor on the current device.

• restore_layer_on_meta –

  Restore a layer to model format with tensors on the meta device.

• to_meta_tensor –

  Convert a tensor to a meta tensor while preserving class and attributes.

CopyCounter

Bases: 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.

Source code in vllm/model_executor/model_loader/reload/meta.py

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)

get_numel_loaded(weight_loader, args)

Determine how many elements would be loaded by a weight loader call.

Parameters:

• weight_loader

  (Callable) –

  used to load weights

• args

  (BoundArguments) –

  bound arguments to weight loader

Returns:

• int –

  number of elements loaded by the weight loader, the return value of the

• object –

  weight loader

Source code in vllm/model_executor/model_loader/reload/meta.py

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, get_tensor_load_numel(param))
    return numel, return_value

materialize_layer(layer, info)

Materialize all meta tensors in a layer to actual tensors.

Source code in vllm/model_executor/model_loader/reload/meta.py

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))

materialize_meta_tensor(meta_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.

Source code in vllm/model_executor/model_loader/reload/meta.py

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

restore_layer_on_meta(layer, info)

Restore a layer to model format with tensors on the meta device.

Source code in vllm/model_executor/model_loader/reload/meta.py

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

    non_persistent = set(layer._non_persistent_buffers_set)
    restore_params, restore_buffers = info.restore_metadata
    tensor_names = (
        get_layer_tensors(layer).keys() | restore_params.keys() | restore_buffers.keys()
    )
    for name in tensor_names:
        if name not in SKIP_TENSORS and hasattr(layer, name):
            delattr(layer, name)

    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, persistent=name not in non_persistent)

to_meta_tensor(tensor)

Convert a tensor to a meta tensor while preserving class and attributes.

Source code in vllm/model_executor/model_loader/reload/meta.py

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