Sync from v0.13

vllm/model_executor/models/transformers/__init__.py

@@ -0,0 +1,127 @@
+# 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.
+"""Wrapper around `transformers` models"""
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.model_executor.models.transformers.base import Base
+from vllm.model_executor.models.transformers.causal import CausalMixin
+from vllm.model_executor.models.transformers.legacy import LegacyMixin
+from vllm.model_executor.models.transformers.moe import MoEMixin
+from vllm.model_executor.models.transformers.multimodal import (
+    DYNAMIC_ARG_DIMS,
+    MultiModalDummyInputsBuilder,
+    MultiModalMixin,
+    MultiModalProcessingInfo,
+    MultiModalProcessor,
+)
+from vllm.model_executor.models.transformers.pooling import (
+    EmbeddingMixin,
+    SequenceClassificationMixin,
+)
+from vllm.model_executor.models.transformers.utils import can_enable_torch_compile
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+
+# Text only models
+@support_torch_compile(enable_if=can_enable_torch_compile)
+class TransformersForCausalLM(CausalMixin, Base): ...
+
+
+@support_torch_compile(enable_if=can_enable_torch_compile)
+class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
+
+
+# Multimodal models
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+@support_torch_compile(
+    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
+)
+class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+@support_torch_compile(
+    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
+)
+class TransformersMultiModalMoEForCausalLM(
+    MoEMixin, MultiModalMixin, CausalMixin, Base
+): ...
+
+
+# Embedding models
+@support_torch_compile(enable_if=can_enable_torch_compile)
+class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ...
+
+
+@support_torch_compile(enable_if=can_enable_torch_compile)
+class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+@support_torch_compile(
+    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
+)
+class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ...
+
+
+# Sequence classification models
+@support_torch_compile(enable_if=can_enable_torch_compile)
+class TransformersForSequenceClassification(
+    SequenceClassificationMixin, LegacyMixin, Base
+): ...
+
+
+@support_torch_compile(enable_if=can_enable_torch_compile)
+class TransformersMoEForSequenceClassification(
+    SequenceClassificationMixin, MoEMixin, Base
+): ...
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MultiModalProcessor,
+    info=MultiModalProcessingInfo,
+    dummy_inputs=MultiModalDummyInputsBuilder,
+)
+@support_torch_compile(
+    dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
+)
+class TransformersMultiModalForSequenceClassification(
+    SequenceClassificationMixin, MultiModalMixin, Base
+): ...
+
+
+def __getattr__(name: str):
+    """Handle imports of non-existent classes with a helpful error message."""
+    if name not in globals():
+        raise AttributeError(
+            "The Transformers modeling backend does not currently have a class to "
+            f"handle the requested model type: {name}. Please open an issue at "
+            "https://github.com/vllm-project/vllm/issues/new"
+        )
+    return globals()[name]

vllm/model_executor/models/transformers/base.py

@@ -0,0 +1,518 @@
+# 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 Iterable
+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.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from vllm.attention.backends.abstract import AttentionType
+from vllm.attention.layer import Attention
+from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
+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.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.utils import (
+    get_feature_request_tip,
+    init_on_device_without_buffers,
+    log_replacement,
+    replace_linear_class,
+    replace_rms_norm_class,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    make_empty_intermediate_tensors_factory,
+    maybe_prefix,
+)
+from vllm.sequence import IntermediateTensors
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+    from vllm.config import VllmConfig
+else:
+    PreTrainedModel = object
+
+logger = init_logger(__name__)
+
+
+def vllm_flash_attention_forward(
+    # Transformers args
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor,
+    # Transformers kwargs
+    scaling: float | None = None,
+    # vLLM kwargs
+    attention_instances: dict[int, Attention] | None = None,
+    **kwargs,
+):
+    self_attn = attention_instances[module.layer_idx]
+    if scaling is not None:
+        self_attn.impl.scale = float(scaling)
+    hidden = query.shape[-2]
+    query, key, value = (x.transpose(1, 2) for x in (query, key, value))
+    query, key, value = (x.reshape(hidden, -1) for x in (query, key, value))
+    return self_attn.forward(query, key, value), None
+
+
+ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward
+
+
+class Base(
+    nn.Module,
+    VllmModel,
+    SupportsQuant,
+    SupportsLoRA,
+    SupportsPP,
+    SupportsEagle,
+    SupportsEagle3,
+):
+    embedding_modules = ["embed_tokens"]  # TODO transformers will have a util to get it
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            # Add `model.` prefix for base model checkpoints,
+            # handling the case where it is already present
+            "": "model.",
+            "model.model.": "model.",
+            # Heads will be adjacent to `model` (pooling included because of adapters)
+            "model.lm_head.": "lm_head.",
+            "model.score.": "classifier.",
+            "model.classifier.": "classifier.",
+        }
+    )
+
+    def __init_subclass__(cls, *args, **kwargs):
+        """Merge hf_to_vllm_mapper in MRO from most specific to least specific."""
+        super().__init_subclass__(*args, **kwargs)
+        hf_to_vllm_mapper = WeightsMapper()
+        for base in cls.__mro__:
+            if base_hf_to_vllm_mapper := getattr(base, "hf_to_vllm_mapper", None):
+                hf_to_vllm_mapper |= base_hf_to_vllm_mapper
+        cls.hf_to_vllm_mapper = hf_to_vllm_mapper
+
+    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.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."""
+
+        # 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 == "mxfp4":
+                raise NotImplementedError(
+                    "Transformers modeling backend does "
+                    "not support MXFP4 quantization yet."
+                )
+            # Skip loading extra bias for GPTQ models.
+            if "gptq" in quant_method_name:
+                self.ignore_unexpected_suffixes.append(".bias")
+
+        # Set correct attn and init on "meta" to delay allocating GPU tensors
+        self.text_config._attn_implementation = "vllm"
+        with init_on_device_without_buffers("meta"):
+            self.model: PreTrainedModel = AutoModel.from_config(
+                self.config,
+                dtype=self.model_config.dtype,
+                trust_remote_code=self.model_config.trust_remote_code,
+            )
+
+        # 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):
+            # Some models scale embeddings inside the input embedding layer
+            self.embed_scale = getattr(input_embeddings, "embed_scale", None)
+            names = ("embedding_size", "hidden_size")
+            embedding_dim = getattr_iter(self.text_config, names, None)
+            assert embedding_dim is not None
+            self.model.set_input_embeddings(
+                VocabParallelEmbedding(
+                    self.text_config.vocab_size,
+                    embedding_dim=embedding_dim,
+                    org_num_embeddings=self.text_config.vocab_size,
+                    quant_config=self.quant_config,
+                )
+            )
+
+        # 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 pipeline_parallel(self):
+        """
+        Apply the model's pipeline parallelization plan.
+        """
+        if self.pp_group.world_size <= 1:
+            return
+
+        if not self.model.supports_pp_plan:
+            tip = get_feature_request_tip(
+                self.model_config.model, self.model_config.trust_remote_code
+            )
+            raise ValueError(
+                f"{type(self.model)} does not support pipeline parallel. {tip}"
+            )
+
+        module_lists = []
+        module_list_idx = None
+        pp_plan = list(self.model._pp_plan.keys())
+        for i, name in enumerate(pp_plan):
+            if isinstance(getattr(self.model, name), 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(self.model)}")
+
+        # Layers before module list
+        for name in pp_plan[:module_list_idx]:
+            if self.pp_group.is_first_rank or (
+                self.text_config.tie_word_embeddings and self.pp_group.is_last_rank
+            ):
+                continue
+            setattr(self.model, name, 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 = pp_plan[module_list_idx]
+        layers = getattr(self.model, layers_name)
+        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 pp_plan[module_list_idx + 1 :]:
+            # Modules that should be on last rank
+            if not self.pp_group.is_last_rank:
+                setattr(self.model, name, PPMissingLayer())
+
+    def recursive_replace(self):
+        """Recursively replace modules in the model as needed.
+
+        Currently, this replaces:
+
+        - `nn.Linear` with vLLM's tensor parallel linear classes
+        - `*RMSNorm` with vLLM's `RMSNorm`
+        """
+        tp_plan = self.model.tp_plan
+
+        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
+            )
+            raise ValueError(
+                f"{type(self.model)} does not support tensor parallel. {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()}
+
+        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)
+                # Populate Eagle3 attrs
+                if (
+                    isinstance(module, nn.ModuleList)
+                    and len(module) == self.text_config.num_hidden_layers
+                ):
+                    self._target_class = type(child_module)
+                    layer_name = qual_name.removeprefix("model.")
+                    self._layer_names[int(child_name)] = layer_name
+                # 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 child_module.__class__.__name__.endswith("RMSNorm"):
+                    new_module = replace_rms_norm_class(
+                        child_module, self.text_config.hidden_size
+                    )
+                else:
+                    _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.dev0", "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_flash_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:
+        inputs_embeds = self.model.get_input_embeddings()(input_ids)
+        if self.embed_scale is not None:
+            inputs_embeds *= self.embed_scale
+        return inputs_embeds
+
+    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"]
+
+        if input_ids is not None:
+            input_ids = input_ids[None, ...]
+        if inputs_embeds is not None:
+            inputs_embeds = inputs_embeds[None, ...]
+
+        # If the model scales embeddings inside the input embedding layer we must
+        # ensure they are scaled here since VocabParallelEmbedding will not do it
+        if (
+            self.embed_scale is not None
+            and input_ids is not None
+            and inputs_embeds is None
+        ):
+            inputs_embeds = self.embed_input_ids(input_ids)
+            input_ids = None
+
+        if self.model_config.uses_mrope:
+            position_ids = positions[:, None]
+        else:
+            position_ids = positions[None, ...]
+
+        outputs = self.model(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            use_cache=False,
+            position_ids=position_ids,
+            attention_instances=self.attention_instances,
+            return_dict=False,
+            **self._output_aux_hidden_states_kwargs,
+            **kwargs,
+        )
+        # We must remove the batch dimension from these outputs
+        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.0.0.dev0", "Eagle3 support")
+        from transformers.utils.generic import OutputRecorder
+
+        # 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
+
+    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
+        num_layers = self.text_config.num_hidden_layers
+        return (2, num_layers // 2, num_layers - 3)

vllm/model_executor/models/transformers/causal.py

@@ -0,0 +1,65 @@
+# 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 mixin for causal language models."""
+
+from typing import TYPE_CHECKING
+
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.models.interfaces_base import VllmModelForTextGeneration
+from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
+
+if TYPE_CHECKING:
+    import torch
+
+    from vllm.config import VllmConfig
+
+
+class CausalMixin(VllmModelForTextGeneration):
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip VllmModelForTextGeneration.__init__ and call the next class in MRO
+        super(VllmModelForTextGeneration, self).__init__(
+            vllm_config=vllm_config, prefix=prefix
+        )
+
+        # Tell `Base.load_weights` to skip
+        # `lm_head` if the model has tied word embeddings
+        if self.text_config.tie_word_embeddings:
+            self.skip_prefixes.append("lm_head.")
+
+        if self.pp_group.is_last_rank:
+            self.lm_head = ParallelLMHead(
+                self.text_config.vocab_size,
+                self.text_config.hidden_size,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+            if self.text_config.tie_word_embeddings:
+                self.lm_head = self.lm_head.tie_weights(
+                    self.model.get_input_embeddings()
+                )
+
+            logit_scale = getattr(self.text_config, "logit_scale", 1.0)
+            self.logits_processor = LogitsProcessor(
+                self.text_config.vocab_size, scale=logit_scale
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+
+    def compute_logits(self, hidden_states: "torch.Tensor") -> "torch.Tensor | None":
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits

vllm/model_executor/models/transformers/legacy.py

@@ -0,0 +1,90 @@
+# 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 mixin for legacy models."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.sequence import IntermediateTensors
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+
+class LegacyMixin:
+    hf_to_vllm_mapper = WeightsMapper(
+        # These are applied in order, so the order matters!
+        orig_to_new_prefix={
+            # Handle BERT-like models
+            "roberta": "model",
+            "bert": "model",
+        },
+        orig_to_new_suffix={
+            # Replace legacy suffixes used for norms
+            ".gamma": ".weight",
+            ".beta": ".bias",
+        },
+    )
+
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+
+        # Skip unsupported/unwanted output embeddings layers
+        self.skip_prefixes.extend(
+            [
+                "model.lm_head.",
+                "model.predictions.",
+                "model.qa_outputs.",
+                "model.embeddings_project.",
+                "model.discriminator_predictions.",
+            ]
+        )
+
+        # Some encoder models have the position_ids buffer in the checkpoint.
+        # vLLM will always pass position_ids as an argument, so we skip loading
+        # the buffer if it exists
+        self.skip_substrs.append("position_ids")
+
+        # Some encoder models have the bias of the final classifier layer
+        # in the checkpoint. vLLM does not use this bias, so we skip loading
+        # it if it exists
+        self.skip_substrs.append("score.bias")
+
+        # roberta-like models an extra padding in positions.
+        # FIXME(Isotr0py): This is quite hacky for roberta edge case,
+        # we should find a better way to handle this.
+        self.is_roberta = "roberta" in self.text_config.model_type
+        self.padding_idx = self.text_config.pad_token_id
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if self.is_roberta:
+            # RoBERTa-specific positions padding
+            positions += self.padding_idx + 1
+        return super().forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )

vllm/model_executor/models/transformers/moe.py

@@ -0,0 +1,325 @@
+# 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 mixin for Mixture of Experts (MoE) models."""
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn as nn
+
+from vllm.config.utils import getattr_iter
+from vllm.distributed import get_dp_group, get_ep_group
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.models.interfaces import MixtureOfExperts
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from .utils import log_replacement
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+
+@CustomOp.register("transformers_fused_moe")
+class TransformersFusedMoE(FusedMoE):
+    """Custom FusedMoE for the Transformers modeling backend."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._topk_ids: torch.Tensor = None
+
+        def custom_routing_function(hidden_states, gating_output, topk, renormalize):
+            """Return `topk_weights` from `gating_output` and the
+            `topk_ids` we stored in the layer earlier."""
+            topk_weights = gating_output
+            topk_ids = self._topk_ids
+            # Handle all gather in expert parallel
+            if topk_ids.size(0) != hidden_states.size(0):
+                dp_metadata = get_forward_context().dp_metadata
+                sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+                is_sp = self.is_sequence_parallel
+                dist_group = get_ep_group() if is_sp else get_dp_group()
+                assert sizes[dist_group.rank_in_group] == topk_ids.shape[0]
+                (topk_ids,) = dist_group.all_gatherv([topk_ids], 0, sizes)
+            return topk_weights, topk_ids
+
+        self.custom_routing_function = custom_routing_function
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """In Transformers `experts.forward` will have this signature.
+
+        We discard any extra kwargs because we cannot use them here."""
+        return torch.ops.vllm.transformers_moe_forward(
+            hidden_states,
+            topk_ids.to(torch.int32),
+            topk_weights.to(torch.float32),
+            self.layer_name,
+        )
+
+
+def transformers_moe_forward(
+    hidden_states: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    """Store the `topk_ids` in the layer and call the actual forward."""
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    self._topk_ids = topk_ids
+    # Clone hidden_states because it will be mutated in-place in FusedMoE
+    return self.forward_impl(hidden_states.clone(), topk_weights)
+
+
+def transformers_moe_forward_fake(
+    hidden_states: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    return torch.empty_like(hidden_states)
+
+
+direct_register_custom_op(
+    op_name="transformers_moe_forward",
+    op_func=transformers_moe_forward,
+    mutates_args=["hidden_states"],
+    fake_impl=transformers_moe_forward_fake,
+    dispatch_key=current_platform.dispatch_key,
+    tags=(torch.Tag.needs_fixed_stride_order,),
+)
+
+
+class MoEMixin(MixtureOfExperts):
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        self.check_version("5.0.0.dev0", "MoE models support")
+        # Skip MixtureOfExperts.__init__ and call the next class in MRO
+        super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix)
+
+    def set_eplb_state(
+        self,
+        expert_load_view: torch.Tensor,
+        logical_to_physical_map: torch.Tensor,
+        logical_replica_count: torch.Tensor,
+    ):
+        for moe_layer_idx, mlp_layer in enumerate(self.mlp_moe_layers):
+            mlp_layer.experts.set_eplb_state(
+                moe_layer_idx=moe_layer_idx,
+                expert_load_view=expert_load_view,
+                logical_to_physical_map=logical_to_physical_map,
+                logical_replica_count=logical_replica_count,
+            )
+
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ):
+        assert self.num_local_physical_experts == num_local_physical_experts
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+        for mlp in self.mlp_moe_layers:
+            mlp.n_local_physical_experts = num_local_physical_experts
+            mlp.n_physical_experts = num_physical_experts
+            mlp.n_redundant_experts = self.num_redundant_experts
+            mlp.experts.update_expert_map()
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        """
+        Params for weights, fp8 weight scales, fp8 activation scales
+        (param_name, weight_name, expert_id, shard_id)
+        """
+        ckpt_names = [
+            # (ckpt_gate_proj_name, ckpt_down_proj_name, ckpt_up_proj_name)
+            ("gate_proj", "down_proj", "up_proj"),  # Most common MoE style
+            ("w1", "w2", "w3"),  # Granite, Mixtral, Phi MoE style
+            ("linear", "linear_1", "linear_v"),  # Grok1 style
+        ]
+        num_experts = self.model_config.get_num_experts()
+        num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
+        expert_mapping = []
+        for gate_proj, down_proj, up_proj in ckpt_names:
+            expert_mapping.extend(
+                FusedMoE.make_expert_params_mapping(
+                    ckpt_gate_proj_name=gate_proj,
+                    ckpt_down_proj_name=down_proj,
+                    ckpt_up_proj_name=up_proj,
+                    num_experts=num_experts,
+                    num_redundant_experts=num_redundant_experts,
+                )
+            )
+        return expert_mapping
+
+    def recursive_replace(self):
+        """Initialize the MoE layers."""
+        text_config = self.text_config
+
+        # Positional arguments
+        num_experts = self.model_config.get_num_experts()
+        top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"], None)
+        assert top_k is not None
+        hidden_size = text_config.hidden_size
+        intermediate_size = getattr_iter(
+            text_config, ["moe_intermediate_size", "intermediate_size"], None
+        )
+        assert intermediate_size is not None
+
+        # If there are shared experts, the results are
+        # reduced after mlp.forward() not inside FusedMoE
+        num_shared_experts = getattr_iter(
+            text_config,
+            [
+                "n_shared_experts",  # DeepSeek, Docs, GLM
+                "moe_num_shared_experts",  # Aria, Ernie
+            ],
+            0,
+        )
+        reduce_results = num_shared_experts == 0
+
+        def add_all_reduce(mlp: nn.Module):
+            """Adds an all-reduce to the output of `mlp.forward()`."""
+
+            class MLPWithAllReduce(mlp.__class__):
+                def forward(self, *args, **kwargs):
+                    output = super().forward(*args, **kwargs)
+                    return self.experts.maybe_all_reduce_tensor_model_parallel(output)
+
+            mlp.__class__ = MLPWithAllReduce
+
+        # Unused kwargs since we use custom_routing_function:
+        # - `scoring_func` and `e_score_correction_bias` only used for grouped
+        #    topk routing inside vLLM and are non-trivial to infer
+        #    and hard code `use_grouped_topk=False`
+        # - `renormalize` passed anyway because it's easy to infer
+        # - `num_expert_group` and `topk_group` used for inferring expert
+        #    placement strategy in FusedMoE
+        # - `apply_router_weight_on_input` is already applied in Transformers
+        renormalize = getattr(text_config, "norm_topk_prob", top_k > 1)
+        num_expert_group = getattr(text_config, "n_group", None)
+        topk_group = getattr(text_config, "topk_group", None)
+
+        # MoE activation function
+        activation = "silu"
+        wrapped_arch = self.config.architectures[0].lower()
+        if "gptoss" in wrapped_arch:
+            activation = "swigluoai"
+        elif "grok1" in wrapped_arch:
+            activation = "gelu"
+
+        # Expert mapping for `AutoWeightsLoader`
+        expert_mapping = self.get_expert_mapping()
+
+        # Expert parallel load balancing kwargs
+        enable_eplb = self.parallel_config.enable_eplb
+        num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
+
+        # MixtureOfExperts mixin settings
+        ep_size = get_ep_group().world_size
+
+        self.mlp_moe_layers = []  # Used for MixtureOfExperts methods
+        self.moe_layers = []
+        self.expert_weights = []
+        self.num_moe_layers = 0
+        self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
+        self.num_logical_experts = num_experts
+        self.num_physical_experts = num_experts + num_redundant_experts
+        self.num_local_physical_experts = self.num_physical_experts // ep_size
+        self.num_routed_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_redundant_experts = num_redundant_experts
+
+        # Recursively fuse MoE layers
+        def _recursive_replace(module: nn.Module, prefix: str):
+            for child_name, child_module in module.named_children():
+                qual_name = maybe_prefix(prefix, child_name)
+                # Naive implementations will have experts as ModuleList
+                is_modulelist = isinstance(child_module, nn.ModuleList)
+                # Packed implementations will have experts as 3D tensors of shapes like:
+                # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size)
+                # down_proj = (num_experts, intermediate_size, hidden_size)
+                params = list(child_module.parameters())
+                is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
+                if child_name == "experts" and (is_modulelist or is_3d):
+                    # Alias for readability
+                    mlp = module
+                    experts = child_module
+                    # Do the experts have biases
+                    has_bias = False
+                    for experts_param_name, _ in experts.named_parameters():
+                        if "bias" in experts_param_name:
+                            has_bias = True
+                            break
+                    # Double check there are no shared experts
+                    nonlocal reduce_results
+                    if reduce_results:
+                        for mlp_param_name, _ in mlp.named_parameters():
+                            if "shared_expert" in mlp_param_name:
+                                reduce_results = False
+                                # If the config does not specify num_shared_experts, but
+                                # the model has shared experts, we assume there is one.
+                                self.num_shared_experts = 1
+                                break
+                    # Replace experts module with FusedMoE
+                    fused_experts = TransformersFusedMoE(
+                        num_experts=num_experts,
+                        top_k=top_k,
+                        hidden_size=hidden_size,
+                        intermediate_size=intermediate_size,
+                        reduce_results=reduce_results,
+                        renormalize=renormalize,
+                        # Hard coded because topk happens in Transformers
+                        use_grouped_topk=False,
+                        num_expert_group=num_expert_group,
+                        topk_group=topk_group,
+                        quant_config=self.quant_config,
+                        prefix=qual_name,
+                        activation=activation,
+                        enable_eplb=enable_eplb,
+                        num_redundant_experts=num_redundant_experts,
+                        has_bias=has_bias,
+                        expert_mapping=expert_mapping,
+                    )
+                    mlp.experts = fused_experts
+                    log_replacement(qual_name, experts, fused_experts)
+                    # Update MixtureOfExperts mixin state
+                    self.mlp_moe_layers.append(mlp)
+                    self.moe_layers.append(fused_experts)
+                    self.expert_weights.append(fused_experts.get_expert_weights())
+                    self.num_moe_layers += 1
+                    # If results are not all-reduced in FusedMoE, ensure they
+                    # are all-reduced at the end of mlp.forward() if tensor
+                    # parallel or expert parallel is enabled
+                    if not reduce_results and (
+                        fused_experts.tp_size > 1 or fused_experts.ep_size > 1
+                    ):
+                        add_all_reduce(mlp)
+                else:
+                    _recursive_replace(child_module, prefix=qual_name)
+
+        _recursive_replace(self.model, prefix="model")
+        # Continue with the replacement of layers in Base
+        super().recursive_replace()

vllm/model_executor/models/transformers/multimodal.py

@@ -0,0 +1,411 @@
+# 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 mixin for multi-modal models."""
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.config.utils import getattr_iter
+from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.multimodal import MultiModalKwargsItems
+from vllm.multimodal.inputs import (
+    MultiModalDataDict,
+    MultiModalFeatureSpec,
+    MultiModalFieldConfig,
+    MultiModalInputs,
+    MultiModalUUIDDict,
+    PlaceholderRange,
+)
+from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
+from vllm.multimodal.processing import BaseMultiModalProcessor, BaseProcessingInfo
+from vllm.multimodal.profiling import BaseDummyInputsBuilder
+from vllm.sequence import IntermediateTensors
+
+if TYPE_CHECKING:
+    from transformers import BatchFeature
+
+    from vllm.config import VllmConfig
+    from vllm.config.multimodal import BaseDummyOptions
+
+DYNAMIC_ARG_DIMS = {
+    "input_ids": 0,
+    # set `positions` to last dim to support Qwen-mrope
+    "positions": -1,
+    "intermediate_tensors": 0,
+    "inputs_embeds": 0,
+}
+
+
+class MultiModalProcessingInfo(BaseProcessingInfo):
+    def get_supported_mm_limits(self):
+        return {"image": None}
+
+    def get_mm_max_tokens_per_item(self, seq_len, mm_counts):
+        return {"image": self.get_max_image_tokens()}
+
+    def get_max_image_tokens(self) -> int:
+        width, height = self.get_max_image_size()
+        processor = self.get_hf_processor()
+        multimodal_config = self.ctx.model_config.multimodal_config
+        mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}
+        mm_tokens = processor._get_num_multimodal_tokens(
+            image_sizes=([height, width],), **mm_processor_kwargs
+        )
+        image_tokens = mm_tokens["num_image_tokens"][0]
+        return image_tokens
+
+    def get_max_image_size(self):
+        return 10_000, 10_000  # hardcode for arbitrary very large size
+
+
+class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_images = mm_counts.get("image", 0)
+
+        processor = self.info.get_hf_processor()
+        if "gemma3" in processor.__class__.__name__.lower():
+            image_token = processor.boi_token
+        else:
+            image_token = getattr(processor, "image_token", "")
+        return image_token * num_images
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, "BaseDummyOptions"] | None = None,
+    ) -> MultiModalDataDict:
+        num_images = mm_counts.get("image", 0)
+
+        target_width, target_height = self.info.get_max_image_size()
+
+        image_overrides = mm_options.get("image") if mm_options else None
+
+        return {
+            "image": self._get_dummy_images(
+                width=target_width,
+                height=target_height,
+                num_images=num_images,
+                overrides=image_overrides,
+            ),
+        }
+
+
+class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ):
+        """
+        Given the original multi-modal items for this modality
+        and HF-processed data, output the updates to perform.
+
+        The information returned by this method is used to update token inputs
+        which bypass the HF processor. It is also used to update the output of
+        HF processor if the HF process does not apply prompt updates to text
+        inputs.
+
+        Moreover, this information is critical to determine the token positions
+        in order to construct  :class:`~vllm-multimodal.input.PlaceholderRange`
+        for each multi-modal item.
+        """
+        return None
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: "BatchFeature",
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        # HF Processors always return a mask but vLLM doesn't need it
+        hf_inputs.pop("attention_mask", None)
+        num_image_patches = hf_inputs.get("num_image_patches")
+        mm_fields = {
+            key: MultiModalFieldConfig.flat_from_sizes("image", num_image_patches)
+            for key in hf_inputs
+        }
+        mm_fields["image_embeds"] = MultiModalFieldConfig.flat_from_sizes(
+            "image", num_image_patches
+        )
+
+        # Keep these as batched, as they always have batch size as first dim
+        mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image")
+        mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image")
+        mm_fields["num_image_patches"] = MultiModalFieldConfig.batched("image")
+        return mm_fields
+
+    def _get_hf_mm_data(
+        self,
+        mm_items: MultiModalDataItems,
+    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+        """
+        In contrast to the base class, this method always adds
+        `return_mm_token_type_ids` to the processor data
+        """
+        processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
+        processor_data["return_mm_token_type_ids"] = True
+        return processor_data, passthrough_data
+
+    def apply(
+        self,
+        prompt: str | list[int],
+        mm_data: MultiModalDataDict,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object] | None = None,
+        mm_uuids: MultiModalUUIDDict | None = None,
+    ) -> MultiModalInputs:
+        """
+        Process multi-modal inputs to be used in vLLM.
+
+        Apply HF Processor on prompt text and multi-modal data together,
+        outputting token IDs and processed tensors.
+        """
+        if tokenization_kwargs is None:
+            tokenization_kwargs = {}
+
+        mm_items = self._to_mm_items(mm_data)
+        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        if not isinstance(prompt, str):
+            # the prompt is the tokenized ids which is not supported
+            # by the hf_processor, which is why we would need to decode the ids
+            # into string
+            prompt = hf_processor.decode(prompt)
+
+        # Bypass cached processor and always apply to the full set of mm inputs
+        # NOTE: we can't just set caching=False because base class method
+        # transforms outputs to `MultiModalKwargs` which is not going to
+        # work for Transformers. We have a lot of logic tied to
+        # `mm_tokens_per_modality` below
+        prompt_ids, processed_data, _ = self._apply_hf_processor_text_mm(
+            prompt_text=prompt,
+            mm_items=mm_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
+        )
+
+        # For gemma3 we check `token_type_ids` as the key
+        token_type_key = (
+            "mm_token_type_ids"
+            if "mm_token_type_ids" in processed_data
+            else "token_type_ids"
+        )
+        mm_token_type_ids = processed_data.pop(token_type_key)
+
+        # We can infer vLLM style placeholder from token type ids, if we split
+        # it for each input `mm_data`.
+        mm_positions = torch.where(mm_token_type_ids == 1)[1]
+        images = mm_items.get_items("image", ImageProcessorItems)
+        multimodal_config = self.info.ctx.model_config.multimodal_config
+        mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}
+        image_sizes = []
+        for item_idx in range(len(images)):
+            image_size = images.get_image_size(item_idx)
+            image_sizes.append((image_size.height, image_size.width))
+
+        mm_tokens_per_modality = hf_processor._get_num_multimodal_tokens(
+            image_sizes=image_sizes, **mm_processor_kwargs
+        )
+
+        mm_placeholders = {}
+        split_sizes = mm_tokens_per_modality["num_image_tokens"]
+        if split_sizes:
+            chunked_mm_positions = torch.split(mm_positions, split_sizes)
+            mm_tokens = torch.tensor(prompt_ids)[mm_token_type_ids[0].bool()]
+            chunked_mm_tokens = torch.split(mm_tokens, split_sizes)
+            ranges = [
+                PlaceholderRange(
+                    offset=positions[0].item(),
+                    length=positions.shape[0],
+                    is_embed=(mm_tokens == hf_processor.image_token_id).bool(),
+                )
+                for positions, mm_tokens in zip(chunked_mm_positions, chunked_mm_tokens)
+            ]
+            mm_placeholders = {"image": ranges}
+
+        processed_data["num_image_patches"] = torch.tensor(
+            mm_tokens_per_modality["num_image_patches"]
+        )
+        mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
+            processed_data,
+            self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs),
+        )
+
+        # Use overrides if provided; fallback to data-dependent hashing.
+        mm_hashes = self._hash_mm_items(
+            mm_items, hf_processor_mm_kwargs, tokenization_kwargs, mm_uuids=mm_uuids
+        )
+
+        return MultiModalInputs(
+            type="multimodal",
+            prompt_token_ids=prompt_ids,
+            mm_kwargs=mm_kwargs,
+            mm_hashes=mm_hashes,
+            mm_placeholders=mm_placeholders,
+        )
+
+
+class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
+    supports_multimodal_raw_input_only = True
+
+    # Backwards compatibility for prev released models. State dicts back then
+    # had different formats and cannot be loaded with `AutoModel` mapping as is
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "language_model.model": "model.language_model",
+            "text_model.model": "model.text_model",
+            "vision_tower": "model.vision_tower",
+            "vqmodel": "model.vqmodel",
+            "visual": "model.visual",
+            "vision_model": "model.vision_model",
+            "vision_embed_tokens": "model.vision_embed_tokens",
+            "image_newline": "model.image_newline",
+            "multi_modal_projector": "model.multi_modal_projector",
+            "text_model.lm_head": "lm_head",
+            "language_model.lm_head": "lm_head",
+            # Qwen models used "model" as the name for the language model.
+            # Therefore, we must map each of submodule explicitly to avoid
+            # conflicts with newer models that use "model.language_model".
+            "model.embed_tokens": "model.language_model.embed_tokens",
+            "model.layers": "model.language_model.layers",
+            "model.norm": "model.language_model.norm",
+        }
+    )
+
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip SupportsMRoPE.__init__ and call the next class in MRO
+        super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        # Gemma3 and PaliGemma needs `token_type_ids` to work correctly
+        # Other models will not have `token_type_ids` in kwargs
+        kwargs = {k: v for k, v in kwargs.items() if k == "token_type_ids"}
+        model_output = super().forward(
+            input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
+        )
+        return model_output
+
+    def get_language_model(self) -> torch.nn.Module:
+        """Transformers modeling backend multimodal classes do not contain a separate
+        vLLM language model class. Therefore, in order to return a language model vLLM
+        class, we use a wrapper to give `self` the same interface as a text model."""
+
+        # Exclude self and object
+        bases = self.__class__.mro()[1:-1]
+        # Keep only classes defined in `vllm.model_executor.models.transformers`
+        bases = [b for b in bases if ".transformers." in b.__module__]
+        # Exclude MultiModalMixin itself
+        bases = [b for b in bases if b is not MultiModalMixin]
+
+        class LanguageModel(*bases):
+            def __init__(self, multimodal_model):
+                # Don't call super().__init__() to avoid re-initialization
+                self.__dict__.update(multimodal_model.__dict__)
+
+            model = getattr_iter(self.model, ("language_model", "text_model"), None)
+
+        return LanguageModel(self)
+
+    def embed_multimodal(self, **kwargs):
+        pixel_values: torch.Tensor | None = kwargs.pop("pixel_values", None)
+        image_embeds: torch.Tensor | None = kwargs.pop("image_embeds", None)
+        # Model might use `image_patches` instead of `pixel_values`
+        if pixel_values is None:
+            pixel_values = kwargs.pop("image_patches", None)
+
+        if image_embeds is not None:
+            return image_embeds
+
+        if pixel_values is None:
+            return None
+
+        num_image_patches = kwargs.pop("num_image_patches")
+        kwargs.pop("token_type_ids", None)  # used only in `forward`
+        if pixel_values is not None:
+            vision_embeddings = self.model.get_image_features(pixel_values, **kwargs)
+
+            if isinstance(vision_embeddings, torch.Tensor):
+                if vision_embeddings.ndim == 2:
+                    vision_embeddings = vision_embeddings.unsqueeze(0)
+
+                # Embeddings have to be 2D tensors of length `num_images`
+                # but transformers returns concat tensors if each patch
+                # is of different size. We split it back to make vLLM happy
+                vision_embeddings = torch.split(
+                    vision_embeddings, num_image_patches.flatten().tolist()
+                )
+                vision_embeddings = [
+                    embed.flatten(start_dim=0, end_dim=-2)
+                    for embed in vision_embeddings
+                ]
+
+            return vision_embeddings
+
+    def get_mrope_input_positions(
+        self,
+        input_tokens: list[int],
+        mm_features: list[MultiModalFeatureSpec],
+    ) -> tuple[torch.Tensor, int]:
+        kwargs = MultiModalFeatureSpec.gather_kwargs(
+            mm_features,
+            {
+                "image_grid_thw",
+                "video_grid_thw",
+                "second_per_grid_ts",
+                "audio_feature_lengths",
+                "use_audio_in_video",
+            },
+        )
+        if any(
+            v
+            for k, v in kwargs.items()
+            if k not in {"image_grid_thw", "video_grid_thw"}
+        ):
+            raise NotImplementedError(
+                "Transformers modeling backend only supports images."
+            )
+
+        image_grid_thw = kwargs.get("image_grid_thw", [])
+        video_grid_thw = kwargs.get("video_grid_thw", [])
+
+        image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(
+            image_grid_thw
+        )
+        video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(
+            video_grid_thw
+        )
+
+        mrope_positions, mrope_position_delta = self.model.get_rope_index(
+            input_ids=torch.tensor(input_tokens).unsqueeze(0),
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+        )
+
+        mrope_positions = mrope_positions[:, 0]
+        mrope_position_delta = mrope_position_delta[0].item()
+
+        return mrope_positions, mrope_position_delta

vllm/model_executor/models/transformers/pooling.py

@@ -0,0 +1,119 @@
+# 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 mixins for pooling models."""
+
+from typing import TYPE_CHECKING
+
+import torch
+from transformers import AutoModelForSequenceClassification
+
+from vllm.config.utils import getattr_iter
+from vllm.model_executor.layers.pooler import (
+    ClassifierPooler,
+    CLSPool,
+    DispatchPooler,
+    Pooler,
+)
+from vllm.model_executor.models.interfaces import SupportsCrossEncoding
+from vllm.model_executor.models.interfaces_base import VllmModelForPooling
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+
+class EmbeddingMixin(VllmModelForPooling):
+    default_pooling_type = "CLS"
+
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip VllmModelForPooling.__init__ and call the next class in MRO
+        super(VllmModelForPooling, self).__init__(
+            vllm_config=vllm_config, prefix=prefix
+        )
+
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+
+        self.pooler = DispatchPooler(
+            {
+                "token_embed": Pooler.for_token_embed(pooler_config),
+                "embed": Pooler.for_embed(pooler_config),
+            }
+        )
+
+
+class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):
+    default_pooling_type = "CLS"
+
+    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
+        # Skip VllmModelForPooling.__init__ and call the next class in MRO
+        super(VllmModelForPooling, self).__init__(
+            vllm_config=vllm_config, prefix=prefix
+        )
+
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+
+        # Certain information about the the model and classifier can only be
+        # inferred from the `ForSequenceClassification` class. Therefore, we
+        # instantiate it on the "meta" device to avoid allocating GPU memory.
+        with torch.device("meta"):
+            seq_cls_model = AutoModelForSequenceClassification.from_config(
+                self.config,
+                dtype=self.model_config.dtype,
+                trust_remote_code=self.model_config.trust_remote_code,
+            )
+
+        # When used for sequence classification, some models have their
+        # pooling layers removed. Make sure this is reflected in vLLM.
+        for module in seq_cls_model.modules():
+            if hasattr(module, "pooler") and module.pooler is None:
+                self.model.pooler = None
+                break
+
+        # Unlike `lm_head`, `classifier` is not always `nn.Linear`.
+        self.classifier = getattr_iter(seq_cls_model, ["classifier", "score"], None)
+        if self.classifier is None:
+            raise ValueError(
+                "Could not find `classifier` or `score` layer in the "
+                "`AutoModelForSequenceClassification` instance."
+            )
+        self.init_parameters(self.classifier, dtype=self.model_config.head_dtype)
+
+        class ClassifierWithReshape(self.classifier.__class__):
+            """CLSPool has already been applied in `pooling`.
+            Add dim to match expected input shape of `classifier.forward`."""
+
+            def forward(self, *args, **kwargs):
+                if len(args) > 0:
+                    args = (args[0].unsqueeze(1), *args[1:])
+                return super().forward(*args, **kwargs)
+
+        self.classifier.__class__ = ClassifierWithReshape
+
+        self.pooler = DispatchPooler(
+            {
+                "token_classify": Pooler.for_token_classify(
+                    pooler_config, classifier=self.classifier
+                ),
+                "classify": ClassifierPooler(
+                    pooling=CLSPool(), classifier=self.classifier, act_fn="classify"
+                ),
+                "score": ClassifierPooler(
+                    pooling=CLSPool(), classifier=self.classifier, act_fn="score"
+                ),
+            }
+        )

vllm/model_executor/models/transformers/utils.py

@@ -0,0 +1,213 @@
+# 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 contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import torch
+from torch import nn
+from transformers.configuration_utils import ALLOWED_LAYER_TYPES
+
+from vllm.config.utils import getattr_iter
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+
+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 = {}
+
+    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", "colwise_rep", "rowwise", "rowwise_rep", "replicate"]
+
+
+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.
+    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, {}),
+        "colwise_rep": (ColumnParallelLinear, {"gather_output": True}),
+        "rowwise": (RowParallelLinear, {}),
+        "rowwise_rep": (RowParallelLinear, {"input_is_parallel": False}),
+        "replicate": (ReplicatedLinear, {}),
+    }.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,
+    )
+
+
+def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm:
+    """Replace a Transformers RMSNorm with vLLM's RMSNorm.
+
+    This method assumes:
+    - Weight is stored as `weight`.
+    - Epsilon is stored as `eps` or `variance_epsilon`.
+    - `with_scale` indicates whether the layer has a weight (Gemma3n only).
+    - `var_hidden_size` is only ever used for Intern vision encoder in vLLM
+    and Transformers doesn't appear to have the same concept.
+    """
+    eps = getattr_iter(rms_norm, ("eps", "variance_epsilon"), 1e-6)
+    kwargs = {"hidden_size": hidden_size, "eps": eps}
+    # Update hidden size if weight is available
+    weight_meta = getattr(rms_norm, "weight", None)
+    if weight_meta is not None:
+        kwargs["hidden_size"] = weight_meta.size(0)
+    # Check if weight is all zeros, which indicates GemmaRMSNorm
+    # We must create a new instance because rms_norm is on meta
+    try:
+        with torch.device("cpu"):
+            weight_test = getattr(rms_norm.__class__(1), "weight", None)
+    except Exception:
+        logger.warning(
+            "Failed to determine if RMSNorm weight is centered on zero or one. "
+            "Defaulting to one."
+        )
+        weight_test = None
+    if weight_test is not None and torch.all(weight_test == 0):
+        return GemmaRMSNorm(**kwargs)
+    # Otherwise assume it's a regular RMSNorm
+    kwargs["has_weight"] = getattr(rms_norm, "with_scale", True)
+    if weight_meta is not None:
+        kwargs["dtype"] = weight_meta.dtype
+    else:
+        # No weight, fall back to weightless RMSNorm
+        kwargs["has_weight"] = False
+    return RMSNorm(**kwargs)
+
+
+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 set(rope_parameters.keys()).issubset(ALLOWED_LAYER_TYPES):
+            rope_parameters = {"": rope_parameters}
+        return all(rp["rope_type"] != "dynamic" for rp in rope_parameters.values())
+    return True