| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2024 The vLLM team. |
| 5 | # |
| 6 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 7 | # you may not use this file except in compliance with the License. |
| 8 | # You may obtain a copy of the License at |
| 9 | # |
| 10 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | # |
| 12 | # Unless required by applicable law or agreed to in writing, software |
| 13 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | # See the License for the specific language governing permissions and |
| 16 | # limitations under the License. |
| 17 | """Transformers modeling backend utilities.""" |
| 18 | |
| 19 | from collections.abc import Callable, Iterator |
| 20 | from contextlib import contextmanager |
| 21 | from itertools import chain |
| 22 | from operator import attrgetter |
| 23 | from pathlib import Path |
| 24 | from typing import TYPE_CHECKING, Literal, TypeAlias, TypeVar |
| 25 | |
| 26 | import torch |
| 27 | from torch import nn |
| 28 | |
| 29 | from vllm.logger import init_logger |
| 30 | from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer |
| 31 | from vllm.model_executor.layers.layernorm import LayerNorm |
| 32 | from vllm.model_executor.layers.linear import ( |
| 33 | ColumnParallelLinear, |
| 34 | ReplicatedLinear, |
| 35 | RowParallelLinear, |
| 36 | ) |
| 37 | from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding |
| 38 | from vllm.model_executor.models.utils import maybe_prefix |
| 39 | from vllm.transformers_utils.config import is_rope_parameters_nested |
| 40 | |
| 41 | if TYPE_CHECKING: |
| 42 | from vllm.config import VllmConfig |
| 43 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 44 | |
| 45 | |
| 46 | logger = init_logger(__name__) |
| 47 | |
| 48 | |
| 49 | # Copied from `accelerate` |
| 50 | @contextmanager |
| 51 | def init_on_device_without_buffers(device: torch.device): |
| 52 | """A context manager under which models are initialized with all |
| 53 | parameters on the specified device. However buffers are not |
| 54 | initialized on specified device. |
| 55 | |
| 56 | Args: |
| 57 | device (`torch.device`): |
| 58 | Device to initialize all parameters on. |
| 59 | |
| 60 | """ |
| 61 | old_register_parameter = nn.Module.register_parameter |
| 62 | |
| 63 | def register_empty_parameter(module, name, param): |
| 64 | old_register_parameter(module, name, param) |
| 65 | if param is not None: |
| 66 | param_cls = type(module._parameters[name]) |
| 67 | kwargs = module._parameters[name].__dict__ |
| 68 | kwargs["requires_grad"] = param.requires_grad |
| 69 | module._parameters[name] = param_cls( |
| 70 | module._parameters[name].to(device), **kwargs |
| 71 | ) |
| 72 | |
| 73 | tensor_constructors_to_patch: dict[str, Callable] = {} |
| 74 | |
| 75 | def patch_tensor_constructor(fn): |
| 76 | def wrapper(*args, **kwargs): |
| 77 | kwargs["device"] = device |
| 78 | return fn(*args, **kwargs) |
| 79 | |
| 80 | return wrapper |
| 81 | |
| 82 | try: |
| 83 | nn.Module.register_parameter = register_empty_parameter |
| 84 | for torch_function_name in tensor_constructors_to_patch: |
| 85 | setattr( |
| 86 | torch, |
| 87 | torch_function_name, |
| 88 | patch_tensor_constructor(getattr(torch, torch_function_name)), |
| 89 | ) |
| 90 | yield |
| 91 | finally: |
| 92 | nn.Module.register_parameter = old_register_parameter |
| 93 | for ( |
| 94 | torch_function_name, |
| 95 | old_torch_function, |
| 96 | ) in tensor_constructors_to_patch.items(): |
| 97 | setattr(torch, torch_function_name, old_torch_function) |
| 98 | |
| 99 | |
| 100 | Style = Literal[ |
| 101 | "colwise", |
| 102 | "rowwise", |
| 103 | "replicate", |
| 104 | "colwise_gather_output", |
| 105 | "rowwise_split_input", |
| 106 | ] |
| 107 | |
| 108 | |
| 109 | def replace_linear_class( |
| 110 | linear: nn.Linear, |
| 111 | style: Style = "replicate", |
| 112 | quant_config: "QuantizationConfig | None" = None, |
| 113 | *, |
| 114 | prefix: str = "", |
| 115 | ) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear: |
| 116 | """Replace nn.Linear with one of vLLM's tensor parallel linear classes. |
| 117 | |
| 118 | Args: |
| 119 | linear: `nn.Linear` to be replaced. |
| 120 | style: Tensor parallel style of the new linear, e.g. "colwise". |
| 121 | quant_config: Quantization config for the new linear. |
| 122 | prefix: Module prefix of the new linear, used for quantization lookup. |
| 123 | |
| 124 | Returns: |
| 125 | The new linear. |
| 126 | |
| 127 | """ |
| 128 | if not isinstance(style, str): |
| 129 | raise ValueError(f"Unsupported parallel style type {type(style)}, expected str") |
| 130 | |
| 131 | vllm_linear_cls, vllm_linear_kwargs = { |
| 132 | "colwise": (ColumnParallelLinear, {}), |
| 133 | "rowwise": (RowParallelLinear, {}), |
| 134 | "replicate": (ReplicatedLinear, {}), |
| 135 | "colwise_gather_output": (ColumnParallelLinear, {"gather_output": True}), |
| 136 | "rowwise_split_input": (RowParallelLinear, {"input_is_parallel": False}), |
| 137 | }.get(style, (ReplicatedLinear, {})) |
| 138 | |
| 139 | return vllm_linear_cls( |
| 140 | input_size=linear.in_features, |
| 141 | output_size=linear.out_features, |
| 142 | bias=linear.bias is not None, |
| 143 | quant_config=quant_config, |
| 144 | prefix=prefix, |
| 145 | return_bias=False, |
| 146 | **vllm_linear_kwargs, |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | TorchConv: TypeAlias = nn.Conv2d | nn.Conv3d |
| 151 | VllmConv: TypeAlias = Conv2dLayer | Conv3dLayer |
| 152 | |
| 153 | |
| 154 | def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv: |
| 155 | """Replace a Transformers Conv2d/Conv3d with vLLM's Conv2d/Conv3d. |
| 156 | |
| 157 | Args: |
| 158 | conv: `nn.Conv2d` or `nn.Conv3d` to be replaced. |
| 159 | |
| 160 | Returns: |
| 161 | The new `Conv2dLayer` or `Conv3dLayer`. If the conv module is not supported, |
| 162 | returns the original conv module. |
| 163 | |
| 164 | """ |
| 165 | # vLLM does not handle non-zero padding modes |
| 166 | if conv.padding_mode != "zeros": |
| 167 | return conv |
| 168 | |
| 169 | vllm_conv_cls = { |
| 170 | nn.Conv2d: Conv2dLayer, |
| 171 | nn.Conv3d: Conv3dLayer, |
| 172 | }.get(type(conv)) |
| 173 | |
| 174 | if vllm_conv_cls is None: |
| 175 | return conv |
| 176 | |
| 177 | return vllm_conv_cls( |
| 178 | in_channels=conv.in_channels, |
| 179 | out_channels=conv.out_channels, |
| 180 | kernel_size=conv.kernel_size, |
| 181 | stride=conv.stride, |
| 182 | padding=conv.padding, |
| 183 | dilation=conv.dilation, |
| 184 | groups=conv.groups, |
| 185 | bias=conv.bias is not None, |
| 186 | padding_mode=conv.padding_mode, |
| 187 | params_dtype=conv.weight.dtype, |
| 188 | ) |
| 189 | |
| 190 | |
| 191 | def replace_layernorm_class(layernorm: nn.LayerNorm) -> nn.Module: |
| 192 | """Replace a standard (mean-centered) `nn.LayerNorm` with vLLM's |
| 193 | `LayerNorm`. |
| 194 | |
| 195 | Args: |
| 196 | layernorm: `nn.LayerNorm` to be replaced. |
| 197 | |
| 198 | Returns: |
| 199 | The new `LayerNorm`. If the layernorm is not supported (a subclass with |
| 200 | its own behavior, multi-dim `normalized_shape`, or not both |
| 201 | elementwise-affine and biased), returns the original module unchanged. |
| 202 | |
| 203 | """ |
| 204 | if ( |
| 205 | type(layernorm) is not nn.LayerNorm |
| 206 | or len(layernorm.normalized_shape) != 1 |
| 207 | or not layernorm.elementwise_affine |
| 208 | or layernorm.bias is None |
| 209 | ): |
| 210 | return layernorm |
| 211 | |
| 212 | return LayerNorm( |
| 213 | layernorm.normalized_shape[0], |
| 214 | eps=layernorm.eps, |
| 215 | dtype=layernorm.weight.dtype, |
| 216 | ) |
| 217 | |
| 218 | |
| 219 | def attrsetter(attr: str) -> Callable[[object, object], None]: |
| 220 | """Set a possibly nested attribute, like the inverse of attrgetter.""" |
| 221 | parent, _, name = attr.rpartition(".") |
| 222 | |
| 223 | def setter(obj: object, value: object): |
| 224 | attr_parent = attrgetter(parent)(obj) if parent else obj |
| 225 | setattr(attr_parent, name, value) |
| 226 | |
| 227 | return setter |
| 228 | |
| 229 | |
| 230 | class _UninitializedEmbedding(nn.Embedding): |
| 231 | """Make `__init__` inert, so that `VocabParallelEmbedding.__init__`' |
| 232 | call to `super().__init__` does not invoke `nn.Embedding.__init__`.""" |
| 233 | |
| 234 | def __init__(self): |
| 235 | pass |
| 236 | |
| 237 | |
| 238 | class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbedding): |
| 239 | """Orders `VocabParallelEmbedding` ahead of `nn.Embedding` in the MRO, so that |
| 240 | `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding.""" |
| 241 | |
| 242 | |
| 243 | _rebased_embedding_classes: dict[type[nn.Embedding], type[VocabParallelEmbedding]] = {} |
| 244 | |
| 245 | |
| 246 | def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbedding]: |
| 247 | """Subclass `cls` so that `VocabParallelEmbedding` supersedes its `nn.Embedding`. |
| 248 | |
| 249 | Args: |
| 250 | cls: The `nn.Embedding` subclass to rebase. Cached, so a given `cls` always |
| 251 | maps to the same class. |
| 252 | |
| 253 | Returns: |
| 254 | The new class, to assign to `__class__` of an instance of `cls`. |
| 255 | |
| 256 | """ |
| 257 | if cls not in _rebased_embedding_classes: |
| 258 | rebased_cls = type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {}) |
| 259 | _rebased_embedding_classes[cls] = rebased_cls |
| 260 | return _rebased_embedding_classes[cls] |
| 261 | |
| 262 | |
| 263 | def replace_embedding_class( |
| 264 | embedding: nn.Embedding, |
| 265 | quant_config: "QuantizationConfig | None" = None, |
| 266 | *, |
| 267 | prefix: str = "", |
| 268 | ) -> nn.Module: |
| 269 | """Replace `nn.Embedding` with `VocabParallelEmbedding`. |
| 270 | |
| 271 | Args: |
| 272 | embedding: The `nn.Embedding` holding a vocab table. |
| 273 | quant_config: Quantization config for the new embedding. |
| 274 | prefix: Qualname of `embedding`, used to look up its quantization method. |
| 275 | |
| 276 | Returns: |
| 277 | The module to install in `embedding`'s place. |
| 278 | |
| 279 | """ |
| 280 | # Tied tables are reached more than once, so don't rebase an already new embedding |
| 281 | if isinstance(embedding, VocabParallelEmbedding): |
| 282 | return embedding |
| 283 | |
| 284 | kwargs = dict( |
| 285 | num_embeddings=embedding.num_embeddings, |
| 286 | embedding_dim=embedding.embedding_dim, |
| 287 | params_dtype=embedding.weight.dtype, |
| 288 | quant_config=quant_config, |
| 289 | prefix=prefix, |
| 290 | ) |
| 291 | # If `embedding` is a bare `nn.Embedding`, simple replace |
| 292 | if type(embedding) is nn.Embedding: |
| 293 | return VocabParallelEmbedding(**kwargs) |
| 294 | |
| 295 | # Otherwise `embedding` inherits `nn.Embedding`, rebase it in place |
| 296 | embedding.__class__ = _rebase_on_vocab_parallel(type(embedding)) |
| 297 | VocabParallelEmbedding.__init__(embedding, **kwargs) |
| 298 | return embedding |
| 299 | |
| 300 | |
| 301 | def recursive_replace_linear( |
| 302 | model: nn.Module, |
| 303 | quant_config: "QuantizationConfig | None", |
| 304 | prefix: str = "", |
| 305 | ): |
| 306 | """Recursively replace linear modules in the model as needed.""" |
| 307 | |
| 308 | def _recursive_replace(module: nn.Module, prefix: str): |
| 309 | for child_name, child_module in module.named_children(): |
| 310 | new_module = child_module |
| 311 | qual_name = maybe_prefix(prefix, child_name) |
| 312 | # Replace modules as needed |
| 313 | if isinstance(child_module, nn.Linear): |
| 314 | new_module = replace_linear_class( |
| 315 | child_module, |
| 316 | "replicate", |
| 317 | quant_config, |
| 318 | prefix=qual_name, |
| 319 | ) |
| 320 | else: |
| 321 | _recursive_replace(child_module, prefix=qual_name) |
| 322 | if new_module is not child_module: |
| 323 | setattr(module, child_name, new_module) |
| 324 | |
| 325 | _recursive_replace(model, prefix=prefix) |
| 326 | |
| 327 | |
| 328 | T = TypeVar("T") |
| 329 | |
| 330 | |
| 331 | def maybe_per_layer(value: T | list[T], layer_idx: int) -> T: |
| 332 | """Pick `layer_idx`'s entry from a config field that may be sized per layer.""" |
| 333 | return value[layer_idx] if isinstance(value, list) else value |
| 334 | |
| 335 | |
| 336 | def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: |
| 337 | """`module`'s own state (i.e. named parameters and buffers).""" |
| 338 | return chain(module.named_parameters(), module.named_buffers()) |
| 339 | |
| 340 | |
| 341 | def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module): |
| 342 | logger.debug("%s: %s -> %s", name, old_module, new_module) |
| 343 | |
| 344 | |
| 345 | def get_feature_request_tip( |
| 346 | model: str, |
| 347 | trust_remote_code: bool, |
| 348 | ) -> str: |
| 349 | hf_url = f"a discussion at https://huggingface.co/{model}/discussions/new" |
| 350 | gh_url = "an issue at https://github.com/huggingface/transformers/issues/new/choose" |
| 351 | url = hf_url if trust_remote_code else gh_url |
| 352 | prefix = f"Please open {url} to request support for this feature. " |
| 353 | if Path(model).exists(): |
| 354 | prefix = "" |
| 355 | doc_url = "https://docs.vllm.ai/en/latest/models/supported_models.html#writing-custom-models" |
| 356 | tip = f"See {doc_url} for instructions on how to add support yourself." |
| 357 | return f"{prefix}{tip}" |
| 358 | |
| 359 | |
| 360 | def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool: |
| 361 | """Callable to be passed to `@support_torch_compile`'s `enable_if` argument. |
| 362 | |
| 363 | Defaults to `True` but is disabled in the following situations: |
| 364 | |
| 365 | - The model uses dynamic rope scaling. |
| 366 | """ |
| 367 | text_config = vllm_config.model_config.hf_config.get_text_config() |
| 368 | # Dynamic rope scaling is not compatible with torch.compile |
| 369 | rope_parameters: dict | None = getattr(text_config, "rope_parameters", None) or {} |
| 370 | if rope_parameters: |
| 371 | # Nest rope_parameters if not nested already to simplify logic |
| 372 | if not is_rope_parameters_nested(rope_parameters): |
| 373 | rope_parameters = {"": rope_parameters} |
| 374 | return all(rp["rope_type"] != "dynamic" for rp in rope_parameters.values()) |
| 375 | return True |