[2/N] Harden Transformers modelling backend multi-modal path (#51657)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Aug 10, 2026 at 13:57 UTC
cf8f3a3bb23ef99d956f3763bc7b6b85cacb125b
3 files changed
+276
-136
tests/models/transformers/test_backend.py
+54
@@ -10,6 +10,13 @@ from typing import Any
10
import pytest
11
import torch
12
import torch.nn as nn
13
+from transformers import AutoConfig, AutoModel
14
+
15
+from vllm.config import ModelConfig, VllmConfig
16
+from vllm.model_executor.models.interfaces import SupportsMultiModal
17
+from vllm.model_executor.models.transformers.base import Base
18
+from vllm.model_executor.models.transformers.multimodal import MultiModalMixin
19
+from vllm.model_executor.models.utils import StageMissingLayer
20
21
from ...conftest import HfRunner, VllmRunner
22
from ...utils import multi_gpu_test, prep_prompts
@@ -468,3 +475,50 @@ def test_replaced_embedding_exposes_one_vpe(vpe):
475
476
lm_head = ParallelLMHead(VOCAB_SIZE, HIDDEN_SIZE)
477
assert lm_head.tie_weights(composed.embed).weight is composed.embed.weight
478
+
479
+
480
+MULTIMODAL_MODEL = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
481
+
482
+
483
+class MarkingStub(SupportsMultiModal, nn.Module):
484
+ """Just enough of the backend to exercise component marking."""
485
+
486
+ _mark_model_components = MultiModalMixin._mark_model_components
487
+ _find_encoder_classes = MultiModalMixin._find_encoder_classes
488
+ _from_config_kwargs = Base._from_config_kwargs
489
+ _pre_trained_model_classes = Base._pre_trained_model_classes
490
+
491
+
492
+def build_marked_model(image_limit: int, skip_tokenizer_init: bool = False):
493
+ """Build the HF model inside the marking context and return it."""
494
+ model_config = ModelConfig(
495
+ model=MULTIMODAL_MODEL,
496
+ model_impl="transformers",
497
+ limit_mm_per_prompt={"image": image_limit},
498
+ )
499
+ # Set after construction: building the config itself needs the tokenizer
500
+ model_config.skip_tokenizer_init = skip_tokenizer_init
501
+ stub = MarkingStub()
502
+ stub.config = AutoConfig.from_pretrained(MULTIMODAL_MODEL)
503
+ stub.model_config = model_config
504
+
505
+ vllm_config = VllmConfig(model_config=model_config)
506
+ with stub._mark_model_components(vllm_config), torch.device("meta"):
507
+ stub.model = AutoModel.from_config(**stub._from_config_kwargs)
508
+ return stub.model
509
+
510
+
511
+@pytest.mark.parametrize(("image_limit", "skipped"), [(0, True), (4, False)])
512
+def test_tower_weights_skipped_when_modality_disabled(image_limit, skipped):
513
+ """`--limit-mm-per-prompt image=0` should drop the vision tower's weights."""
514
+ vision_tower = build_marked_model(image_limit).vision_tower
515
+ assert isinstance(vision_tower, StageMissingLayer) is skipped
516
+
517
+
518
+def test_marking_skipped_without_tokenizer():
519
+ """Marking needs the HF processor, which needs a tokenizer, so it is skipped.
520
+
521
+ The tower is built as normal rather than the model failing to load.
522
+ """
523
+ vision_tower = build_marked_model(0, skip_tokenizer_init=True).vision_tower
524
+ assert not isinstance(vision_tower, StageMissingLayer)
vllm/model_executor/models/transformers/base.py
+52
-20
@@ -18,9 +18,11 @@
18
19
import os
20
from collections.abc import Callable, Iterable
21
+from contextlib import contextmanager
22
+from functools import cached_property
23
from itertools import chain
24
from operator import attrgetter
23
-from typing import TYPE_CHECKING
25
+from typing import TYPE_CHECKING, Any, NamedTuple
26
27
import regex as re
28
import torch
@@ -85,6 +87,12 @@ if TYPE_CHECKING:
87
logger = init_logger(__name__)
88
89
90
+class PreTrainedModelClasses(NamedTuple):
91
+ decoder: type["PreTrainedModel"]
92
+ encoders: dict[str, type["PreTrainedModel"]]
93
+ """Modality -> encoder class, for each modality that has one."""
94
+
95
+
96
class Base(
97
nn.Module,
98
VllmModel,
@@ -147,14 +155,13 @@ class Base(
155
)
156
157
self._patch_config()
150
- from_config_kwargs = dict(
151
- config=self.config,
152
- dtype=self.model_config.dtype,
153
- trust_remote_code=self.model_config.trust_remote_code,
154
- )
155
- self._decorate_for_torch_compile(**from_config_kwargs)
158
+ self._decorate_for_torch_compile()
159
# Init on "meta" to delay allocating GPU tensors
157
- with init_on_device_without_buffers("meta"):
160
+ with (
161
+ self._mark_model_components(vllm_config),
162
+ init_on_device_without_buffers("meta"),
163
+ ):
164
+ from_config_kwargs = self._from_config_kwargs
165
self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs)
166
167
# Create weight name to module qualname mapper
@@ -196,22 +203,47 @@ class Base(
203
self.text_config._attn_implementation = "vllm"
204
self.config.dtype = torch.get_default_dtype()
205
199
- def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]:
206
+ @contextmanager
207
+ def _mark_model_components(self, vllm_config: "VllmConfig"):
208
+ """Mark language model and tower submodules as `self.model` is created.
209
+
210
+ Nothing to do in `Base`, `MultiModalMixin` will override."""
211
+ yield
212
+
213
+ @cached_property
214
+ def _from_config_kwargs(self) -> dict[str, Any]:
215
+ """The kwargs used to create `self.model`."""
216
+ return dict(
217
+ config=self.config,
218
+ dtype=self.model_config.dtype,
219
+ trust_remote_code=self.model_config.trust_remote_code,
220
+ )
221
+
222
+ def _find_encoder_classes(
223
+ self, model: "PreTrainedModel"
224
+ ) -> dict[str, type["PreTrainedModel"]]:
225
+ """Find the encoder class of each modality `model` has one for.
226
+
227
+ Text models have none. Multi-modal ones override this.
228
"""
201
- Get the decoder class from the model.
229
+ return {}
230
203
- Args:
204
- kwargs: The kwargs to create the model.
231
+ @cached_property
232
+ def _pre_trained_model_classes(self) -> PreTrainedModelClasses:
233
+ """The decoder class and the encoder class of each modality that has one.
234
206
- Returns:
207
- The decoder class.
235
+ Both come from a single throwaway model, since building one is not cheap.
236
"""
237
with torch.device("meta"):
210
- model: PreTrainedModel = AutoModel.from_config(**kwargs)
211
- decoder_cls = type(model.get_decoder())
212
- logger.debug("Identified decoder class as: %s", decoder_cls)
238
+ model: PreTrainedModel = AutoModel.from_config(**self._from_config_kwargs)
239
+ model_classes = PreTrainedModelClasses(
240
+ decoder=type(model.get_decoder()),
241
+ encoders=self._find_encoder_classes(model),
242
+ )
243
del model
214
- return decoder_cls
244
+
245
+ logger.debug("Identified model classes as: %s", model_classes)
246
+ return model_classes
247
248
def _decorate_cls_for_torch_compile(
249
self,
@@ -246,7 +278,7 @@ class Base(
278
is_encoder=is_encoder,
279
)(cls)
280
249
- def _decorate_for_torch_compile(self, **kwargs: dict):
281
+ def _decorate_for_torch_compile(self):
282
"""
283
Decorate the model's decoder class to indicate to vLLM that it supports torch
284
compile if `can_enable_torch_compile` is True.
@@ -256,7 +288,7 @@ class Base(
288
class.
289
"""
290
self._decorate_cls_for_torch_compile(
259
- cls=self._get_decoder_cls(**kwargs),
291
+ cls=self._pre_trained_model_classes.decoder,
292
# Applied to a PreTrainedModel so the batch dimension will exist
293
dynamic_arg_dims=dict[str, int](
294
input_ids=1, # shape: [1, seq_len]
vllm/model_executor/models/transformers/multimodal.py
+170
-116
@@ -16,20 +16,24 @@
16
# limitations under the License.
17
"""Transformers modeling backend mixin for multi-modal models."""
18
19
-from collections.abc import Mapping
20
-from contextlib import nullcontext
19
+from collections import defaultdict
20
+from collections.abc import Mapping, Sequence
21
+from contextlib import ExitStack, contextmanager, nullcontext
22
from typing import TYPE_CHECKING, Any
23
24
import torch
24
-from transformers import AutoModel
25
26
from vllm.compilation.decorators import should_torch_compile_mm_encoder
27
from vllm.config.utils import getattr_iter
28
from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input
29
from vllm.logger import init_logger
30
-from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal
30
+from vllm.model_executor.models.interfaces import (
31
+ MultiModalEmbeddings,
32
+ SupportsMRoPE,
33
+ SupportsMultiModal,
34
+)
35
from vllm.model_executor.models.module_mapping import MultiModelKeys
32
-from vllm.multimodal import MultiModalKwargsItems
36
+from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
37
from vllm.multimodal.inputs import (
38
MultiModalFeatureSpec,
39
MultiModalFieldConfig,
@@ -37,6 +41,7 @@ from vllm.multimodal.inputs import (
41
)
42
from vllm.multimodal.parse import (
43
ImageProcessorItems,
44
+ ImageSize,
45
MultiModalDataItems,
46
MultiModalDataParser,
47
)
@@ -45,6 +50,7 @@ from vllm.multimodal.processing import (
50
BaseMultiModalProcessor,
51
BaseProcessingInfo,
52
ProcessorInputs,
53
+ PromptUpdate,
54
TimingContext,
55
)
56
from vllm.platforms import current_platform
@@ -74,8 +80,19 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
80
def _is_image_model(self) -> bool:
81
return hasattr(self.get_hf_processor(), "image_processor")
82
77
- def _is_video_model(self) -> bool:
78
- return hasattr(self.get_hf_processor(), "video_processor")
83
+ def _get_supported_modalities(self) -> list[str]:
84
+ modalities = []
85
+ if self._is_audio_model():
86
+ modalities.append("audio")
87
+ if self._is_image_model():
88
+ modalities.append("image")
89
+ if not modalities:
90
+ raise ValueError(
91
+ f"{type(self.get_hf_processor()).__name__} exposes neither an image "
92
+ "processor nor an audio processor, so the Transformers modeling "
93
+ "backend cannot serve this model as multi-modal."
94
+ )
95
+ return modalities
96
97
def _get_audio_token_id(self) -> int:
98
processor = self.get_hf_processor()
@@ -107,32 +124,16 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
124
return super().get_data_parser()
125
126
def get_supported_mm_limits(self):
110
- limits = {}
111
- if self._is_audio_model():
112
- limits["audio"] = None
113
- if self._is_image_model():
114
- limits["image"] = None
115
- if not limits:
116
- raise ValueError(
117
- f"Unable to detect a supported modality on "
118
- f"{type(self.get_hf_processor()).__name__}. "
119
- "Checked `_is_audio_model` and `_is_image_model`."
120
- )
121
- return limits
127
+ return dict.fromkeys(self._get_supported_modalities())
128
129
def get_mm_max_tokens_per_item(self, seq_len, mm_counts):
124
- result = {}
125
- if self._is_audio_model():
126
- result["audio"] = self.get_max_audio_tokens()
127
- if self._is_image_model():
128
- result["image"] = self.get_max_image_tokens()
129
- if not result:
130
- raise ValueError(
131
- f"Unable to detect a supported modality on "
132
- f"{type(self.get_hf_processor()).__name__}. "
133
- "Checked `_is_audio_model` and `_is_image_model`."
134
- )
135
- return result
130
+ modalities = self._get_supported_modalities()
131
+ max_tokens = {}
132
+ if "audio" in modalities:
133
+ max_tokens["audio"] = self.get_max_audio_tokens()
134
+ if "image" in modalities:
135
+ max_tokens["image"] = self.get_max_image_tokens()
136
+ return max_tokens
137
138
def get_max_audio_tokens(self) -> int:
139
config = self.get_hf_config()
@@ -148,7 +149,7 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
149
)
150
151
def get_max_image_tokens(self) -> int:
151
- width, height = self.get_max_image_size()
152
+ width, height = self.get_image_size_with_most_features()
153
processor = self.get_hf_processor()
154
multimodal_config = self.ctx.model_config.multimodal_config
155
mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}
@@ -158,8 +159,8 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
159
image_tokens = mm_tokens["num_image_tokens"][0]
160
return image_tokens
161
161
- def get_max_image_size(self):
162
- return 10_000, 10_000 # hardcode for arbitrary very large size
162
+ def get_image_size_with_most_features(self) -> ImageSize:
163
+ return ImageSize(width=10_000, height=10_000) # arbitrary very large size
164
165
166
class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingInfo]):
@@ -199,7 +200,7 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn
200
overrides=mm_options.get("audio"),
201
)
202
if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)):
202
- target_width, target_height = self.info.get_max_image_size()
203
+ target_width, target_height = self.info.get_image_size_with_most_features()
204
data["image"] = self._get_dummy_images(
205
width=target_width,
206
height=target_height,
@@ -215,21 +216,14 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
216
mm_items: MultiModalDataItems,
217
hf_processor_mm_kwargs: Mapping[str, object],
218
out_mm_kwargs: MultiModalKwargsItems,
218
- ):
219
- """
220
- Given the original multi-modal items for this modality
221
- and HF-processed data, output the updates to perform.
222
-
223
- The information returned by this method is used to update token inputs
224
- which bypass the HF processor. It is also used to update the output of
225
- HF processor if the HF process does not apply prompt updates to text
226
- inputs.
219
+ ) -> Sequence[PromptUpdate]:
220
+ """No updates: `apply` locates placeholders via `mm_token_type_ids` instead.
221
228
- Moreover, this information is critical to determine the token positions
229
- in order to construct :class:`~vllm-multimodal.input.PlaceholderRange`
230
- for each multi-modal item.
222
+ HF processors have no generic contract for the token sequence they insert,
223
+ so it cannot be expressed as a `PromptUpdate`. Returning nothing is only
224
+ safe because `apply` is overridden; the base class would reject it.
225
"""
232
- return None
226
+ return []
227
228
def _get_modality_field_names(self, modality: str) -> set[str]:
229
"""Field names the sub-processor for `modality` produces."""
@@ -485,6 +479,21 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
479
self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs),
480
)
481
482
+ # Bypassing `_maybe_apply_prompt_updates` also bypasses its validation.
483
+ # `_validate_mm_placeholders` can't be reused because it is typed for the
484
+ # `PlaceholderFeaturesInfo` the prompt update machinery produces.
485
+ mm_item_counts = mm_items.get_all_counts()
486
+ self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
487
+ for modality, item_count in mm_item_counts.items():
488
+ num_placeholders = len(mm_placeholders.get(modality, []))
489
+ if num_placeholders != item_count:
490
+ raise RuntimeError(
491
+ f"Expected there to be {item_count} prompt placeholders "
492
+ f"corresponding to {item_count} {modality} items, but instead "
493
+ f"found {num_placeholders} prompt placeholders! Make sure the "
494
+ "prompt contains a placeholder token for each item."
495
+ )
496
+
497
return mm_input(
498
prompt_token_ids=prompt_ids,
499
mm_kwargs=mm_kwargs,
@@ -498,65 +507,109 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
507
# Skip SupportsMRoPE.__init__ and call the next class in MRO
508
super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
509
501
- def _get_encoder_cls(
502
- self, modality: str = "image", **kwargs: dict
503
- ) -> type["PreTrainedModel"]:
504
- """
505
- Get the encoder class from the model.
506
-
507
- Args:
508
- kwargs: The kwargs to create the model.
509
-
510
- Returns:
511
- The encoder class.
512
- """
513
- with torch.device("meta"):
514
- model: PreTrainedModel = AutoModel.from_config(**kwargs)
515
- encoder_cls = type(model.get_encoder(modality=modality))
516
- logger.debug("Identified encoder class as: %s", encoder_cls)
517
- if type(model) is encoder_cls:
518
- raise ValueError(
519
- "Unable to infer vision encoder class from the model. "
520
- "You must either: update the model so that "
521
- "https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder"
522
- " can detect the vision encoder correctly, or remove "
523
- "'compile_mm_encoder'."
510
+ def _find_encoder_classes(
511
+ self, model: "PreTrainedModel"
512
+ ) -> dict[str, type["PreTrainedModel"]]:
513
+ """Modalities whose encoder cannot be told apart from the model itself are
514
+ omitted, as are those `get_encoder` rejects."""
515
+ encoder_classes: dict[str, type[PreTrainedModel]] = {}
516
+ for modality in _MODALITY_TO_TOKEN_TYPE_ID:
517
+ try:
518
+ encoder_cls = type(model.get_encoder(modality=modality))
519
+ except (TypeError, ValueError):
520
+ continue
521
+ if encoder_cls is not type(model):
522
+ encoder_classes[modality] = encoder_cls
523
+ return encoder_classes
524
+
525
+ @contextmanager
526
+ def _mark_model_components(self, vllm_config: "VllmConfig"):
527
+ model_config = vllm_config.model_config
528
+ encoder_classes = self._pre_trained_model_classes.encoders
529
+ if not encoder_classes:
530
+ logger.debug("No encoders identified, so no components will be marked")
531
+ yield
532
+ return
533
+
534
+ if model_config.skip_tokenizer_init:
535
+ # Determining the supported modalities needs the HF processor, which in
536
+ # turn needs a tokenizer
537
+ mm_config = model_config.multimodal_config
538
+ if mm_config.mm_encoder_only or any(
539
+ mm_config.get_limit_per_prompt(modality) == 0
540
+ for modality in encoder_classes
541
+ ):
542
+ logger.warning_once(
543
+ "Unable to determine the supported modalities without a "
544
+ "tokenizer, so no model components will be skipped."
545
+ )
546
+ yield
547
+ return
548
+
549
+ # Modalities we don't serve report a limit of 999, which would stop their
550
+ # encoder ever being skipped
551
+ supported_modalities = MULTIMODAL_REGISTRY.get_processing_info(
552
+ model_config
553
+ ).supported_mm_limits
554
+
555
+ # One encoder often serves several modalities, and may only be skipped when
556
+ # all of them are disabled, so mark it once for the whole set
557
+ modalities_by_encoder = defaultdict(set)
558
+ for modality, encoder_cls in encoder_classes.items():
559
+ if modality in supported_modalities:
560
+ modalities_by_encoder[encoder_cls].add(modality)
561
+
562
+ with ExitStack() as stack:
563
+ stack.enter_context(
564
+ self._mark_language_model(
565
+ vllm_config, targets=self._pre_trained_model_classes.decoder
566
+ )
567
)
525
- del model
526
- return encoder_cls
568
+ for encoder_cls, modalities in modalities_by_encoder.items():
569
+ stack.enter_context(
570
+ self._mark_tower_model(vllm_config, modalities, targets=encoder_cls)
571
+ )
572
+ yield
573
528
- def _decorate_for_torch_compile(self, **kwargs: dict):
574
+ def _decorate_for_torch_compile(self):
575
"""
576
Decorate the model's decoder and encoder classes to indicate to vLLM that they
577
support torch compile if `can_enable_torch_compile` and
578
`should_torch_compile_mm_encoder` are True respectively.
533
-
534
- Args:
535
- kwargs: The kwargs to create the model, which are needed to get the decoder
536
- and encoder classes.
579
"""
538
- super()._decorate_for_torch_compile(**kwargs)
539
- # Decorate the vision encoder model class to support torch compile if needed
580
+ super()._decorate_for_torch_compile()
581
+ # Decorate the encoder model classes to support torch compile if needed
582
if self.compilation_config.compile_mm_encoder:
583
self.check_version("5.0.0", "multimodal encoder compilation support")
584
+ encoder_classes = self._pre_trained_model_classes.encoders
585
+ if not encoder_classes:
586
+ raise ValueError(
587
+ "Unable to infer any encoder classes from the model. "
588
+ "You must either: update the model so that "
589
+ "https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder"
590
+ " can detect the encoders correctly, or remove "
591
+ "'compile_mm_encoder'."
592
+ )
593
logger.warning_once(
594
"Multimodal encoder compilation with the Transformers modeling backend "
595
"is an experimental feature. It relies on:\n"
545
- "- The vision encoder being torch compilable.\n"
546
- "- All vision encoder tensor inputs must be type hinted as either "
596
+ "- The encoder being torch compilable.\n"
597
+ "- All encoder tensor inputs must be type hinted as either "
598
"`torch.Tensor` or `torch.FloatTensor`.\n"
548
- "- The 0-th dimension of all tensor inputs to the vision encoder being "
549
- "the dynamic dimension (i.e., sequence length or number of patches).\n"
599
+ "- The 0-th dimension of all tensor inputs to the encoder being the "
600
+ "dynamic dimension (e.g. sequence length, number of patches).\n"
601
"Please report any issues you encounter to help us improve it."
602
)
552
- self._decorate_cls_for_torch_compile(
553
- cls=self._get_encoder_cls(**kwargs),
554
- # TODO: properly infer dynamic_arg_dims based on the encoder's forward
555
- # method signature. Currently we assume dim 0 for all tensor inputs.
556
- dynamic_arg_dims=None,
557
- enable_if=should_torch_compile_mm_encoder,
558
- is_encoder=True,
559
- )
603
+ # One encoder can serve several modalities, and must only be decorated once
604
+ for encoder_cls in dict.fromkeys(encoder_classes.values()):
605
+ self._decorate_cls_for_torch_compile(
606
+ cls=encoder_cls,
607
+ # TODO: properly infer dynamic_arg_dims based on the encoder's
608
+ # forward method signature. We assume dim 0 for all tensor inputs.
609
+ dynamic_arg_dims=None,
610
+ enable_if=should_torch_compile_mm_encoder,
611
+ is_encoder=True,
612
+ )
613
614
def forward(
615
self,
@@ -625,13 +678,25 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
678
tokens_per_item = total_tokens // total_expected
679
token_split_sizes = [s * tokens_per_item for s in split_sizes]
680
elif total_expected > 0:
628
- # Mismatch (profiling with dummy data) - pad/truncate
681
+ # TODO: make this an error once we know profiling never relies on it
682
if total_tokens == 0:
683
raise ValueError(
684
"Encoder returned empty embeddings. "
685
f"Expected {total_expected} tokens from "
686
f"split_sizes={split_sizes}"
687
)
688
+ # Keep the counts out of the message: `warning_once` keys its cache on
689
+ # the args, so varying them would log on every new pair
690
+ logger.warning_once(
691
+ "Encoder returned a different number of tokens than expected; "
692
+ "padding or truncating to fit. The embeddings are not trustworthy "
693
+ "outside of memory profiling."
694
+ )
695
+ logger.debug(
696
+ "Encoder returned %s tokens but %s were expected",
697
+ total_tokens,
698
+ total_expected,
699
+ )
700
if total_tokens < total_expected:
701
repeat_factor = (total_expected + total_tokens - 1) // total_tokens
702
embeddings = embeddings.repeat(repeat_factor, 1)
@@ -642,14 +707,14 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
707
708
return list(torch.split(embeddings, token_split_sizes, dim=0))
709
645
- def _embed_audio(self, **kwargs) -> list[torch.Tensor] | None:
646
- self.check_version("5.13.0", "audio models support")
710
+ def _process_audio_input(self, **kwargs) -> list[torch.Tensor] | None:
711
input_features: torch.Tensor | None = kwargs.pop("input_features", None)
712
if input_features is None:
713
input_features = kwargs.pop("input_values", None)
714
if input_features is None:
715
return None
716
717
+ self.check_version("5.13.0", "audio models support")
718
num_audio_tokens = kwargs.pop("num_audio_tokens")
719
kwargs.pop("token_type_ids", None)
720
kwargs.pop("mm_token_type_ids", None)
@@ -668,7 +733,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
733
split_sizes = num_audio_tokens.flatten().tolist()
734
return self._split_embeddings(audio_embeddings, split_sizes)
735
671
- def _embed_vision(self, **kwargs) -> list[torch.Tensor] | torch.Tensor | None:
736
+ def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None:
737
pixel_values: torch.Tensor | None = kwargs.pop("pixel_values", None)
738
image_embeds: torch.Tensor | None = kwargs.pop("image_embeds", None)
739
# Model might use `image_patches` instead of `pixel_values`
@@ -676,7 +741,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
741
pixel_values = kwargs.pop("image_patches", None)
742
743
if image_embeds is not None:
679
- return image_embeds
744
+ return [image_embeds]
745
746
if pixel_values is None:
747
return None
@@ -714,25 +779,14 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
779
split_sizes = num_image_patches.flatten().tolist()
780
return self._split_embeddings(vision_embeddings, split_sizes)
781
717
- return vision_embeddings
718
-
719
- def embed_multimodal(self, **kwargs):
720
- embeddings: tuple[torch.Tensor, ...] = ()
721
- if "input_features" in kwargs or "input_values" in kwargs:
722
- audio_embeddings = self._embed_audio(**kwargs)
723
- if audio_embeddings is not None:
724
- embeddings += tuple(audio_embeddings)
725
- if (
726
- "pixel_values" in kwargs
727
- or "image_embeds" in kwargs
728
- or "image_patches" in kwargs
729
- ):
730
- vision_embeddings = self._embed_vision(**kwargs)
731
- if vision_embeddings is not None:
732
- if isinstance(vision_embeddings, torch.Tensor):
733
- embeddings += (vision_embeddings,)
734
- else:
735
- embeddings += tuple(vision_embeddings)
782
+ return list(vision_embeddings)
783
+
784
+ def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings:
785
+ # Each helper detects its own inputs. We are called once per modality, so the
786
+ # leftovers a helper forwards to the HF model can't belong to the other one.
787
+ embeddings: list[torch.Tensor] = []
788
+ for process_input in (self._process_audio_input, self._process_image_input):
789
+ embeddings.extend(process_input(**kwargs) or [])
790
return embeddings
791
792
def get_mrope_input_positions(