diff --git a/tests/models/multimodal/pooling/test_clip.py b/tests/models/multimodal/pooling/test_clip.py index 54dd8556e9..9c500fd6a3 100644 --- a/tests/models/multimodal/pooling/test_clip.py +++ b/tests/models/multimodal/pooling/test_clip.py @@ -52,6 +52,31 @@ def _run_test( vllm_model.embed(texts) vllm_model.embed([""], images=images) + # Mixed image+text batch must not skip the text encoder (#53091). + mixed_outputs = vllm_outputs_per_case[2] + check_embeddings_close( + embeddings_0_lst=[vllm_outputs_per_case[0][0]], + embeddings_1_lst=[mixed_outputs[0]], + name_0="text_only", + name_1="mixed_text", + ) + check_embeddings_close( + embeddings_0_lst=[vllm_outputs_per_case[1][0]], + embeddings_1_lst=[mixed_outputs[1]], + name_0="image_only", + name_1="mixed_image", + ) + empty_text = vllm_model.embed([""]) + empty_sim = torch.nn.functional.cosine_similarity( + torch.tensor(mixed_outputs[0]), + torch.tensor(empty_text[0]), + dim=0, + ) + assert empty_sim < 0.99, ( + "Mixed-batch text embedding collapsed to the empty-string vector " + f"(cosine={empty_sim:.4f})" + ) + with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model: hf_outputs_per_case = [] for input_texts, input_images in input_cases: @@ -101,6 +126,7 @@ def test_models( input_cases = [ (HF_TEXT_PROMPTS, text_images), (HF_IMAGE_PROMPTS, images), + ([HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]], [None, images[0]]), ] _run_test( diff --git a/tests/models/multimodal/pooling/test_dual_encoder_routing.py b/tests/models/multimodal/pooling/test_dual_encoder_routing.py new file mode 100644 index 0000000000..9e7ad1909f --- /dev/null +++ b/tests/models/multimodal/pooling/test_dual_encoder_routing.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.model_executor.models.clip import ( + dual_encoder_has_text_tokens, + merge_dual_encoder_text_and_vision, +) + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + + +def test_text_only_batch_runs_text_encoder(): + assert dual_encoder_has_text_tokens(False, None) + assert dual_encoder_has_text_tokens(False, torch.zeros(4, dtype=torch.bool)) + + +def test_vision_only_batch_skips_text_encoder(): + is_multimodal = torch.ones(8, dtype=torch.bool) + assert not dual_encoder_has_text_tokens(True, is_multimodal) + + +def test_mixed_batch_runs_text_encoder(): + # Image tokens followed by a text sequence: the old batch-wide flag + # treated this as vision-only and skipped the text encoder (#53091). + is_multimodal = torch.tensor([True, True, True, False, False, False]) + assert dual_encoder_has_text_tokens(True, is_multimodal) + + +def test_missing_token_mask_with_mm_embeddings_is_vision_only(): + assert not dual_encoder_has_text_tokens(True, None) + + +def test_merge_keeps_vision_on_mm_tokens(): + text = torch.zeros(4, 2) + vision = torch.ones(4, 2) + is_multimodal = torch.tensor([True, True, False, False]) + out = merge_dual_encoder_text_and_vision(text, vision, is_multimodal) + assert torch.equal(out[:2], vision[:2]) + assert torch.equal(out[2:], text[2:]) diff --git a/tests/models/multimodal/pooling/test_siglip.py b/tests/models/multimodal/pooling/test_siglip.py index 8eeb594db2..df9bb61353 100644 --- a/tests/models/multimodal/pooling/test_siglip.py +++ b/tests/models/multimodal/pooling/test_siglip.py @@ -63,6 +63,20 @@ def _run_test( vllm_model.embed(texts) vllm_model.embed([""], images=images) + mixed_outputs = vllm_outputs_per_case[2] + check_embeddings_close( + embeddings_0_lst=[vllm_outputs_per_case[0][0]], + embeddings_1_lst=[mixed_outputs[0]], + name_0="text_only", + name_1="mixed_text", + ) + check_embeddings_close( + embeddings_0_lst=[vllm_outputs_per_case[1][0]], + embeddings_1_lst=[mixed_outputs[1]], + name_0="image_only", + name_1="mixed_image", + ) + with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model: hf_outputs_per_case = [] for input_texts, input_images, tokenization_kwargs in input_cases: @@ -122,6 +136,14 @@ def test_models( }, ), (HF_IMAGE_PROMPTS, images, {}), + ( + [HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]], + [None, images[0]], + { + "padding": "max_length", + "max_length": 64, + }, + ), ] _run_test( diff --git a/vllm/model_executor/models/clip.py b/vllm/model_executor/models/clip.py index 4294a203df..f0be4890ac 100644 --- a/vllm/model_executor/models/clip.py +++ b/vllm/model_executor/models/clip.py @@ -119,6 +119,38 @@ def _get_vision_feature_select_strategy(pooling_type: str): ) from None +def dual_encoder_has_text_tokens( + has_mm_embeddings: bool, + is_multimodal: torch.Tensor | None, +) -> bool: + """Whether a dual-encoder pooling batch still contains text tokens. + + ``embed_input_ids`` used to treat "the batch has any image embeddings" as + "the whole batch is vision-only". Mixed image+text batches then skipped + the text encoder, so completion-style text embeddings collapsed (issue + #53091). Honor the per-token ``is_multimodal`` mask instead. + """ + if not has_mm_embeddings or is_multimodal is None: + return not has_mm_embeddings + return bool((~is_multimodal).any().item()) + + +def merge_dual_encoder_text_and_vision( + text_features: torch.Tensor, + vision_embeds: torch.Tensor, + is_multimodal: torch.Tensor | None, +) -> torch.Tensor: + """Restore vision embeddings on multimodal tokens after the text encoder. + + The model runner builds ``is_multimodal`` on CPU; text and vision features + are on the compute device. + """ + if is_multimodal is None or not bool(is_multimodal.any().item()): + return text_features + mm_mask = is_multimodal.to(device=vision_embeds.device, non_blocking=True) + return torch.where(mm_mask.unsqueeze(-1), vision_embeds, text_features) + + class CLIPProcessingInfo(BaseProcessingInfo): def get_hf_config(self): return self.ctx.get_hf_config(CLIPConfig) @@ -809,8 +841,9 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): self.pooler = DispatchPooler.for_embedding(pooler_config) - # Assumes that self.forward is called after self.embed_input_ids - self._is_text_input = True + # Set in embed_input_ids; consumed by forward. + self._has_text_tokens = True + self._mm_token_mask: torch.Tensor | None = None def get_text_features( self, @@ -908,8 +941,12 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): *, is_multimodal: torch.Tensor | None = None, ) -> torch.Tensor: - self._is_text_input = ( - multimodal_embeddings is None or len(multimodal_embeddings) == 0 + has_mm_embeddings = ( + multimodal_embeddings is not None and len(multimodal_embeddings) > 0 + ) + self._mm_token_mask = is_multimodal + self._has_text_tokens = dual_encoder_has_text_tokens( + has_mm_embeddings, is_multimodal ) # This is to satisfy the type checker for each overload @@ -943,9 +980,10 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): # Multimodal inputs assert inputs_embeds is not None - if not self._is_text_input: + if not self._has_text_tokens: return inputs_embeds + vision_embeds = inputs_embeds # NOTE: inputs_embeds in model runner has size text_config.projection_dim # (instead of text_config.hidden_size) to accommodate image embeddings hidden_size = self.text_embed_dim @@ -955,7 +993,10 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): # No need to handle this case for now raise NotImplementedError - return self.get_text_features(input_ids, positions, inputs_embeds) + text_features = self.get_text_features(input_ids, positions, inputs_embeds) + return merge_dual_encoder_text_and_vision( + text_features, vision_embeds, self._mm_token_mask + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader( diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index 73f3c35637..cf325b64f9 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -57,6 +57,7 @@ from vllm.sequence import IntermediateTensors from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape +from .clip import dual_encoder_has_text_tokens, merge_dual_encoder_text_and_vision from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant from .interfaces_base import default_pooling_type from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -943,7 +944,9 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): self.pooler = DispatchPooler.for_embedding(pooler_config) - self._is_text_input = True + # Set in embed_input_ids; consumed by forward. + self._has_text_tokens = True + self._mm_token_mask: torch.Tensor | None = None def get_text_features( self, @@ -1088,8 +1091,12 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): *, is_multimodal: torch.Tensor | None = None, ) -> torch.Tensor: - self._is_text_input = ( - multimodal_embeddings is None or len(multimodal_embeddings) == 0 + has_mm_embeddings = ( + multimodal_embeddings is not None and len(multimodal_embeddings) > 0 + ) + self._mm_token_mask = is_multimodal + self._has_text_tokens = dual_encoder_has_text_tokens( + has_mm_embeddings, is_multimodal ) if multimodal_embeddings is None or is_multimodal is None: @@ -1121,9 +1128,11 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): raise RuntimeError("PP is not supported for this model") # Multimodal inputs (image embeddings) - if not self._is_text_input: + assert inputs_embeds is not None + if not self._has_text_tokens: return inputs_embeds + vision_embeds = inputs_embeds # NOTE: inputs_embeds in model runner has size text_config.projection_size # (instead of text_config.hidden_size) to accommodate image embeddings hidden_size = self.text_embed_dim @@ -1133,7 +1142,10 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): # No need to handle this case for now raise NotImplementedError - return self.get_text_features(input_ids, positions, inputs_embeds) + text_features = self.get_text_features(input_ids, positions, inputs_embeds) + return merge_dual_encoder_text_and_vision( + text_features, vision_embeds, self._mm_token_mask + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader(