| 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 base class.""" |
| 18 | |
| 19 | import os |
| 20 | from collections.abc import Callable, Iterable |
| 21 | from contextlib import contextmanager |
| 22 | from functools import cached_property |
| 23 | from itertools import chain |
| 24 | from operator import attrgetter |
| 25 | from typing import TYPE_CHECKING, Any, NamedTuple |
| 26 | |
| 27 | import regex as re |
| 28 | import torch |
| 29 | import transformers |
| 30 | from packaging.version import Version |
| 31 | from torch import nn |
| 32 | from transformers import AutoModel, PreTrainedConfig |
| 33 | from transformers.conversion_mapping import ( |
| 34 | WeightRenaming, |
| 35 | get_checkpoint_conversion_mapping, |
| 36 | get_model_conversion_mapping, |
| 37 | ) |
| 38 | |
| 39 | from vllm.compilation.decorators import DynamicArgDims, support_torch_compile |
| 40 | from vllm.config.utils import getattr_iter |
| 41 | from vllm.distributed import get_pp_group, get_tp_group |
| 42 | from vllm.distributed.utils import get_pp_indices |
| 43 | from vllm.logger import init_logger |
| 44 | from vllm.model_executor.layers.attention import ( |
| 45 | Attention, |
| 46 | EncoderOnlyAttention, |
| 47 | MLAAttention, |
| 48 | ) |
| 49 | from vllm.model_executor.layers.attention.mla_attention import get_mla_dims |
| 50 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 51 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 52 | from vllm.model_executor.models.interfaces import ( |
| 53 | SupportsEagle, |
| 54 | SupportsEagle3, |
| 55 | SupportsLoRA, |
| 56 | SupportsPP, |
| 57 | SupportsQuant, |
| 58 | ) |
| 59 | from vllm.model_executor.models.interfaces_base import VllmModel |
| 60 | from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers |
| 61 | from vllm.model_executor.models.transformers.fusers import AttentionFuser, MLAFuser |
| 62 | from vllm.model_executor.models.transformers.fusers.attention import ( |
| 63 | VLLM_ATTN_IMPL, |
| 64 | VLLM_MLA_ATTN_IMPL, |
| 65 | ) |
| 66 | from vllm.model_executor.models.transformers.utils import ( |
| 67 | attrsetter, |
| 68 | can_enable_torch_compile, |
| 69 | get_feature_request_tip, |
| 70 | init_on_device_without_buffers, |
| 71 | log_replacement, |
| 72 | named_state, |
| 73 | replace_conv_class, |
| 74 | replace_embedding_class, |
| 75 | replace_layernorm_class, |
| 76 | replace_linear_class, |
| 77 | ) |
| 78 | from vllm.model_executor.models.utils import ( |
| 79 | AutoWeightsLoader, |
| 80 | PPMissingLayer, |
| 81 | ShardId, |
| 82 | WeightsMapper, |
| 83 | make_empty_intermediate_tensors_factory, |
| 84 | maybe_prefix, |
| 85 | ) |
| 86 | from vllm.sequence import IntermediateTensors |
| 87 | |
| 88 | if TYPE_CHECKING: |
| 89 | from transformers import PreTrainedModel |
| 90 | |
| 91 | from vllm.config import VllmConfig |
| 92 | |
| 93 | logger = init_logger(__name__) |
| 94 | |
| 95 | VLLM_ATTN_ATTR = "attn" |
| 96 | """Attribute in which to store the vLLM `Attention` instance that serves a Transformers |
| 97 | attention module. `"attn"` is chosen deliberately to match the mappings generated by |
| 98 | `maybe_remap_kv_scale_name`.""" |
| 99 | |
| 100 | |
| 101 | class PreTrainedModelClasses(NamedTuple): |
| 102 | decoder: type["PreTrainedModel"] |
| 103 | encoders: dict[str, type["PreTrainedModel"]] |
| 104 | """Modality -> encoder class, for each modality that has one.""" |
| 105 | |
| 106 | |
| 107 | class Base( |
| 108 | nn.Module, |
| 109 | VllmModel, |
| 110 | SupportsQuant, |
| 111 | SupportsLoRA, |
| 112 | SupportsPP, |
| 113 | SupportsEagle, |
| 114 | SupportsEagle3, |
| 115 | ): |
| 116 | hf_to_vllm_mapper: WeightsMapper |
| 117 | |
| 118 | # TODO transformers will have a util to get these |
| 119 | embedding_modules = {"embed_tokens": "input_embeddings"} |
| 120 | |
| 121 | def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): |
| 122 | nn.Module.__init__(self) |
| 123 | logger.info("Using Transformers modeling backend.") |
| 124 | |
| 125 | self.vllm_config = vllm_config |
| 126 | self.config = vllm_config.model_config.hf_config |
| 127 | self.text_config = self.config.get_text_config() |
| 128 | self.cache_config = vllm_config.cache_config |
| 129 | self.compilation_config = vllm_config.compilation_config |
| 130 | self.device_config = vllm_config.device_config |
| 131 | self.model_config = vllm_config.model_config |
| 132 | self.parallel_config = vllm_config.parallel_config |
| 133 | self.quant_config = vllm_config.quant_config |
| 134 | |
| 135 | self.pp_group = get_pp_group() |
| 136 | self.tp_group = get_tp_group() |
| 137 | |
| 138 | # Attrs for weight loading (see self.load_weights) |
| 139 | self.ignore_unexpected_prefixes: list[str] = [] |
| 140 | """Ignore unexpected weights whose qualname starts with these prefixes.""" |
| 141 | self.ignore_unexpected_suffixes: list[str] = [] |
| 142 | """Ignore unexpected weights whose qualname ends with these suffixes.""" |
| 143 | self.packed_modules_mapping: dict[str, list[str]] = {} |
| 144 | """Fused module -> constituent projections, populated by `recursive_replace` |
| 145 | for the quantization machinery and loaders (e.g. bitsandbytes).""" |
| 146 | self.fusers: dict[str, list[BaseFuser]] = {} |
| 147 | """Module qualname -> the fusers applied to it, populated |
| 148 | by `recursive_replace` for `_create_attention_instances`.""" |
| 149 | self.attention_fusers: dict[int, tuple[str, AttentionFuser]] = {} |
| 150 | """`layer_idx` -> the qualname and fuser of the module computing that |
| 151 | layer's attention, populated by `recursive_replace` for |
| 152 | `_create_attention_instances`.""" |
| 153 | |
| 154 | # Attrs for Eagle3 (see self.set_aux_hidden_state_layers) |
| 155 | self._target_class: type[nn.Module] = nn.Module |
| 156 | """Target class for Eagle3 aux hidden state recording.""" |
| 157 | self._layer_names: dict[int, str] = {} |
| 158 | """Mapping from layer index to layer name for Eagle3.""" |
| 159 | self._output_aux_hidden_states_kwargs: dict[str, bool] = {} |
| 160 | """Kwargs to pass to model forward for Eagle3 aux hidden states.""" |
| 161 | |
| 162 | if self.quant_config: |
| 163 | quant_method_name = self.quant_config.get_name() |
| 164 | # Check for unsupported quantization methods. |
| 165 | if quant_method_name in ("mxfp4", "gpt_oss_mxfp4"): |
| 166 | raise NotImplementedError( |
| 167 | "Transformers modeling backend does " |
| 168 | "not support MXFP4 quantization yet." |
| 169 | ) |
| 170 | |
| 171 | self._patch_config() |
| 172 | self._decorate_for_torch_compile() |
| 173 | # Init on "meta" to delay allocating GPU tensors |
| 174 | with ( |
| 175 | self._mark_model_components(vllm_config), |
| 176 | init_on_device_without_buffers("meta"), |
| 177 | ): |
| 178 | from_config_kwargs = self._from_config_kwargs |
| 179 | self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs) |
| 180 | |
| 181 | # Create weight name to module qualname mapper |
| 182 | self._create_hf_to_vllm_mapper() |
| 183 | # Remove layers not on this pipeline parallel rank |
| 184 | self.pipeline_parallel() |
| 185 | # Substitute remaining layers with vLLM's layers as needed |
| 186 | self.recursive_replace() |
| 187 | # Create attention instances for KV cache allocation |
| 188 | self._create_attention_instances(prefix) |
| 189 | |
| 190 | # Initialize any parameters that have not had their modules replaced |
| 191 | self.init_parameters(self.model) |
| 192 | |
| 193 | # Upcast weights Transformers always keeps in fp32 |
| 194 | self.keep_in_fp32(self.model) |
| 195 | |
| 196 | # Pipeline parallel intermediate tensors |
| 197 | self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( |
| 198 | ["hidden_states"], self.text_config.hidden_size |
| 199 | ) |
| 200 | |
| 201 | def _patch_config(self): |
| 202 | """Patch the config to ensure that the model is created correctly: |
| 203 | |
| 204 | - Sets the attention implementation to "vllm" so the attention instances from |
| 205 | `_create_attention_instances` are used |
| 206 | - Sets the dtype to the default torch dtype set by vLLM because Transformers |
| 207 | uses the config dtype when creating the model |
| 208 | """ |
| 209 | self.text_config._attn_implementation = VLLM_ATTN_IMPL |
| 210 | self.config.dtype = torch.get_default_dtype() |
| 211 | |
| 212 | @contextmanager |
| 213 | def _mark_model_components(self, vllm_config: "VllmConfig"): |
| 214 | """Mark language model and tower submodules as `self.model` is created. |
| 215 | |
| 216 | Nothing to do in `Base`, `MultiModalMixin` will override.""" |
| 217 | yield |
| 218 | |
| 219 | @cached_property |
| 220 | def _from_config_kwargs(self) -> dict[str, Any]: |
| 221 | """The kwargs used to create `self.model`.""" |
| 222 | return dict( |
| 223 | config=self.config, |
| 224 | dtype=self.model_config.dtype, |
| 225 | trust_remote_code=self.model_config.trust_remote_code, |
| 226 | ) |
| 227 | |
| 228 | def _find_encoder_classes( |
| 229 | self, model: "PreTrainedModel" |
| 230 | ) -> dict[str, type["PreTrainedModel"]]: |
| 231 | """Find the encoder class of each modality `model` has one for. |
| 232 | |
| 233 | Text models have none. Multi-modal ones override this. |
| 234 | """ |
| 235 | return {} |
| 236 | |
| 237 | @cached_property |
| 238 | def _pre_trained_model_classes(self) -> PreTrainedModelClasses: |
| 239 | """The decoder class and the encoder class of each modality that has one. |
| 240 | |
| 241 | Both come from a single throwaway model, since building one is not cheap. |
| 242 | """ |
| 243 | with torch.device("meta"): |
| 244 | model: PreTrainedModel = AutoModel.from_config(**self._from_config_kwargs) |
| 245 | model_classes = PreTrainedModelClasses( |
| 246 | decoder=type(model.get_decoder()), |
| 247 | encoders=self._find_encoder_classes(model), |
| 248 | ) |
| 249 | del model |
| 250 | |
| 251 | logger.debug("Identified model classes as: %s", model_classes) |
| 252 | return model_classes |
| 253 | |
| 254 | def _decorate_cls_for_torch_compile( |
| 255 | self, |
| 256 | cls: type["PreTrainedModel"], |
| 257 | dynamic_arg_dims: DynamicArgDims | None, |
| 258 | enable_if: Callable[["VllmConfig"], bool], |
| 259 | is_encoder: bool, |
| 260 | ): |
| 261 | """Decorate `cls` to indicate to vLLM that it supports torch compile. |
| 262 | |
| 263 | Args: |
| 264 | cls: The PreTrainedModel class to decorate. |
| 265 | dynamic_arg_dims: A mapping from argument name to the dynamic dimensions |
| 266 | of the argument. If None, default dynamic arg dims will be used. See |
| 267 | [`support_torch_compile`][vllm.compilation.decorators.support_torch_compile] |
| 268 | for more details. |
| 269 | enable_if: A function which takes in the vLLM config and returns whether |
| 270 | torch compile should be enabled for this class. |
| 271 | is_encoder: Whether the class being decorated is an encoder. |
| 272 | |
| 273 | """ |
| 274 | logger.debug( |
| 275 | "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s", |
| 276 | cls.__name__, |
| 277 | "encoder" if is_encoder else "decoder", |
| 278 | dynamic_arg_dims, |
| 279 | ) |
| 280 | |
| 281 | support_torch_compile( |
| 282 | dynamic_arg_dims=dynamic_arg_dims, |
| 283 | enable_if=enable_if, |
| 284 | is_encoder=is_encoder, |
| 285 | )(cls) |
| 286 | |
| 287 | def _decorate_for_torch_compile(self): |
| 288 | """Decorate the model's decoder class to indicate to vLLM that it |
| 289 | supports torch compile if `can_enable_torch_compile` is True.""" |
| 290 | self._decorate_cls_for_torch_compile( |
| 291 | cls=self._pre_trained_model_classes.decoder, |
| 292 | # Applied to a PreTrainedModel so the batch dimension will exist |
| 293 | dynamic_arg_dims=DynamicArgDims( |
| 294 | input_ids=1, # shape: [1, seq_len] |
| 295 | inputs_embeds=1, # shape: [1, seq_len, hidden_size] |
| 296 | position_ids=-1, # shape: [1, seq_len] or [3, 1, seq_len] for mrope |
| 297 | ), |
| 298 | enable_if=can_enable_torch_compile, |
| 299 | is_encoder=False, |
| 300 | ) |
| 301 | |
| 302 | def _create_hf_to_vllm_mapper(self): |
| 303 | """Create a WeightsMapper to map checkpoint weight names to module qualnames. |
| 304 | |
| 305 | This handles: |
| 306 | |
| 307 | - Transformers weight renaming from `WeightRenaming` |
| 308 | - Checkpoints saved with a base model prefix that is not `model` |
| 309 | - Checkpoints saved with no base model prefix |
| 310 | - Any quantization config specific mappings |
| 311 | """ |
| 312 | orig_to_new_renaming: list[WeightRenaming] = [] |
| 313 | orig_to_new_regex: dict[re.Pattern, str | None] = {} |
| 314 | |
| 315 | mappings = get_model_conversion_mapping(self.model) |
| 316 | # Ensure layers that are part of the model, but not AutoModel, get their mapping |
| 317 | instantiated = {type(m).__name__ for m in self.model.modules()} |
| 318 | for arch in self.config.architectures or []: |
| 319 | if arch not in instantiated: |
| 320 | mappings.extend(get_checkpoint_conversion_mapping(arch) or []) |
| 321 | |
| 322 | for mapping in mappings: |
| 323 | # Handle weights which have been renamed in Transformers |
| 324 | if isinstance(mapping, WeightRenaming): |
| 325 | orig_to_new_renaming.append(mapping) |
| 326 | # TODO: Handle WeightConverter to enable layer merging |
| 327 | |
| 328 | # Handle unexpected weights which should be ignored |
| 329 | if self.model._keys_to_ignore_on_load_unexpected is not None: |
| 330 | for key in self.model._keys_to_ignore_on_load_unexpected: |
| 331 | orig_to_new_regex[re.compile(key)] = None |
| 332 | |
| 333 | # Standardise base model prefix |
| 334 | bmp = self.model.base_model_prefix |
| 335 | expected_bmp = r"model.\1" |
| 336 | # Handle checkpoints saved with different base model prefix |
| 337 | if bmp and bmp != "model": |
| 338 | different_bmp_pattern = re.compile(rf"^{bmp}\.(.+)") |
| 339 | orig_to_new_regex[different_bmp_pattern] = expected_bmp |
| 340 | # Handle direct children of self.model which were saved without the model prefix |
| 341 | direct_children = chain( |
| 342 | self.model.named_children(), |
| 343 | self.model.named_parameters(recurse=False), |
| 344 | self.model.named_buffers(recurse=False), |
| 345 | ) |
| 346 | model_children = "|".join(name for name, _ in direct_children) |
| 347 | missing_bmp_pattern = re.compile(rf"^(?!model\.)(({model_children}).*)") |
| 348 | orig_to_new_regex[missing_bmp_pattern] = expected_bmp |
| 349 | # Handle weights saved as direct children of self.model which no longer are |
| 350 | unexpected_bmp_pattern = re.compile(rf"^(model\.)((?!{model_children}).+)") |
| 351 | orig_to_new_regex[unexpected_bmp_pattern] = r"\2" |
| 352 | # Handle lm_head which was saved inside the base model |
| 353 | nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)") |
| 354 | orig_to_new_regex[nested_lm_head_pattern] = r"\2" |
| 355 | |
| 356 | self.hf_to_vllm_mapper = WeightsMapper( |
| 357 | orig_to_new_renaming=orig_to_new_renaming, |
| 358 | orig_to_new_regex=orig_to_new_regex, |
| 359 | ) |
| 360 | |
| 361 | # Apply mapping to quantization config if needed |
| 362 | self._maybe_apply_model_mapping() |
| 363 | |
| 364 | def _get_tie_word_embeddings(self): |
| 365 | """Check if the model has tied word embeddings.""" |
| 366 | # Models created with Transformers v4 and v5 will store this in different places |
| 367 | tie_word_embeddings_v4 = getattr(self.text_config, "tie_word_embeddings", False) |
| 368 | tie_word_embeddings_v5 = getattr(self.config, "tie_word_embeddings", False) |
| 369 | return tie_word_embeddings_v4 or tie_word_embeddings_v5 |
| 370 | |
| 371 | def pipeline_parallel(self): |
| 372 | """Apply the model's pipeline parallelization plan.""" |
| 373 | if self.pp_group.world_size <= 1: |
| 374 | return |
| 375 | |
| 376 | if self.model.supports_pp_plan: |
| 377 | module = self.model |
| 378 | names = list(module._pp_plan.keys()) |
| 379 | else: |
| 380 | module = self.model.get_decoder() |
| 381 | has_parameters = lambda m: next(m.parameters(), None) is not None |
| 382 | names = [n for n, c in module.named_children() if has_parameters(c)] |
| 383 | tip = get_feature_request_tip( |
| 384 | self.model_config.model, self.model_config.trust_remote_code |
| 385 | ) |
| 386 | logger.warning_once( |
| 387 | "%s does not define a pipeline parallel plan. The Transformers " |
| 388 | "modeling backend will infer the split from the layers of %s in order " |
| 389 | "of declaration and keep parameter-free modules on every rank. This " |
| 390 | "may fail if the model's structure is non-standard. %s", |
| 391 | type(self.model).__name__, |
| 392 | type(module).__name__, |
| 393 | tip, |
| 394 | ) |
| 395 | |
| 396 | module_lists = [] |
| 397 | module_list_idx = None |
| 398 | for i, name in enumerate(names): |
| 399 | # attrgetter in case the module is nested (e.g. "text_model.layers") |
| 400 | if isinstance(attrgetter(name)(module), nn.ModuleList): |
| 401 | module_lists.append(name) |
| 402 | module_list_idx = i |
| 403 | |
| 404 | if len(module_lists) > 1: |
| 405 | raise ValueError( |
| 406 | "Pipeline parallel of models with multiple `ModuleList`s " |
| 407 | "in the base model are not supported yet!" |
| 408 | ) |
| 409 | if module_list_idx is None: |
| 410 | raise ValueError(f"Could not find `ModuleList` in {type(module)}") |
| 411 | |
| 412 | # Layers before module list |
| 413 | for name in names[:module_list_idx]: |
| 414 | if self.pp_group.is_first_rank or ( |
| 415 | self._get_tie_word_embeddings() and self.pp_group.is_last_rank |
| 416 | ): |
| 417 | continue |
| 418 | # attrsetter in case the module is nested (e.g. "text_model.embed_tokens") |
| 419 | attrsetter(name)(module, PPMissingLayer()) |
| 420 | |
| 421 | # Module list |
| 422 | start_layer, end_layer = get_pp_indices( |
| 423 | self.text_config.num_hidden_layers, |
| 424 | self.pp_group.rank_in_group, |
| 425 | self.pp_group.world_size, |
| 426 | ) |
| 427 | layers_name = names[module_list_idx] |
| 428 | # attrgetter in case the module is nested (e.g. "text_model.layers") |
| 429 | layers = attrgetter(layers_name)(module) |
| 430 | for i in range(len(layers)): |
| 431 | if start_layer <= i and i < end_layer: |
| 432 | continue |
| 433 | layers[i] = PPMissingLayer() |
| 434 | |
| 435 | # Layers after module list |
| 436 | for name in names[module_list_idx + 1 :]: |
| 437 | # Modules that should be on last rank |
| 438 | if not self.pp_group.is_last_rank: |
| 439 | # attrsetter in case the module is nested (e.g. "text_model.norm") |
| 440 | attrsetter(name)(module, PPMissingLayer()) |
| 441 | |
| 442 | def _vocab_embeddings(self) -> set[nn.Embedding]: |
| 443 | """The `nn.Embedding`s in `self.model` which hold a vocab table.""" |
| 444 | |
| 445 | def vocab_sizes(config: PreTrainedConfig): |
| 446 | for key, value in vars(config).items(): |
| 447 | if isinstance(value, PreTrainedConfig): |
| 448 | yield from vocab_sizes(value) |
| 449 | elif "vocab_size" in key and isinstance(value, int): |
| 450 | yield value |
| 451 | |
| 452 | sizes = set(vocab_sizes(self.config)) |
| 453 | # `get_input_embeddings` may return a module which composes the `nn.Embedding`, |
| 454 | # or a `PPMissingLayer` if the embeddings are not on this pipeline stage |
| 455 | embeddings = { |
| 456 | module |
| 457 | for module in self.model.get_input_embeddings().modules() |
| 458 | if isinstance(module, nn.Embedding) |
| 459 | } |
| 460 | embeddings.update( |
| 461 | module |
| 462 | for module in self.model.modules() |
| 463 | if isinstance(module, nn.Embedding) and module.num_embeddings in sizes |
| 464 | ) |
| 465 | return embeddings |
| 466 | |
| 467 | def recursive_replace(self): |
| 468 | """Recursively replace modules in the model as needed. |
| 469 | |
| 470 | Currently, this replaces: |
| 471 | |
| 472 | - GLUs with a fused `MergedColumnParallelLinear` + `...AndMul` |
| 473 | - Attention QKV projections with a fused `QKVParallelLinear` + split |
| 474 | - `nn.Linear` with vLLM's tensor parallel linear classes |
| 475 | - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d` |
| 476 | - Vocab `nn.Embedding`s with vLLM's `VocabParallelEmbedding` |
| 477 | - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm` |
| 478 | - `nn.LayerNorm` with vLLM's `LayerNorm` |
| 479 | """ |
| 480 | tp_plan = self.model.tp_plan or {} |
| 481 | |
| 482 | if not tp_plan and self.tp_group.world_size > 1: |
| 483 | tip = get_feature_request_tip( |
| 484 | self.model_config.model, self.model_config.trust_remote_code |
| 485 | ) |
| 486 | logger.warning_once( |
| 487 | "%s does not define a tensor parallel plan. The Transformers modeling " |
| 488 | "backend will shard the model the best it can during graph fusion and " |
| 489 | "replicate the rest. This may be suboptimal or fail if the model does " |
| 490 | "not fuse cleanly. %s", |
| 491 | type(self.model).__name__, |
| 492 | tip, |
| 493 | ) |
| 494 | |
| 495 | # Prefix the patterns because we always start from `self.model` |
| 496 | tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()} |
| 497 | # Detect fusable patterns once per module class (cached, so this is cheap) |
| 498 | fusers = Fusers(self.model, self.vllm_config) |
| 499 | |
| 500 | vocab_embeddings = self._vocab_embeddings() |
| 501 | |
| 502 | orig_to_new_stacked: dict[str, tuple[str, ShardId]] = {} |
| 503 | |
| 504 | def register_fusion(fuser: BaseFuser, prefix: str, module: nn.Module): |
| 505 | """Register a fused layer's mappings just before it is built.""" |
| 506 | self.fusers.setdefault(prefix, []).append(fuser) |
| 507 | |
| 508 | if ( |
| 509 | isinstance(fuser, AttentionFuser) |
| 510 | and (index := fuser.layer_index(module)) is not None |
| 511 | ): |
| 512 | if index in self.attention_fusers: |
| 513 | other, _ = self.attention_fusers[index] |
| 514 | raise ValueError( |
| 515 | f"Layer {index} dispatches attention from more than one " |
| 516 | f"module ({other} and {prefix}), which one vLLM attention " |
| 517 | "layer cannot serve." |
| 518 | ) |
| 519 | self.attention_fusers[index] = (prefix, fuser) |
| 520 | |
| 521 | orig_to_new_stacked.update(fuser.orig_to_new_stacked(prefix)) |
| 522 | |
| 523 | packed_modules_mapping = fuser.packed_modules_mapping |
| 524 | self.packed_modules_mapping.update(packed_modules_mapping) |
| 525 | if self.quant_config is not None: |
| 526 | self.quant_config.packed_modules_mapping.update(packed_modules_mapping) |
| 527 | |
| 528 | def _recursive_replace(module: nn.Module, prefix: str): |
| 529 | for child_name, child_module in module.named_children(): |
| 530 | new_module = child_module |
| 531 | qual_name = maybe_prefix(prefix, child_name) |
| 532 | if ( |
| 533 | isinstance(module, nn.ModuleList) |
| 534 | and len(module) == self.text_config.num_hidden_layers |
| 535 | ): |
| 536 | # Populate Eagle3 attrs |
| 537 | self._target_class = type(child_module) |
| 538 | layer_name = qual_name.removeprefix("model.") |
| 539 | self._layer_names[int(child_name)] = layer_name |
| 540 | # MTP weights should not be loaded into the base model |
| 541 | num_hidden_layers = self.text_config.num_hidden_layers |
| 542 | names = ( |
| 543 | "n_predict", # Override from SpeculativeConfig |
| 544 | "num_nextn_predict_layers", # Most models |
| 545 | "mtp_num_hidden_layers", # Qwen 3.5 |
| 546 | ) |
| 547 | n_predict = getattr_iter(self.text_config, names, 0) |
| 548 | for i in range(num_hidden_layers, num_hidden_layers + n_predict): |
| 549 | mtp_prefix = f"{prefix}.{i}." |
| 550 | if mtp_prefix not in self.ignore_unexpected_prefixes: |
| 551 | self.ignore_unexpected_prefixes.append(mtp_prefix) |
| 552 | # Replace modules as needed |
| 553 | if isinstance(child_module, nn.Linear): |
| 554 | generator = (p for p in tp_plan if re.match(p, qual_name)) |
| 555 | pattern = next(generator, None) |
| 556 | # Some weight loaders expect all linear layers to inherit |
| 557 | # LinearBase, so we set a default style which causes any |
| 558 | # unspecified layers to be replaced with ReplicatedLinear |
| 559 | style = tp_plan.get(pattern, "replicate") |
| 560 | new_module = replace_linear_class( |
| 561 | child_module, style, self.quant_config, prefix=qual_name |
| 562 | ) |
| 563 | elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)): |
| 564 | new_module = replace_conv_class(child_module) |
| 565 | elif child_module in vocab_embeddings: |
| 566 | new_module = replace_embedding_class( |
| 567 | child_module, self.quant_config, prefix=qual_name |
| 568 | ) |
| 569 | elif isinstance(child_module, nn.LayerNorm): |
| 570 | new_module = replace_layernorm_class(child_module) |
| 571 | elif child_module_fusers := fusers[child_module]: |
| 572 | for fuser in child_module_fusers: |
| 573 | register_fusion(fuser, qual_name, child_module) |
| 574 | new_module = fuser.fuse(new_module, qual_name, self.vllm_config) |
| 575 | logger.info_once(fuser.info(child_name)) |
| 576 | _recursive_replace(new_module, prefix=qual_name) |
| 577 | elif not isinstance(child_module, MoERunner): |
| 578 | # MoERunner can contain aliases of shared experts and gates, |
| 579 | # so we don't want to recurse into it and break weight loading. |
| 580 | _recursive_replace(child_module, prefix=qual_name) |
| 581 | |
| 582 | if new_module is not child_module: |
| 583 | setattr(module, child_name, new_module) |
| 584 | log_replacement(qual_name, child_module, new_module) |
| 585 | |
| 586 | _recursive_replace(self.model, prefix="model") |
| 587 | |
| 588 | self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked) |
| 589 | |
| 590 | def _create_attention_instances(self, prefix: str = ""): |
| 591 | """Create `Attention` instances to inform KV cache allocation.""" |
| 592 | text_config = self.text_config |
| 593 | attn_cls = self._get_attn_cls() |
| 594 | |
| 595 | # kv_lora_rank indicates that this is an MLA model |
| 596 | if getattr(text_config, "kv_lora_rank", None) is not None: |
| 597 | if attn_cls is MLAAttention: |
| 598 | text_config._attn_implementation = VLLM_MLA_ATTN_IMPL |
| 599 | else: |
| 600 | # MLA model not using MLAAttention: recompute head_size for full attn |
| 601 | qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", 0) |
| 602 | qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", 0) |
| 603 | if qk_head_dim := qk_nope_head_dim + qk_rope_head_dim: |
| 604 | self.model_config.model_arch_config.head_size = qk_head_dim |
| 605 | |
| 606 | logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None) |
| 607 | layer_types = getattr(text_config, "layer_types", None) |
| 608 | |
| 609 | pp_rank = self.pp_group.rank_in_group |
| 610 | pp_size = self.pp_group.world_size |
| 611 | start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size) |
| 612 | |
| 613 | for i in range(start, end): |
| 614 | if i not in self.attention_fusers: |
| 615 | layer = ( |
| 616 | f"{i} ({layer_types[i]})" |
| 617 | if layer_types and i < len(layer_types) |
| 618 | else str(i) |
| 619 | ) |
| 620 | raise ValueError( |
| 621 | f"Layer {layer} does not dispatch through the Transformers " |
| 622 | "attention interface and vLLM has no other way to handle it." |
| 623 | ) |
| 624 | attn_prefix, attn_fuser = self.attention_fusers[i] |
| 625 | attn_module = self.get_submodule(attn_prefix) |
| 626 | |
| 627 | # `[i]` is the whole-model config unless the checkpoint is |
| 628 | # heterogeneous, in which case it is this layer's own geometry. |
| 629 | arch_config = self.model_config.model_arch_config[i] |
| 630 | num_heads = self.model_config.get_num_attention_heads( |
| 631 | self.parallel_config, arch_config |
| 632 | ) |
| 633 | head_size = arch_config.head_size |
| 634 | if (scale := attn_fuser.scale(attn_module)) is None: |
| 635 | # Default to Llama scale if AttentionFuser couldn't identify it |
| 636 | scale = head_size**-0.5 |
| 637 | |
| 638 | kwargs = dict( |
| 639 | num_heads=num_heads, |
| 640 | scale=scale, |
| 641 | cache_config=self.cache_config, |
| 642 | quant_config=self.quant_config, |
| 643 | prefix=maybe_prefix(prefix, f"{attn_prefix}.{VLLM_ATTN_ATTR}"), |
| 644 | ) |
| 645 | |
| 646 | if attn_cls is MLAAttention: |
| 647 | mla_fuser = next( |
| 648 | (f for f in self.fusers[attn_prefix] if isinstance(f, MLAFuser)), |
| 649 | None, |
| 650 | ) |
| 651 | if mla_fuser is None: |
| 652 | raise ValueError( |
| 653 | f"Layer {i} ({attn_prefix}) did not fuse as MLA, but the rest " |
| 654 | "of the model did, so its latent KV cache cannot be allocated." |
| 655 | ) |
| 656 | dims = get_mla_dims(self.model_config) |
| 657 | kwargs.update( |
| 658 | qk_nope_head_dim=dims.qk_nope_head_dim, |
| 659 | qk_rope_head_dim=dims.qk_rope_head_dim, |
| 660 | v_head_dim=dims.v_head_dim, |
| 661 | q_lora_rank=dims.q_lora_rank, |
| 662 | kv_lora_rank=dims.kv_lora_rank, |
| 663 | kv_b_proj=attn_module.get_submodule(mla_fuser.kv_b_proj_name), |
| 664 | ) |
| 665 | else: |
| 666 | kwargs.update( |
| 667 | head_size=head_size, |
| 668 | num_kv_heads=self.model_config.get_num_kv_heads( |
| 669 | self.parallel_config, arch_config |
| 670 | ), |
| 671 | logits_soft_cap=logits_soft_cap, |
| 672 | ) |
| 673 | |
| 674 | # Handle interleaved sliding window attention |
| 675 | if layer_types and layer_types[i] == "sliding_attention": |
| 676 | kwargs["per_layer_sliding_window"] = text_config.sliding_window |
| 677 | # Handle attention sinks |
| 678 | if (sinks := attn_fuser.sinks(attn_module)) is not None: |
| 679 | kwargs["sinks"] = sinks |
| 680 | |
| 681 | attn_instance = attn_cls(**kwargs) |
| 682 | |
| 683 | # Attach the instance to the module it serves, which is the only way it |
| 684 | # is in the module tree at all: a plain dict is invisible to |
| 685 | # `nn.Module.__setattr__`. Being in the tree runs its |
| 686 | # `process_weights_after_loading` and moves its KV-scale buffers with |
| 687 | # `model.to(device)`, and lets the attention interface read it off the |
| 688 | # module rather than index by `module.layer_idx`, an `int` that bakes |
| 689 | # layer identity into the traced graph and so costs one compiled |
| 690 | # artifact per layer. |
| 691 | setattr(attn_module, VLLM_ATTN_ATTR, attn_instance) |
| 692 | |
| 693 | def _get_attn_cls(self) -> type[AttentionLayerBase]: |
| 694 | """Return the `Attention` class to use for this model's layers.""" |
| 695 | # In encoder models, the attention layers will have `is_causal=False` |
| 696 | is_encoder = lambda module: not getattr(module, "is_causal", True) |
| 697 | has_encoder = lambda model: any(is_encoder(m) for m in model.modules()) |
| 698 | is_multimodal = lambda config: config != config.get_text_config() |
| 699 | # vLLM does not support encoder-decoder models, so if any encoder layer is |
| 700 | # found in a text only model, we assume the whole model is an encoder model |
| 701 | if has_encoder(self.model) and not is_multimodal(self.config): |
| 702 | return EncoderOnlyAttention |
| 703 | if self.model_config.use_mla: |
| 704 | if any( |
| 705 | isinstance(fuser, MLAFuser) |
| 706 | for fusers in self.fusers.values() |
| 707 | for fuser in fusers |
| 708 | ): |
| 709 | return MLAAttention |
| 710 | logger.warning_once( |
| 711 | "This model uses MLA but `MLAFuser` failed to match and/or fuse any " |
| 712 | "MLA attention layers. Falling back to full attention with a padded " |
| 713 | "`value` head dimension." |
| 714 | ) |
| 715 | os.environ["VLLM_MLA_DISABLE"] = "1" |
| 716 | return Attention |
| 717 | |
| 718 | def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None): |
| 719 | """If a `parameter` is on the `meta` device, then its parent |
| 720 | `module` is the original module created by: |
| 721 | |
| 722 | ```python |
| 723 | with torch.device("meta"): |
| 724 | self.model: "PreTrainedModel" = AutoModel.from_config(...) |
| 725 | ``` |
| 726 | """ |
| 727 | dtype = dtype or self.model_config.dtype |
| 728 | device = self.device_config.device |
| 729 | |
| 730 | def _init_parameters(module: nn.Module): |
| 731 | for name, param in module.named_parameters(recurse=False): |
| 732 | # Already on device, nothing to do |
| 733 | if param.device != torch.device("meta"): |
| 734 | continue |
| 735 | # Already a vLLM parameter, nothing to do |
| 736 | if hasattr(param, "weight_loader"): |
| 737 | continue |
| 738 | data = torch.empty_like(param.data, dtype=dtype, device=device) |
| 739 | setattr(module, name, nn.Parameter(data=data)) |
| 740 | for child in module.children(): |
| 741 | _init_parameters(child) |
| 742 | |
| 743 | _init_parameters(module) |
| 744 | |
| 745 | def keep_in_fp32(self, module: nn.Module): |
| 746 | """Honor `_keep_in_fp32_modules_strict` as `from_pretrained` would.""" |
| 747 | if self.model_config.dtype not in (torch.float16, torch.bfloat16): |
| 748 | return |
| 749 | fragments = getattr(module, "_keep_in_fp32_modules_strict", None) |
| 750 | if not fragments: |
| 751 | return |
| 752 | pattern = re.compile("|".join(re.escape(f) for f in fragments)) |
| 753 | for name, tensor in named_state(module): |
| 754 | if not hasattr(tensor, "weight_loader") and pattern.search(name): |
| 755 | tensor.data = tensor.data.to(torch.float32) |
| 756 | |
| 757 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 758 | return self.model.get_input_embeddings()(input_ids) |
| 759 | |
| 760 | def forward( |
| 761 | self, |
| 762 | input_ids: torch.Tensor | None, |
| 763 | positions: torch.Tensor, |
| 764 | intermediate_tensors: IntermediateTensors | None = None, |
| 765 | inputs_embeds: torch.Tensor | None = None, |
| 766 | **kwargs, |
| 767 | ) -> torch.Tensor | IntermediateTensors: |
| 768 | if not self.pp_group.is_first_rank: |
| 769 | assert intermediate_tensors is not None |
| 770 | input_ids = None |
| 771 | inputs_embeds = intermediate_tensors["hidden_states"] |
| 772 | |
| 773 | # Add batch dimension before entering Transformers model |
| 774 | if input_ids is not None and input_ids.ndim == 1: |
| 775 | # [seq_len] -> [1, seq_len] |
| 776 | input_ids = input_ids[None, ...] |
| 777 | if inputs_embeds is not None and inputs_embeds.ndim == 2: |
| 778 | # [seq_len, hidden_size] -> [1, seq_len, hidden_size] |
| 779 | inputs_embeds = inputs_embeds[None, ...] |
| 780 | if positions.ndim == 1: |
| 781 | # [seq_len] -> [1, seq_len] |
| 782 | positions = positions[None, ...] |
| 783 | |
| 784 | # Transformers models expect either input_ids or inputs_embeds, but not both |
| 785 | if input_ids is not None and inputs_embeds is not None: |
| 786 | input_ids = None |
| 787 | |
| 788 | outputs = self.model( |
| 789 | input_ids=input_ids, |
| 790 | inputs_embeds=inputs_embeds, |
| 791 | use_cache=False, |
| 792 | position_ids=positions, |
| 793 | return_dict=False, |
| 794 | **self._output_aux_hidden_states_kwargs, |
| 795 | **kwargs, |
| 796 | ) |
| 797 | |
| 798 | # Remove batch dimension after exiting Transformers model |
| 799 | hidden_states = outputs[0][0, ...] |
| 800 | if self._output_aux_hidden_states_kwargs: |
| 801 | aux_hidden_states = [x[0][0, ...] for x in outputs[1:]] |
| 802 | |
| 803 | if not self.pp_group.is_last_rank: |
| 804 | return IntermediateTensors({"hidden_states": hidden_states}) |
| 805 | |
| 806 | if self._output_aux_hidden_states_kwargs and len(aux_hidden_states) > 0: |
| 807 | return hidden_states, aux_hidden_states |
| 808 | return hidden_states |
| 809 | |
| 810 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 811 | loader = AutoWeightsLoader( |
| 812 | self, |
| 813 | ignore_unexpected_prefixes=self.ignore_unexpected_prefixes, |
| 814 | ignore_unexpected_suffixes=self.ignore_unexpected_suffixes, |
| 815 | ) |
| 816 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 817 | |
| 818 | @staticmethod |
| 819 | def check_version(min_version: str, feature: str): |
| 820 | installed = Version(transformers.__version__) |
| 821 | required = Version(min_version) |
| 822 | if installed < required: |
| 823 | raise ImportError( |
| 824 | f"Transformers modeling backend requires transformers>={required} " |
| 825 | f"for {feature}, but got {installed}" |
| 826 | ) |
| 827 | |
| 828 | def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: |
| 829 | from transformers.utils.output_capturing import ( |
| 830 | OutputRecorder, |
| 831 | maybe_install_capturing_hooks, |
| 832 | ) |
| 833 | |
| 834 | # The default value in PreTrainedModel is None |
| 835 | if self.model._can_record_outputs is None: |
| 836 | self.model._can_record_outputs = {} |
| 837 | |
| 838 | target_class = self._target_class |
| 839 | for layer in layers: |
| 840 | # layer - 1 because we want the input to the layer |
| 841 | layer_name = self._layer_names[layer - 1] |
| 842 | layer_key = f"aux_hidden_state_{layer}" |
| 843 | aux_hidden_state_i = OutputRecorder(target_class, layer_name=layer_name) |
| 844 | self.model._can_record_outputs[layer_key] = aux_hidden_state_i |
| 845 | self._output_aux_hidden_states_kwargs[f"output_{layer_key}"] = True |
| 846 | |
| 847 | # Ensure that the capture hooks are installed before dynamo traces the model |
| 848 | maybe_install_capturing_hooks(self.model) |
| 849 | |
| 850 | def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: |
| 851 | num_layers = self.text_config.num_hidden_layers |
| 852 | return (2, num_layers // 2, num_layers - 3) |