vllm.model_executor.models.transformers.utils

Transformers modeling backend utilities.

Functions:

• attrsetter –

  Set a possibly nested attribute, like the inverse of attrgetter.

• can_enable_torch_compile –

  Callable to be passed to @support_torch_compile's enable_if argument.

• init_on_device_without_buffers –

  A context manager under which models are initialized with all

• maybe_per_layer –

  Pick layer_idx's entry from a config field that may be sized per layer.

• named_state –

  module's own state (i.e. named parameters and buffers).

• recursive_replace_linear –

  Recursively replace linear modules in the model as needed.

• replace_conv_class –

  Replace a Transformers Conv2d/Conv3d with vLLM's Conv2d/Conv3d.

• replace_embedding_class –

  Replace nn.Embedding with VocabParallelEmbedding.

• replace_layernorm_class –

  Replace a standard (mean-centered) nn.LayerNorm with vLLM's

• replace_linear_class –

  Replace nn.Linear with one of vLLM's tensor parallel linear classes.

_UninitializedEmbedding

Bases: Embedding

Make __init__ inert, so that VocabParallelEmbedding.__init__' call to super().__init__ does not invoke nn.Embedding.__init__.

Source code in vllm/model_executor/models/transformers/utils.py

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

_VocabParallelEmbeddingBase

Bases: VocabParallelEmbedding, _UninitializedEmbedding

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

Source code in vllm/model_executor/models/transformers/utils.py

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."""

_rebase_on_vocab_parallel(cls)

Subclass cls so that VocabParallelEmbedding supersedes its nn.Embedding.

Parameters:

• cls

  (type[Embedding]) –

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

Returns:

• type[VocabParallelEmbedding] –

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

Source code in vllm/model_executor/models/transformers/utils.py

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]

attrsetter(attr)

Set a possibly nested attribute, like the inverse of attrgetter.

Source code in vllm/model_executor/models/transformers/utils.py

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

can_enable_torch_compile(vllm_config)

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.

Source code in vllm/model_executor/models/transformers/utils.py

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

init_on_device_without_buffers(device)

A context manager under which models are initialized with all parameters on the specified device. However buffers are not initialized on specified device.

Parameters:

• device

  (`torch.device`) –

  Device to initialize all parameters on.

Source code in vllm/model_executor/models/transformers/utils.py

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

maybe_per_layer(value, layer_idx)

Pick layer_idx's entry from a config field that may be sized per layer.

Source code in vllm/model_executor/models/transformers/utils.py

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

named_state(module)

module's own state (i.e. named parameters and buffers).

Source code in vllm/model_executor/models/transformers/utils.py

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

recursive_replace_linear(model, quant_config, prefix='')

Recursively replace linear modules in the model as needed.

Source code in vllm/model_executor/models/transformers/utils.py

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)

replace_conv_class(conv)

Replace a Transformers Conv2d/Conv3d with vLLM's Conv2d/Conv3d.

Parameters:

• conv

  (TorchConv) –

  nn.Conv2d or nn.Conv3d to be replaced.

Returns:

• VllmConv | TorchConv –

  The new Conv2dLayer or Conv3dLayer. If the conv module is not supported,

• VllmConv | TorchConv –

  returns the original conv module.

Source code in vllm/model_executor/models/transformers/utils.py

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

replace_embedding_class(embedding, quant_config=None, *, prefix='')

Replace nn.Embedding with VocabParallelEmbedding.

Parameters:

• embedding

  (Embedding) –

  The nn.Embedding holding a vocab table.

• quant_config

  (QuantizationConfig | None, default: None ) –

  Quantization config for the new embedding.

• prefix

  (str, default: '' ) –

  Qualname of embedding, used to look up its quantization method.

Returns:

• Module –

  The module to install in embedding's place.

Source code in vllm/model_executor/models/transformers/utils.py

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

replace_layernorm_class(layernorm)

Replace a standard (mean-centered) nn.LayerNorm with vLLM's LayerNorm.

Parameters:

• layernorm

  (LayerNorm) –

  nn.LayerNorm to be replaced.

Returns:

• Module –

  The new LayerNorm. If the layernorm is not supported (a subclass with

• Module –

  its own behavior, multi-dim normalized_shape, or not both

• Module –

  elementwise-affine and biased), returns the original module unchanged.

Source code in vllm/model_executor/models/transformers/utils.py

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

replace_linear_class(linear, style='replicate', quant_config=None, *, prefix='')

Replace nn.Linear with one of vLLM's tensor parallel linear classes.

Parameters:

• linear

  (Linear) –

  nn.Linear to be replaced.

• style

  (Style, default: 'replicate' ) –

  Tensor parallel style of the new linear, e.g. "colwise".

• quant_config

  (QuantizationConfig | None, default: None ) –

  Quantization config for the new linear.

• prefix

  (str, default: '' ) –

  Module prefix of the new linear, used for quantization lookup.

Returns:

• ColumnParallelLinear | RowParallelLinear | ReplicatedLinear –

  The new linear.

Source code in vllm/model_executor/models/transformers/utils.py

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