# 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.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 is_rope_parameters_nested

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 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 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.
    """
    text_config = vllm_config.model_config.hf_config.get_text_config()
    # Dynamic rope scaling is not compatible with torch.compile
    rope_parameters: dict | None = getattr(text_config, "rope_parameters", None) or {}
    if rope_parameters:
        # Nest rope_parameters if not nested already to simplify logic
        if not is_rope_parameters_nested(rope_parameters):
            rope_parameters = {"": rope_parameters}
        return all(rp["rope_type"] != "dynamic" for rp in rope_parameters.values())
    return True