[Bugfix] Honor OPT token embedding width for prompt embeddings (#61736)

Signed-off-by: maithilijoshi20 <maithilij2003@gmail.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

tests/entrypoints/openai/completion/test_prompt_validation.py

@@ -64,6 +64,7 @@ def test_load_prompt_embeds(
     model_config = Mock(spec=ModelConfig)
     model_config.enable_prompt_embeds = True
     model_config.get_hidden_size.return_value = hidden_size
+    model_config.hf_text_config = object()
     model_config.dtype = dtype
 
     # construct arbitrary tensors of various dtypes, layouts, and sizes.

tests/models/language/generation/test_common.py

@@ -75,6 +75,9 @@ AITER_MODEL_LIST = [
             "facebook/opt-125m",  # opt
             marks=[pytest.mark.core_model, pytest.mark.cpu_model],
         ),
+        pytest.param(
+            "facebook/opt-350m",  # opt with projected token embeddings
+        ),
         pytest.param(
             "microsoft/phi-2",  # phi
             marks=[pytest.mark.core_model, pytest.mark.slow_test],

tests/models/registry.py

@@ -462,7 +462,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "OlmoHybridForCausalLM": _HfExamplesInfo("allenai/Olmo-Hybrid-7B"),
     "OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),
     "OPTForCausalLM": _HfExamplesInfo(
-        "facebook/opt-125m", {"1b": "facebook/opt-iml-max-1.3b"}
+        "facebook/opt-125m",
+        {"350m": "facebook/opt-350m", "1b": "facebook/opt-iml-max-1.3b"},
     ),
     "OrionForCausalLM": _HfExamplesInfo(
         "OrionStarAI/Orion-14B-Chat", trust_remote_code=True

tests/renderers/chat_utils/test_prompt_embeds.py

@@ -105,6 +105,7 @@ def _make_mock_model_config(*, enable_prompt_embeds: bool = True) -> mock.MagicM
     # `safe_load_prompt_embeds` pins each tensor to the model's hidden_size
     # and dtype, so the mock must return concrete values.
     mc.get_hidden_size.return_value = _MOCK_HIDDEN_SIZE
+    mc.hf_text_config = object()
     mc.dtype = _MOCK_DTYPE
     return mc
 

tests/renderers/test_completions.py

@@ -34,6 +34,7 @@ class MockModelConfig:
     tokenizer_revision = None
     tokenizer_mode = "auto"
     hf_config = MockHFConfig()
+    hf_text_config = hf_config
     encoder_config: dict[str, Any] | None = None
     enable_prompt_embeds: bool = True
     skip_tokenizer_init: bool = False

tests/renderers/test_prompt_embeds_payload_validation.py

@@ -65,6 +65,40 @@ def _saved_tensor(tensor: torch.Tensor) -> bytes:
     return buffer.getvalue()
 
 
+@pytest.mark.parametrize(
+    "model, hidden_size, input_size, buffer_size",
+    [
+        ("facebook/opt-350m", 1024, 512, 512),
+        ("facebook/opt-125m", 768, 768, 768),
+        ("openai/clip-vit-base-patch32", 512, 512, 1024),
+        ("google/siglip-base-patch16-224", 768, 768, 1024),
+    ],
+)
+def test_prompt_embeds_input_width(model, hidden_size, input_size, buffer_size):
+    """OPT uses its token width; dual encoders retain their text input width."""
+    from vllm.config import ModelConfig
+
+    config = ModelConfig(
+        model=model,
+        dtype="float32",
+        enable_prompt_embeds=True,
+    )
+    if "opt-" not in model:
+        config.hf_text_config.projection_dim = buffer_size
+    assert config.get_hidden_size() == hidden_size
+    assert config.get_inputs_embeds_size() == buffer_size
+
+    tensor = torch.randn(4, input_size)
+    loaded = safe_load_prompt_embeds(config, _encode(_saved_tensor(tensor)))
+    torch.testing.assert_close(loaded, tensor)
+
+    wrong_size = hidden_size if hidden_size != input_size else input_size + 1
+    with pytest.raises(VLLMValidationError, match="does not match"):
+        safe_load_prompt_embeds(
+            config, _encode(_saved_tensor(torch.randn(4, wrong_size)))
+        )
+
+
 UNPARSABLE_PAYLOADS = [
     pytest.param(b"", id="empty"),
     pytest.param(b"\x80", id="single-byte"),

tests/tokenizers_/test_detokenize.py

@@ -15,6 +15,7 @@ from vllm.tokenizers.detokenizer_utils import (
 )
 from vllm.tokenizers.mistral import MistralTokenizer
 from vllm.tokenizers.protocol import TokenizerLike
+from vllm.transformers_utils.repo_utils import with_retry
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.detokenizer import (
     FastIncrementalDetokenizer,
@@ -95,10 +96,13 @@ def _run_incremental_decode(
 
 @pytest.fixture
 def tokenizer(tokenizer_name):
-    return (
-        MistralTokenizer.from_pretrained(tokenizer_name)
-        if "mistral" in tokenizer_name
-        else AutoTokenizer.from_pretrained(tokenizer_name)
+    return with_retry(
+        lambda: (
+            MistralTokenizer.from_pretrained(tokenizer_name)
+            if "mistral" in tokenizer_name
+            else AutoTokenizer.from_pretrained(tokenizer_name)
+        ),
+        log_msg=f"Error loading tokenizer {tokenizer_name}",
     )
 
 

vllm/config/model.py

@@ -1622,8 +1622,8 @@ class ModelConfig:
     def get_inputs_embeds_size(self) -> int:
         # The size of inputs_embeds is usually identical to the size
         # of the hidden states, however there are exceptions, such as
-        # embedding models like CLIP and SigLIP
-        names = ("projection_dim", "projection_size")
+        # OPT's projected token embeddings and embedding models like CLIP/SigLIP.
+        names = ("word_embed_proj_dim", "projection_dim", "projection_size")
         return getattr_iter(
             self.hf_text_config, names, default_factory=self.get_hidden_size
         )

vllm/renderers/embed_utils.py

@@ -91,11 +91,17 @@ def safe_load_prompt_embeds(
             parameter="prompt_embeds",
         )
 
-    # Pin each tensor to the model's hidden_size. Validating here
+    # Pin each tensor to the model's text embedding width. Validating here
     # also transitively guarantees cross-tensor consistency for requests that
     # include multiple `prompt_embeds` parts, which is required by downstream
     # concatenation in `_build_mixed_prompt_embeds`.
-    expected_hidden_size = model_config.get_hidden_size()
+    # OPT projects token embeddings to hidden_size inside its decoder. CLIP
+    # and SigLIP still accept text embeddings at hidden_size, not projection_dim.
+    expected_hidden_size = getattr(
+        model_config.hf_text_config,
+        "word_embed_proj_dim",
+        model_config.get_hidden_size(),
+    )
     if tensor.shape[1] != expected_hidden_size:
         raise VLLMValidationError(
             f"`prompt_embeds` hidden_size {tensor.shape[1]} does not match "