# 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 base class.""" from collections.abc import Callable, Iterable from itertools import chain from operator import attrgetter from typing import TYPE_CHECKING import regex as re import torch import transformers from packaging.version import Version from torch import nn from transformers import AutoModel from transformers.conversion_mapping import ( WeightRenaming, get_model_conversion_mapping, ) from vllm.compilation.decorators import support_torch_compile from vllm.config.utils import getattr_iter from vllm.distributed import get_pp_group, get_tp_group from vllm.distributed.utils import get_pp_indices from vllm.logger import init_logger from vllm.model_executor.layers.attention import ( Attention, EncoderOnlyAttention, ) from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.models.interfaces import ( SupportsEagle, SupportsEagle3, SupportsLoRA, SupportsPP, SupportsQuant, ) from vllm.model_executor.models.interfaces_base import VllmModel from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers from vllm.model_executor.models.transformers.utils import ( can_enable_torch_compile, get_feature_request_tip, init_on_device_without_buffers, log_replacement, replace_conv_class, replace_linear_class, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix, ) from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType if TYPE_CHECKING: from transformers import PreTrainedModel from vllm.config import VllmConfig logger = init_logger(__name__) class ScaledVocabParallelEmbedding(VocabParallelEmbedding): """`VocabParallelEmbedding` that scales its output.""" def __init__(self, *args, embed_scale: float, **kwargs): super().__init__(*args, **kwargs) self.embed_scale = embed_scale def forward(self, input_: torch.Tensor) -> torch.Tensor: return super().forward(input_) * self.embed_scale class Base( nn.Module, VllmModel, SupportsQuant, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3, ): embedding_modules = ["embed_tokens"] # TODO transformers will have a util to get it def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): super().__init__() logger.info("Using Transformers modeling backend.") self.config = vllm_config.model_config.hf_config self.text_config = self.config.get_text_config() self.cache_config = vllm_config.cache_config self.compilation_config = vllm_config.compilation_config self.device_config = vllm_config.device_config self.model_config = vllm_config.model_config self.parallel_config = vllm_config.parallel_config self.quant_config = vllm_config.quant_config self.pp_group = get_pp_group() self.tp_group = get_tp_group() # Attrs for weight loading (see self.load_weights) self.skip_prefixes: list[str] = [] """Skip loading weights whose qualname starts with these prefixes.""" self.skip_substrs: list[str] = [] """Skip loading weights whose qualname contains these substrings.""" self.ignore_unexpected_prefixes: list[str] = [] """Ignore unexpected weights whose qualname starts with these prefixes.""" self.ignore_unexpected_suffixes: list[str] = [] """Ignore unexpected weights whose qualname ends with these suffixes.""" self.packed_modules_mapping: dict[str, list[str]] = {} """Fused module -> constituent projections, populated by `recursive_replace` for the quantization machinery and loaders (e.g. bitsandbytes).""" # Attrs for Eagle3 (see self.set_aux_hidden_state_layers) self._target_class: type[nn.Module] = nn.Module """Target class for Eagle3 aux hidden state recording.""" self._layer_names: dict[int, str] = {} """Mapping from layer index to layer name for Eagle3.""" self._output_aux_hidden_states_kwargs: dict[str, bool] = {} """Kwargs to pass to model forward for Eagle3 aux hidden states.""" if self.quant_config: quant_method_name = self.quant_config.get_name() # Check for unsupported quantization methods. if quant_method_name in ("mxfp4", "gpt_oss_mxfp4"): raise NotImplementedError( "Transformers modeling backend does " "not support MXFP4 quantization yet." ) self._patch_config() from_config_kwargs = dict( config=self.config, dtype=self.model_config.dtype, trust_remote_code=self.model_config.trust_remote_code, ) self._decorate_for_torch_compile(**from_config_kwargs) # Init on "meta" to delay allocating GPU tensors with init_on_device_without_buffers("meta"): self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs) # Create weight name to module qualname mapper self._create_hf_to_vllm_mapper() # Remove layers not on this pipeline parallel rank self.pipeline_parallel() # Substitute remaining layers with vLLM's layers as needed self.recursive_replace() # Create attention instances for KV cache allocation self.attention_instances = self.create_attention_instances() # Input embeddings input_embeddings = self.model.get_input_embeddings() if not isinstance(input_embeddings, PPMissingLayer): names = ("embedding_size", "hidden_size") embedding_dim = getattr_iter(self.text_config, names, None) assert embedding_dim is not None embedding_kwargs = dict( num_embeddings=self.text_config.vocab_size, embedding_dim=embedding_dim, org_num_embeddings=self.text_config.vocab_size, quant_config=self.quant_config, ) embed_scale = getattr(input_embeddings, "embed_scale", None) if embed_scale is not None: # Some models scale embeddings inside the input embedding layer new_input_embeddings = ScaledVocabParallelEmbedding( **embedding_kwargs, embed_scale=float(embed_scale) ) else: new_input_embeddings = VocabParallelEmbedding(**embedding_kwargs) self.model.set_input_embeddings(new_input_embeddings) # Initialize any parameters that have not had their modules replaced self.init_parameters(self.model) # Pipeline parallel intermediate tensors self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states"], self.text_config.hidden_size ) def _patch_config(self): """ Patch the config to ensure that the model is created correctly: - Sets the attention implementation to "vllm" so the attention instances from `create_attention_instances` are used - Sets the dtype to the default torch dtype set by vLLM because Transformers uses the config dtype when creating the model """ self.text_config._attn_implementation = "vllm" self.config.dtype = torch.get_default_dtype() def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]: """ Get the decoder class from the model. Args: kwargs: The kwargs to create the model. Returns: The decoder class. """ with torch.device("meta"): model: PreTrainedModel = AutoModel.from_config(**kwargs) decoder_cls = type(model.get_decoder()) logger.debug("Identified decoder class as: %s", decoder_cls) del model return decoder_cls def _decorate_cls_for_torch_compile( self, cls: type["PreTrainedModel"], dynamic_arg_dims: dict[str, int] | None, enable_if: Callable[["VllmConfig"], bool], is_encoder: bool, ): """ Decorate `cls` to indicate to vLLM that it supports torch compile. Args: cls: The PreTrainedModel class to decorate. dynamic_arg_dims: A mapping from argument name to the dynamic dimensions of the argument. If None, default dynamic arg dims will be used. See [`support_torch_compile`][vllm.compilation.decorators.support_torch_compile] for more details. enable_if: A function which takes in the vLLM config and returns whether torch compile should be enabled for this class. is_encoder: Whether the class being decorated is an encoder. """ logger.debug( "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s", cls.__name__, "encoder" if is_encoder else "decoder", dynamic_arg_dims, ) support_torch_compile( dynamic_arg_dims=dynamic_arg_dims, enable_if=enable_if, is_encoder=is_encoder, )(cls) def _decorate_for_torch_compile(self, **kwargs: dict): """ Decorate the model's decoder class to indicate to vLLM that it supports torch compile if `can_enable_torch_compile` is True. Args: kwargs: The kwargs to create the model, which are needed to get the decoder class. """ self._decorate_cls_for_torch_compile( cls=self._get_decoder_cls(**kwargs), # Applied to a PreTrainedModel so the batch dimension will exist dynamic_arg_dims=dict[str, int]( input_ids=1, # shape: [1, seq_len] inputs_embeds=1, # shape: [1, seq_len, hidden_size] position_ids=-1, # shape: [1, seq_len] or [3, 1, seq_len] for mrope ), enable_if=can_enable_torch_compile, is_encoder=False, ) def _create_hf_to_vllm_mapper(self): """ Create a WeightsMapper to map checkpoint weight names to module qualnames. This handles: - Transformers weight renaming from `WeightRenaming` - Checkpoints saved with a base model prefix that is not `model` - Checkpoints saved with no base model prefix - Any quantization config specific mappings """ self.hf_to_vllm_mapper = WeightsMapper() orig_to_new_renaming = self.hf_to_vllm_mapper.orig_to_new_renaming orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex for mapping in get_model_conversion_mapping(self.model): # Handle weights which have been renamed in Transformers if isinstance(mapping, WeightRenaming): orig_to_new_renaming.append(mapping) # TODO: Handle WeightConverter to enable layer merging # Handle unexpected weights which should be ignored if self.model._keys_to_ignore_on_load_unexpected is not None: for key in self.model._keys_to_ignore_on_load_unexpected: orig_to_new_regex[re.compile(key)] = None # Standardise base model prefix bmp = self.model.base_model_prefix expected_bmp = r"model.\1" # Handle checkpoints saved with different base model prefix if bmp and bmp != "model": different_bmp_pattern = re.compile(rf"^{bmp}\.(.+)") orig_to_new_regex[different_bmp_pattern] = expected_bmp # Handle direct children of self.model which were saved without the model prefix direct_children = chain( self.model.named_children(), self.model.named_parameters(recurse=False), self.model.named_buffers(recurse=False), ) model_children = "|".join(name for name, _ in direct_children) missing_bmp_pattern = re.compile(rf"^(?!model\.)(({model_children}).*)") orig_to_new_regex[missing_bmp_pattern] = expected_bmp # Handle weights saved as direct children of self.model which no longer are unexpected_bmp_pattern = re.compile(rf"^(model\.)((?!{model_children}).+)") orig_to_new_regex[unexpected_bmp_pattern] = r"\2" # Handle lm_head which was saved inside the base model nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)") orig_to_new_regex[nested_lm_head_pattern] = r"\2" # Apply mapping to quantization config if needed self._maybe_apply_model_mapping() def _get_tie_word_embeddings(self): """ Check if the model has tied word embeddings. """ # Models created with Transformers v4 and v5 will store this in different places tie_word_embeddings_v4 = getattr(self.text_config, "tie_word_embeddings", False) tie_word_embeddings_v5 = getattr(self.config, "tie_word_embeddings", False) return tie_word_embeddings_v4 or tie_word_embeddings_v5 def pipeline_parallel(self): """ Apply the model's pipeline parallelization plan. """ if self.pp_group.world_size <= 1: return if self.model.supports_pp_plan: module = self.model names = list(module._pp_plan.keys()) else: module = self.model.get_decoder() has_parameters = lambda m: next(m.parameters(), None) is not None names = [n for n, c in module.named_children() if has_parameters(c)] tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) logger.warning( "%s does not define a pipeline parallel plan. The Transformers " "modeling backend will infer the split from the layers of %s in order " "of declaration and keep parameter-free modules on every rank. This " "may fail if the model's structure is non-standard. %s", type(self.model), type(module), tip, ) 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 module_lists = [] module_list_idx = None for i, name in enumerate(names): # attrgetter in case the module is nested (e.g. "text_model.layers") if isinstance(attrgetter(name)(module), nn.ModuleList): module_lists.append(name) module_list_idx = i if len(module_lists) > 1: raise ValueError( "Pipeline parallel of models with multiple `ModuleList`s " "in the base model are not supported yet!" ) if module_list_idx is None: raise ValueError(f"Could not find `ModuleList` in {type(module)}") # Layers before module list for name in names[:module_list_idx]: if self.pp_group.is_first_rank or ( self._get_tie_word_embeddings() and self.pp_group.is_last_rank ): continue # attrsetter in case the module is nested (e.g. "text_model.embed_tokens") attrsetter(name)(module, PPMissingLayer()) # Module list start_layer, end_layer = get_pp_indices( self.text_config.num_hidden_layers, self.pp_group.rank_in_group, self.pp_group.world_size, ) layers_name = names[module_list_idx] # attrgetter in case the module is nested (e.g. "text_model.layers") layers = attrgetter(layers_name)(module) for i in range(len(layers)): if start_layer <= i and i < end_layer: continue layers[i] = PPMissingLayer() # Layers after module list for name in names[module_list_idx + 1 :]: # Modules that should be on last rank if not self.pp_group.is_last_rank: # attrsetter in case the module is nested (e.g. "text_model.norm") attrsetter(name)(module, PPMissingLayer()) def recursive_replace(self): """Recursively replace modules in the model as needed. Currently, this replaces: - GLUs with a fused `MergedColumnParallelLinear` + `...AndMul` - Attention QKV projections with a fused `QKVParallelLinear` + split - `nn.Linear` with vLLM's tensor parallel linear classes - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d` - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm` """ tp_plan = self.model.tp_plan or {} if not tp_plan and self.tp_group.world_size > 1: tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) logger.warning_once( "%s does not define a tensor parallel plan. The Transformers modeling " "backend will shard the model the best it can during graph fusion and " "replicate the rest. This may be suboptimal or fail if the model does " "not fuse cleanly. %s", type(self.model), tip, ) # Prefix the patterns because we always start from `self.model` tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()} # Detect fusable patterns once per module class (cached, so this is cheap) fusers = Fusers(self.model, self.model_config) def register_fusion(fuser: BaseFuser, prefix: str): """Register a fused layer's mappings just before it is built.""" orig_to_new_stacked = fuser.orig_to_new_stacked(prefix) self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked) packed_modules_mapping = fuser.packed_modules_mapping self.packed_modules_mapping.update(packed_modules_mapping) if self.quant_config is not None: self.quant_config.packed_modules_mapping.update(packed_modules_mapping) 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) if ( isinstance(module, nn.ModuleList) and len(module) == self.text_config.num_hidden_layers ): # Populate Eagle3 attrs self._target_class = type(child_module) layer_name = qual_name.removeprefix("model.") self._layer_names[int(child_name)] = layer_name # MTP weights should not be loaded into the base model num_hidden_layers = self.text_config.num_hidden_layers names = ( "n_predict", # Override from SpeculativeConfig "num_nextn_predict_layers", # Most models "mtp_num_hidden_layers", # Qwen 3.5 ) n_predict = getattr_iter(self.text_config, names, 0) for i in range(num_hidden_layers, num_hidden_layers + n_predict): mtp_prefix = f"{prefix}.{i}." if mtp_prefix not in self.ignore_unexpected_prefixes: self.ignore_unexpected_prefixes.append(mtp_prefix) # Replace modules as needed if isinstance(child_module, nn.Linear): generator = (p for p in tp_plan if re.match(p, qual_name)) pattern = next(generator, None) # Some weight loaders expect all linear layers to inherit # LinearBase, so we set a default style which causes any # unspecified layers to be replaced with ReplicatedLinear style = tp_plan.get(pattern, "replicate") new_module = replace_linear_class( child_module, style, self.quant_config, prefix=qual_name ) elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)): new_module = replace_conv_class(child_module) elif (fuser := fusers[child_module]) is not None: register_fusion(fuser, qual_name) new_module = fuser.fuse( child_module, qual_name, self.model_config, self.quant_config ) logger.info_once(fuser.info(child_name)) _recursive_replace(new_module, prefix=qual_name) elif not isinstance(child_module, MoERunner): # MoERunner can contain aliases of shared experts and gates, # so we don't want to recurse into it and break weight loading. _recursive_replace(child_module, prefix=qual_name) if new_module is not child_module: setattr(module, child_name, new_module) log_replacement(qual_name, child_module, new_module) _recursive_replace(self.model, prefix="model") def create_attention_instances(self) -> dict[int, Attention]: """ Create `Attention` instances to inform KV cache allocation. """ text_config = self.text_config num_heads = self.model_config.get_num_attention_heads(self.parallel_config) head_size = self.model_config.get_head_size() num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config) logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None) # In encoder models, the attention layers will have `is_causal=False` is_encoder = lambda module: not getattr(module, "is_causal", True) has_encoder = lambda model: any(is_encoder(m) for m in model.modules()) is_multimodal = lambda config: config != config.get_text_config() # vLLM does not support encoder-decoder models, so if any encoder layer is # found in a text only model, we assume the whole model is an encoder model if has_encoder(self.model) and not is_multimodal(self.config): self.check_version("5.0.0", "encoder models support") attn_type = AttentionType.ENCODER_ONLY else: attn_type = AttentionType.DECODER pp_rank = self.pp_group.rank_in_group pp_size = self.pp_group.world_size start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size) attention_instances = {} for i in range(start, end): # Handle interleaved sliding window attention per_layer_sliding_window = None if ( hasattr(self.config, "layer_types") and self.config.layer_types[i] == "sliding_attention" ): per_layer_sliding_window = self.config.sliding_window attn_cls = ( EncoderOnlyAttention if attn_type == AttentionType.ENCODER_ONLY else Attention ) attention_instances[i] = attn_cls( num_heads=num_heads, head_size=head_size, # NOTE: We use Llama scale as default, if it's set by # Transformers, it's updated in vllm_attention_forward scale=head_size**-0.5, num_kv_heads=num_kv_heads, cache_config=self.cache_config, quant_config=self.quant_config, logits_soft_cap=logits_soft_cap, per_layer_sliding_window=per_layer_sliding_window, prefix=f"{i}.attn", attn_type=attn_type, ) return attention_instances def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None): """ If a `parameter` is on the `meta` device, then its parent `module` is the original module created by: ```python with torch.device("meta"): self.model: "PreTrainedModel" = AutoModel.from_config(...) ``` """ def _init_parameters(module: nn.Module, dtype: torch.dtype | None): for name, param in module.named_parameters(recurse=False): if param.device == torch.device("meta"): new_param = nn.Parameter( torch.empty_like( param.data, dtype=dtype or self.model_config.dtype, device=self.device_config.device, ) ) setattr(module, name, new_param) for child in module.children(): _init_parameters(child, dtype) _init_parameters(module, dtype) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.get_input_embeddings()(input_ids) def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | IntermediateTensors: if not self.pp_group.is_first_rank: assert intermediate_tensors is not None input_ids = None inputs_embeds = intermediate_tensors["hidden_states"] # Add batch dimension before entering Transformers model if input_ids is not None and input_ids.ndim == 1: # [seq_len] -> [1, seq_len] input_ids = input_ids[None, ...] if inputs_embeds is not None and inputs_embeds.ndim == 2: # [seq_len, hidden_size] -> [1, seq_len, hidden_size] inputs_embeds = inputs_embeds[None, ...] if positions.ndim == 1: # [seq_len] -> [1, seq_len] positions = positions[None, ...] outputs = self.model( input_ids=input_ids, inputs_embeds=inputs_embeds, use_cache=False, position_ids=positions, attention_instances=self.attention_instances, return_dict=False, **self._output_aux_hidden_states_kwargs, **kwargs, ) # Remove batch dimension after exiting Transformers model hidden_states = outputs[0][0, ...] if self._output_aux_hidden_states_kwargs: aux_hidden_states = [x[0][0, ...] for x in outputs[1:]] if not self.pp_group.is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) if self._output_aux_hidden_states_kwargs and len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, skip_prefixes=self.skip_prefixes, skip_substrs=self.skip_substrs, ignore_unexpected_prefixes=self.ignore_unexpected_prefixes, ignore_unexpected_suffixes=self.ignore_unexpected_suffixes, ) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) @staticmethod def check_version(min_version: str, feature: str): installed = Version(transformers.__version__) required = Version(min_version) if installed < required: raise ImportError( f"Transformers modeling backend requires transformers>={required} " f"for {feature}, but got {installed}" ) def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: self.check_version("5.2.0", "Eagle3 support") from transformers.utils.output_capturing import ( OutputRecorder, maybe_install_capturing_hooks, ) # The default value in PreTrainedModel is None if self.model._can_record_outputs is None: self.model._can_record_outputs = {} target_class = self._target_class for layer in layers: # layer - 1 because we want the input to the layer layer_name = self._layer_names[layer - 1] layer_key = f"aux_hidden_state_{layer}" aux_hidden_state_i = OutputRecorder(target_class, layer_name=layer_name) self.model._can_record_outputs[layer_key] = aux_hidden_state_i self._output_aux_hidden_states_kwargs[f"output_{layer_key}"] = True # Ensure that the capture hooks are installed before dynamo traces the model maybe_install_capturing_hooks(self.model) def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: num_layers = self.text_config.num_hidden_layers return (2, num_layers // 2, num_layers - 3)