iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[2/N] Harden Transformers modelling backend multi-modal path (#51657)

Browse Source 
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Harry Mellor
2026-08-10 12:57:12 +00:00
committed by GitHub
parent 436be94e13
commit cf8f3a3bb2
3 changed files with 273 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/transformers/test_backend.py
+54
View File
@@ -10,6 +10,13 @@ from typing import Any
import pytest
import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel
from vllm.config import ModelConfig, VllmConfig
from vllm.model_executor.models.interfaces import SupportsMultiModal
from vllm.model_executor.models.transformers.base import Base
from vllm.model_executor.models.transformers.multimodal import MultiModalMixin
from vllm.model_executor.models.utils import StageMissingLayer
from ...conftest import HfRunner, VllmRunner
from ...utils import multi_gpu_test, prep_prompts
@@ -468,3 +475,50 @@ def test_replaced_embedding_exposes_one_vpe(vpe):
lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
assert lm_head.tie_weights(composed.embed).weight is composed.embed.weight
MULTIMODAL_MODEL = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
class MarkingStub(SupportsMultiModal, nn.Module):
"""Just enough of the backend to exercise component marking."""
_mark_model_components = MultiModalMixin._mark_model_components
_find_encoder_classes = MultiModalMixin._find_encoder_classes
_from_config_kwargs = Base._from_config_kwargs
_pre_trained_model_classes = Base._pre_trained_model_classes
def build_marked_model(image_limit: int, skip_tokenizer_init: bool = False):
"""Build the HF model inside the marking context and return it."""
model_config = ModelConfig(
model=MULTIMODAL_MODEL,
model_impl="transformers",
limit_mm_per_prompt={"image": image_limit},
)
# Set after construction: building the config itself needs the tokenizer
model_config.skip_tokenizer_init = skip_tokenizer_init
stub = MarkingStub()
stub.config = AutoConfig.from_pretrained(MULTIMODAL_MODEL)
stub.model_config = model_config
vllm_config = VllmConfig(model_config=model_config)
with stub._mark_model_components(vllm_config), torch.device("meta"):
stub.model = AutoModel.from_config(**stub._from_config_kwargs)
return stub.model
@pytest.mark.parametrize(("image_limit", "skipped"), [(0, True), (4, False)])
def test_tower_weights_skipped_when_modality_disabled(image_limit, skipped):
"""`--limit-mm-per-prompt image=0` should drop the vision tower's weights."""
vision_tower = build_marked_model(image_limit).vision_tower
assert isinstance(vision_tower, StageMissingLayer) is skipped
def test_marking_skipped_without_tokenizer():
"""Marking needs the HF processor, which needs a tokenizer, so it is skipped.
The tower is built as normal rather than the model failing to load.
"""
vision_tower = build_marked_model(0, skip_tokenizer_init=True).vision_tower
assert not isinstance(vision_tower, StageMissingLayer)
vllm/model_executor/models/transformers/base.py
+52 -20
View File
@@ -18,9 +18,11 @@
import os
from collections.abc import Callable, Iterable
from contextlib import contextmanager
from functools import cached_property
from itertools import chain
from operator import attrgetter
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any, NamedTuple
import regex as re
import torch
@@ -85,6 +87,12 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
class PreTrainedModelClasses(NamedTuple):
decoder: type["PreTrainedModel"]
encoders: dict[str, type["PreTrainedModel"]]
"""Modality -> encoder class, for each modality that has one."""
class Base(
nn.Module,
VllmModel,
@@ -147,14 +155,13 @@ class Base(
)
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)
self._decorate_for_torch_compile()
# Init on "meta" to delay allocating GPU tensors
with init_on_device_without_buffers("meta"):
with (
self._mark_model_components(vllm_config),
init_on_device_without_buffers("meta"),
):
from_config_kwargs = self._from_config_kwargs
self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs)
# Create weight name to module qualname mapper
@@ -196,22 +203,47 @@ class Base(
self.text_config._attn_implementation = "vllm"
self.config.dtype = torch.get_default_dtype()
def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]:
@contextmanager
def _mark_model_components(self, vllm_config: "VllmConfig"):
"""Mark language model and tower submodules as `self.model` is created.
Nothing to do in `Base`, `MultiModalMixin` will override."""
yield
@cached_property
def _from_config_kwargs(self) -> dict[str, Any]:
"""The kwargs used to create `self.model`."""
return dict(
config=self.config,
dtype=self.model_config.dtype,
trust_remote_code=self.model_config.trust_remote_code,
)
def _find_encoder_classes(
self, model: "PreTrainedModel"
) -> dict[str, type["PreTrainedModel"]]:
"""Find the encoder class of each modality `model` has one for.
Text models have none. Multi-modal ones override this.
"""
Get the decoder class from the model.
return {}
Args:
kwargs: The kwargs to create the model.
@cached_property
def _pre_trained_model_classes(self) -> PreTrainedModelClasses:
"""The decoder class and the encoder class of each modality that has one.
Returns:
The decoder class.
Both come from a single throwaway model, since building one is not cheap.
"""
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)
model: PreTrainedModel = AutoModel.from_config(**self._from_config_kwargs)
model_classes = PreTrainedModelClasses(
decoder=type(model.get_decoder()),
encoders=self._find_encoder_classes(model),
)
del model
return decoder_cls
logger.debug("Identified model classes as: %s", model_classes)
return model_classes
def _decorate_cls_for_torch_compile(
self,
@@ -246,7 +278,7 @@ class Base(
is_encoder=is_encoder,
)(cls)
def _decorate_for_torch_compile(self, **kwargs: dict):
def _decorate_for_torch_compile(self):
"""
Decorate the model's decoder class to indicate to vLLM that it supports torch
compile if `can_enable_torch_compile` is True.
@@ -256,7 +288,7 @@ class Base(
class.
"""
self._decorate_cls_for_torch_compile(
cls=self._get_decoder_cls(**kwargs),
cls=self._pre_trained_model_classes.decoder,
# Applied to a PreTrainedModel so the batch dimension will exist
dynamic_arg_dims=dict[str, int](
input_ids=1, # shape: [1, seq_len]
vllm/model_executor/models/transformers/multimodal.py
+167 -113
View File
@@ -16,20 +16,24 @@
# limitations under the License.
"""Transformers modeling backend mixin for multi-modal models."""
from collections.abc import Mapping
from contextlib import nullcontext
from collections import defaultdict
from collections.abc import Mapping, Sequence
from contextlib import ExitStack, contextmanager, nullcontext
from typing import TYPE_CHECKING, Any
import torch
from transformers import AutoModel
from vllm.compilation.decorators import should_torch_compile_mm_encoder
from vllm.config.utils import getattr_iter
from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal
from vllm.model_executor.models.interfaces import (
MultiModalEmbeddings,
SupportsMRoPE,
SupportsMultiModal,
)
from vllm.model_executor.models.module_mapping import MultiModelKeys
from vllm.multimodal import MultiModalKwargsItems
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalFieldConfig,
@@ -37,6 +41,7 @@ from vllm.multimodal.inputs import (
)
from vllm.multimodal.parse import (
ImageProcessorItems,
ImageSize,
MultiModalDataItems,
MultiModalDataParser,
)
@@ -45,6 +50,7 @@ from vllm.multimodal.processing import (
BaseMultiModalProcessor,
BaseProcessingInfo,
ProcessorInputs,
PromptUpdate,
TimingContext,
)
from vllm.platforms import current_platform
@@ -74,8 +80,19 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
def _is_image_model(self) -> bool:
return hasattr(self.get_hf_processor(), "image_processor")
def _is_video_model(self) -> bool:
return hasattr(self.get_hf_processor(), "video_processor")
def _get_supported_modalities(self) -> list[str]:
modalities = []
if self._is_audio_model():
modalities.append("audio")
if self._is_image_model():
modalities.append("image")
if not modalities:
raise ValueError(
f"{type(self.get_hf_processor()).__name__} exposes neither an image "
"processor nor an audio processor, so the Transformers modeling "
"backend cannot serve this model as multi-modal."
)
return modalities
def _get_audio_token_id(self) -> int:
processor = self.get_hf_processor()
@@ -107,32 +124,16 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
return super().get_data_parser()
def get_supported_mm_limits(self):
limits = {}
if self._is_audio_model():
limits["audio"] = None
if self._is_image_model():
limits["image"] = None
if not limits:
raise ValueError(
f"Unable to detect a supported modality on "
f"{type(self.get_hf_processor()).__name__}. "
"Checked `_is_audio_model` and `_is_image_model`."
)
return limits
return dict.fromkeys(self._get_supported_modalities())
def get_mm_max_tokens_per_item(self, seq_len, mm_counts):
result = {}
if self._is_audio_model():
result["audio"] = self.get_max_audio_tokens()
if self._is_image_model():
result["image"] = self.get_max_image_tokens()
if not result:
raise ValueError(
f"Unable to detect a supported modality on "
f"{type(self.get_hf_processor()).__name__}. "
"Checked `_is_audio_model` and `_is_image_model`."
)
return result
modalities = self._get_supported_modalities()
max_tokens = {}
if "audio" in modalities:
max_tokens["audio"] = self.get_max_audio_tokens()
if "image" in modalities:
max_tokens["image"] = self.get_max_image_tokens()
return max_tokens
def get_max_audio_tokens(self) -> int:
config = self.get_hf_config()
@@ -148,7 +149,7 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
)
def get_max_image_tokens(self) -> int:
width, height = self.get_max_image_size()
width, height = self.get_image_size_with_most_features()
processor = self.get_hf_processor()
multimodal_config = self.ctx.model_config.multimodal_config
mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}
@@ -158,8 +159,8 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
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
def get_image_size_with_most_features(self) -> ImageSize:
return ImageSize(width=10_000, height=10_000) # arbitrary very large size
class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingInfo]):
@@ -199,7 +200,7 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn
overrides=mm_options.get("audio"),
)
if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)):
target_width, target_height = self.info.get_max_image_size()
target_width, target_height = self.info.get_image_size_with_most_features()
data["image"] = self._get_dummy_images(
width=target_width,
height=target_height,
@@ -215,21 +216,14 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
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.
) -> Sequence[PromptUpdate]:
"""No updates: `apply` locates placeholders via `mm_token_type_ids` instead.
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.
HF processors have no generic contract for the token sequence they insert,
so it cannot be expressed as a `PromptUpdate`. Returning nothing is only
safe because `apply` is overridden; the base class would reject it.
"""
return None
return []
def _get_modality_field_names(self, modality: str) -> set[str]:
"""Field names the sub-processor for `modality` produces."""
@@ -485,6 +479,21 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs),
)
# Bypassing `_maybe_apply_prompt_updates` also bypasses its validation.
# `_validate_mm_placeholders` can't be reused because it is typed for the
# `PlaceholderFeaturesInfo` the prompt update machinery produces.
mm_item_counts = mm_items.get_all_counts()
self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
for modality, item_count in mm_item_counts.items():
num_placeholders = len(mm_placeholders.get(modality, []))
if num_placeholders != item_count:
raise RuntimeError(
f"Expected there to be {item_count} prompt placeholders "
f"corresponding to {item_count} {modality} items, but instead "
f"found {num_placeholders} prompt placeholders! Make sure the "
"prompt contains a placeholder token for each item."
)
return mm_input(
prompt_token_ids=prompt_ids,
mm_kwargs=mm_kwargs,
@@ -498,65 +507,109 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
# Skip SupportsMRoPE.__init__ and call the next class in MRO
super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
def _get_encoder_cls(
self, modality: str = "image", **kwargs: dict
) -> type["PreTrainedModel"]:
"""
Get the encoder class from the model.
def _find_encoder_classes(
self, model: "PreTrainedModel"
) -> dict[str, type["PreTrainedModel"]]:
"""Modalities whose encoder cannot be told apart from the model itself are
omitted, as are those `get_encoder` rejects."""
encoder_classes: dict[str, type[PreTrainedModel]] = {}
for modality in _MODALITY_TO_TOKEN_TYPE_ID:
try:
encoder_cls = type(model.get_encoder(modality=modality))
except (TypeError, ValueError):
continue
if encoder_cls is not type(model):
encoder_classes[modality] = encoder_cls
return encoder_classes
Args:
kwargs: The kwargs to create the model.
@contextmanager
def _mark_model_components(self, vllm_config: "VllmConfig"):
model_config = vllm_config.model_config
encoder_classes = self._pre_trained_model_classes.encoders
if not encoder_classes:
logger.debug("No encoders identified, so no components will be marked")
yield
return
Returns:
The encoder class.
"""
with torch.device("meta"):
model: PreTrainedModel = AutoModel.from_config(**kwargs)
encoder_cls = type(model.get_encoder(modality=modality))
logger.debug("Identified encoder class as: %s", encoder_cls)
if type(model) is encoder_cls:
raise ValueError(
"Unable to infer vision encoder class from the model. "
"You must either: update the model so that "
"https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder"
" can detect the vision encoder correctly, or remove "
"'compile_mm_encoder'."
if model_config.skip_tokenizer_init:
# Determining the supported modalities needs the HF processor, which in
# turn needs a tokenizer
mm_config = model_config.multimodal_config
if mm_config.mm_encoder_only or any(
mm_config.get_limit_per_prompt(modality) == 0
for modality in encoder_classes
):
logger.warning_once(
"Unable to determine the supported modalities without a "
"tokenizer, so no model components will be skipped."
)
yield
return
# Modalities we don't serve report a limit of 999, which would stop their
# encoder ever being skipped
supported_modalities = MULTIMODAL_REGISTRY.get_processing_info(
model_config
).supported_mm_limits
# One encoder often serves several modalities, and may only be skipped when
# all of them are disabled, so mark it once for the whole set
modalities_by_encoder = defaultdict(set)
for modality, encoder_cls in encoder_classes.items():
if modality in supported_modalities:
modalities_by_encoder[encoder_cls].add(modality)
with ExitStack() as stack:
stack.enter_context(
self._mark_language_model(
vllm_config, targets=self._pre_trained_model_classes.decoder
)
)
del model
return encoder_cls
for encoder_cls, modalities in modalities_by_encoder.items():
stack.enter_context(
self._mark_tower_model(vllm_config, modalities, targets=encoder_cls)
)
yield
def _decorate_for_torch_compile(self, **kwargs: dict):
def _decorate_for_torch_compile(self):
"""
Decorate the model's decoder and encoder classes to indicate to vLLM that they
support torch compile if `can_enable_torch_compile` and
`should_torch_compile_mm_encoder` are True respectively.
Args:
kwargs: The kwargs to create the model, which are needed to get the decoder
and encoder classes.
"""
super()._decorate_for_torch_compile(**kwargs)
# Decorate the vision encoder model class to support torch compile if needed
super()._decorate_for_torch_compile()
# Decorate the encoder model classes to support torch compile if needed
if self.compilation_config.compile_mm_encoder:
self.check_version("5.0.0", "multimodal encoder compilation support")
encoder_classes = self._pre_trained_model_classes.encoders
if not encoder_classes:
raise ValueError(
"Unable to infer any encoder classes from the model. "
"You must either: update the model so that "
"https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder"
" can detect the encoders correctly, or remove "
"'compile_mm_encoder'."
)
logger.warning_once(
"Multimodal encoder compilation with the Transformers modeling backend "
"is an experimental feature. It relies on:\n"
"- The vision encoder being torch compilable.\n"
"- All vision encoder tensor inputs must be type hinted as either "
"- The encoder being torch compilable.\n"
"- All encoder tensor inputs must be type hinted as either "
"`torch.Tensor` or `torch.FloatTensor`.\n"
"- The 0-th dimension of all tensor inputs to the vision encoder being "
"the dynamic dimension (i.e., sequence length or number of patches).\n"
"- The 0-th dimension of all tensor inputs to the encoder being the "
"dynamic dimension (e.g. sequence length, number of patches).\n"
"Please report any issues you encounter to help us improve it."
)
self._decorate_cls_for_torch_compile(
cls=self._get_encoder_cls(**kwargs),
# TODO: properly infer dynamic_arg_dims based on the encoder's forward
# method signature. Currently we assume dim 0 for all tensor inputs.
dynamic_arg_dims=None,
enable_if=should_torch_compile_mm_encoder,
is_encoder=True,
)
# One encoder can serve several modalities, and must only be decorated once
for encoder_cls in dict.fromkeys(encoder_classes.values()):
self._decorate_cls_for_torch_compile(
cls=encoder_cls,
# TODO: properly infer dynamic_arg_dims based on the encoder's
# forward method signature. We assume dim 0 for all tensor inputs.
dynamic_arg_dims=None,
enable_if=should_torch_compile_mm_encoder,
is_encoder=True,
)
def forward(
self,
@@ -625,13 +678,25 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
tokens_per_item = total_tokens // total_expected
token_split_sizes = [s * tokens_per_item for s in split_sizes]
elif total_expected > 0:
# Mismatch (profiling with dummy data) - pad/truncate
# TODO: make this an error once we know profiling never relies on it
if total_tokens == 0:
raise ValueError(
"Encoder returned empty embeddings. "
f"Expected {total_expected} tokens from "
f"split_sizes={split_sizes}"
)
# Keep the counts out of the message: `warning_once` keys its cache on
# the args, so varying them would log on every new pair
logger.warning_once(
"Encoder returned a different number of tokens than expected; "
"padding or truncating to fit. The embeddings are not trustworthy "
"outside of memory profiling."
)
logger.debug(
"Encoder returned %s tokens but %s were expected",
total_tokens,
total_expected,
)
if total_tokens < total_expected:
repeat_factor = (total_expected + total_tokens - 1) // total_tokens
embeddings = embeddings.repeat(repeat_factor, 1)
@@ -642,14 +707,14 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
return list(torch.split(embeddings, token_split_sizes, dim=0))
def _embed_audio(self, **kwargs) -> list[torch.Tensor] | None:
self.check_version("5.13.0", "audio models support")
def _process_audio_input(self, **kwargs) -> list[torch.Tensor] | None:
input_features: torch.Tensor | None = kwargs.pop("input_features", None)
if input_features is None:
input_features = kwargs.pop("input_values", None)
if input_features is None:
return None
self.check_version("5.13.0", "audio models support")
num_audio_tokens = kwargs.pop("num_audio_tokens")
kwargs.pop("token_type_ids", None)
kwargs.pop("mm_token_type_ids", None)
@@ -668,7 +733,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
split_sizes = num_audio_tokens.flatten().tolist()
return self._split_embeddings(audio_embeddings, split_sizes)
def _embed_vision(self, **kwargs) -> list[torch.Tensor] | torch.Tensor | None:
def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None:
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`
@@ -676,7 +741,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
pixel_values = kwargs.pop("image_patches", None)
if image_embeds is not None:
return image_embeds
return [image_embeds]
if pixel_values is None:
return None
@@ -714,25 +779,14 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
split_sizes = num_image_patches.flatten().tolist()
return self._split_embeddings(vision_embeddings, split_sizes)
return vision_embeddings
return list(vision_embeddings)
def embed_multimodal(self, **kwargs):
embeddings: tuple[torch.Tensor, ...] = ()
if "input_features" in kwargs or "input_values" in kwargs:
audio_embeddings = self._embed_audio(**kwargs)
if audio_embeddings is not None:
embeddings += tuple(audio_embeddings)
if (
"pixel_values" in kwargs
or "image_embeds" in kwargs
or "image_patches" in kwargs
):
vision_embeddings = self._embed_vision(**kwargs)
if vision_embeddings is not None:
if isinstance(vision_embeddings, torch.Tensor):
embeddings += (vision_embeddings,)
else:
embeddings += tuple(vision_embeddings)
def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings:
# Each helper detects its own inputs. We are called once per modality, so the
# leftovers a helper forwards to the HF model can't belong to the other one.
embeddings: list[torch.Tensor] = []
for process_input in (self._process_audio_input, self._process_image_input):
embeddings.extend(process_input(**kwargs) or [])
return embeddings
def get_mrope_input_positions(