[Bugfix][Multimodal] Encode text in mixed CLIP/SigLIP pooling batches (#53165)
Signed-off-by: Prudhvivuda <prudhvi12042001@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
4f686e182a
commit
342b8ebd8b
@@ -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(
|
||||
|
||||
@@ -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:])
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user