vLLM/vllm/model_executor/models/transformers/utils.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# Copyright 2024 The vLLM team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Transformers modeling backend utilities."""

from collections.abc import Callable, Iterator
from contextlib import contextmanager
from itertools import chain
from operator import attrgetter
from pathlib import Path
from typing import TYPE_CHECKING, Literal, TypeAlias, TypeVar

import torch
from torch import nn

from vllm.logger import init_logger
from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
from vllm.model_executor.layers.layernorm import LayerNorm
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
from vllm.model_executor.models.utils import maybe_prefix
from vllm.transformers_utils.config import iter_rope_parameters

if TYPE_CHECKING:
    from vllm.config import VllmConfig
    from vllm.model_executor.layers.quantization import QuantizationConfig


logger = init_logger(__name__)


# Copied from `accelerate`
@contextmanager
def init_on_device_without_buffers(device: torch.device):
    """A context manager under which models are initialized with all
    parameters on the specified device. However buffers are not
    initialized on specified device.

    Args:
        device (`torch.device`):
            Device to initialize all parameters on.

    """
    old_register_parameter = nn.Module.register_parameter

    def register_empty_parameter(module, name, param):
        old_register_parameter(module, name, param)
        if param is not None:
            param_cls = type(module._parameters[name])
            kwargs = module._parameters[name].__dict__
            kwargs["requires_grad"] = param.requires_grad
            module._parameters[name] = param_cls(
                module._parameters[name].to(device), **kwargs
            )

    tensor_constructors_to_patch: dict[str, Callable] = {}

    def patch_tensor_constructor(fn):
        def wrapper(*args, **kwargs):
            kwargs["device"] = device
            return fn(*args, **kwargs)

        return wrapper

    try:
        nn.Module.register_parameter = register_empty_parameter
        for torch_function_name in tensor_constructors_to_patch:
            setattr(
                torch,
                torch_function_name,
                patch_tensor_constructor(getattr(torch, torch_function_name)),
            )
        yield
    finally:
        nn.Module.register_parameter = old_register_parameter
        for (
            torch_function_name,
            old_torch_function,
        ) in tensor_constructors_to_patch.items():
            setattr(torch, torch_function_name, old_torch_function)


Style = Literal[
    "colwise",
    "rowwise",
    "replicate",
    "colwise_gather_output",
    "rowwise_split_input",
]


def replace_linear_class(
    linear: nn.Linear,
    style: Style = "replicate",
    quant_config: "QuantizationConfig | None" = None,
    *,
    prefix: str = "",
) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear:
    """Replace nn.Linear with one of vLLM's tensor parallel linear classes.

    Args:
        linear: `nn.Linear` to be replaced.
        style: Tensor parallel style of the new linear, e.g. "colwise".
        quant_config: Quantization config for the new linear.
        prefix: Module prefix of the new linear, used for quantization lookup.

    Returns:
        The new linear.

    """
    if not isinstance(style, str):
        raise ValueError(f"Unsupported parallel style type {type(style)}, expected str")

    vllm_linear_cls, vllm_linear_kwargs = {
        "colwise": (ColumnParallelLinear, {}),
        "rowwise": (RowParallelLinear, {}),
        "replicate": (ReplicatedLinear, {}),
        "colwise_gather_output": (ColumnParallelLinear, {"gather_output": True}),
        "rowwise_split_input": (RowParallelLinear, {"input_is_parallel": False}),
    }.get(style, (ReplicatedLinear, {}))

    return vllm_linear_cls(
        input_size=linear.in_features,
        output_size=linear.out_features,
        bias=linear.bias is not None,
        quant_config=quant_config,
        prefix=prefix,
        return_bias=False,
        **vllm_linear_kwargs,
    )


TorchConv: TypeAlias = nn.Conv2d | nn.Conv3d
VllmConv: TypeAlias = Conv2dLayer | Conv3dLayer


def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
    """Replace a Transformers Conv2d/Conv3d with vLLM's Conv2d/Conv3d.

    Args:
        conv: `nn.Conv2d` or `nn.Conv3d` to be replaced.

    Returns:
        The new `Conv2dLayer` or `Conv3dLayer`. If the conv module is not supported,
        returns the original conv module.

    """
    # vLLM does not handle non-zero padding modes
    if conv.padding_mode != "zeros":
        return conv

    vllm_conv_cls = {
        nn.Conv2d: Conv2dLayer,
        nn.Conv3d: Conv3dLayer,
    }.get(type(conv))

    if vllm_conv_cls is None:
        return conv

    return vllm_conv_cls(
        in_channels=conv.in_channels,
        out_channels=conv.out_channels,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        dilation=conv.dilation,
        groups=conv.groups,
        bias=conv.bias is not None,
        padding_mode=conv.padding_mode,
        params_dtype=conv.weight.dtype,
    )


def replace_layernorm_class(layernorm: nn.LayerNorm) -> nn.Module:
    """Replace a standard (mean-centered) `nn.LayerNorm` with vLLM's
    `LayerNorm`.

    Args:
        layernorm: `nn.LayerNorm` to be replaced.

    Returns:
        The new `LayerNorm`. If the layernorm is not supported (a subclass with
        its own behavior, multi-dim `normalized_shape`, or not both
        elementwise-affine and biased), returns the original module unchanged.

    """
    if (
        type(layernorm) is not nn.LayerNorm
        or len(layernorm.normalized_shape) != 1
        or not layernorm.elementwise_affine
        or layernorm.bias is None
    ):
        return layernorm

    return LayerNorm(
        layernorm.normalized_shape[0],
        eps=layernorm.eps,
        dtype=layernorm.weight.dtype,
    )


def attrsetter(attr: str) -> Callable[[object, object], None]:
    """Set a possibly nested attribute, like the inverse of attrgetter."""
    parent, _, name = attr.rpartition(".")

    def setter(obj: object, value: object):
        attr_parent = attrgetter(parent)(obj) if parent else obj
        setattr(attr_parent, name, value)

    return setter


class _UninitializedEmbedding(nn.Embedding):
    """Make `__init__` inert, so that `VocabParallelEmbedding.__init__`'
    call to `super().__init__` does not invoke `nn.Embedding.__init__`."""

    def __init__(self):
        pass


class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbedding):
    """Orders `VocabParallelEmbedding` ahead of `nn.Embedding` in the MRO, so that
    `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding."""


_rebased_embedding_classes: dict[type[nn.Embedding], type[VocabParallelEmbedding]] = {}


def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbedding]:
    """Subclass `cls` so that `VocabParallelEmbedding` supersedes its `nn.Embedding`.

    Args:
        cls: The `nn.Embedding` subclass to rebase. Cached, so a given `cls` always
            maps to the same class.

    Returns:
        The new class, to assign to `__class__` of an instance of `cls`.

    """
    if cls not in _rebased_embedding_classes:
        rebased_cls = type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {})
        _rebased_embedding_classes[cls] = rebased_cls
    return _rebased_embedding_classes[cls]


def replace_embedding_class(
    embedding: nn.Embedding,
    quant_config: "QuantizationConfig | None" = None,
    *,
    prefix: str = "",
) -> nn.Module:
    """Replace `nn.Embedding` with `VocabParallelEmbedding`.

    Args:
        embedding: The `nn.Embedding` holding a vocab table.
        quant_config: Quantization config for the new embedding.
        prefix: Qualname of `embedding`, used to look up its quantization method.

    Returns:
        The module to install in `embedding`'s place.

    """
    # Tied tables are reached more than once, so don't rebase an already new embedding
    if isinstance(embedding, VocabParallelEmbedding):
        return embedding

    kwargs = dict(
        num_embeddings=embedding.num_embeddings,
        embedding_dim=embedding.embedding_dim,
        params_dtype=embedding.weight.dtype,
        quant_config=quant_config,
        prefix=prefix,
    )
    # If `embedding` is a bare `nn.Embedding`, simple replace
    if type(embedding) is nn.Embedding:
        return VocabParallelEmbedding(**kwargs)

    # Otherwise `embedding` inherits `nn.Embedding`, rebase it in place
    embedding.__class__ = _rebase_on_vocab_parallel(type(embedding))
    VocabParallelEmbedding.__init__(embedding, **kwargs)
    return embedding


def recursive_replace_linear(
    model: nn.Module,
    quant_config: "QuantizationConfig | None",
    prefix: str = "",
):
    """Recursively replace linear modules in the model as needed."""

    def _recursive_replace(module: nn.Module, prefix: str):
        for child_name, child_module in module.named_children():
            new_module = child_module
            qual_name = maybe_prefix(prefix, child_name)
            # Replace modules as needed
            if isinstance(child_module, nn.Linear):
                new_module = replace_linear_class(
                    child_module,
                    "replicate",
                    quant_config,
                    prefix=qual_name,
                )
            else:
                _recursive_replace(child_module, prefix=qual_name)
            if new_module is not child_module:
                setattr(module, child_name, new_module)

    _recursive_replace(model, prefix=prefix)


T = TypeVar("T")


def maybe_per_layer(value: T | list[T], layer_idx: int) -> T:
    """Pick `layer_idx`'s entry from a config field that may be sized per layer."""
    return value[layer_idx] if isinstance(value, list) else value


def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
    """`module`'s own state (i.e. named parameters and buffers)."""
    return chain(module.named_parameters(), module.named_buffers())


def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module):
    logger.debug("%s: %s -> %s", name, old_module, new_module)


def get_feature_request_tip(
    model: str,
    trust_remote_code: bool,
) -> str:
    hf_url = f"a discussion at https://huggingface.co/{model}/discussions/new"
    gh_url = "an issue at https://github.com/huggingface/transformers/issues/new/choose"
    url = hf_url if trust_remote_code else gh_url
    prefix = f"Please open {url} to request support for this feature. "
    if Path(model).exists():
        prefix = ""
    doc_url = "https://docs.vllm.ai/en/latest/models/supported_models.html#writing-custom-models"
    tip = f"See {doc_url} for instructions on how to add support yourself."
    return f"{prefix}{tip}"


def dynamic_rope_update_is_noop(vllm_config: "VllmConfig") -> bool:
    """Whether `max_model_len` keeps Transformers' dynamic rope from ever rescaling.

    `@dynamic_rope_update` only recomputes `inv_freq` once positions grow past
    `max_position_embeddings`, which vLLM never feeds beyond `max_model_len`.
    """
    text_config = vllm_config.model_config.hf_config.get_text_config()
    max_position_embeddings = getattr(text_config, "max_position_embeddings", None)
    if max_position_embeddings is None:
        return False
    return vllm_config.model_config.max_model_len <= max_position_embeddings


def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool:
    """Callable to be passed to `@support_torch_compile`'s `enable_if` argument.

    Defaults to `True` but is disabled in the following situations:

    - The model uses dynamic rope scaling that can still rescale within
      `max_model_len` (otherwise `disable_noop_dynamic_rope` removes it).
    """
    text_config = vllm_config.model_config.hf_config.get_text_config()
    rope_types = {rp["rope_type"] for rp in iter_rope_parameters(text_config)}
    if "dynamic" not in rope_types:
        return True
    # Layer types using longrope still sync, even once dynamic rope is disabled
    return "longrope" not in rope_types and dynamic_rope_update_is_noop(vllm_config)


def disable_noop_dynamic_rope(model: nn.Module, vllm_config: "VllmConfig") -> None:
    """Skip `@dynamic_rope_update` in rotary modules where it can never rescale.

    The update syncs on `torch.max(position_ids)` every forward, which breaks
    torch.compile and CUDA graph capture. Relabelling the module's `rope_type` as
    "default" skips it while keeping the `inv_freq` computed at init.
    """
    if not dynamic_rope_update_is_noop(vllm_config):
        return
    for module in model.modules():
        rope_type = getattr(module, "rope_type", None)
        if not hasattr(module, "max_seq_len_cached") or rope_type is None:
            continue
        if isinstance(rope_type, dict):
            module.rope_type = {
                k: "default" if v == "dynamic" else v for k, v in rope_type.items()
            }
        elif rope_type == "dynamic":
            module.rope_type = "default"