support hybrid dtype

Author: noooop

.buildkite/test-pipeline.yaml

@@ -475,6 +475,7 @@ steps:
     - pytest -v -s models/test_utils.py
     - pytest -v -s models/test_vision.py
     - pytest -v -s models/test_initialization.py
+    - pytest -v -s models/test_hybrid_dtype.py
 
 - label: Language Models Test (Standard)
   mirror_hardwares: [amdexperimental]

tests/models/language/pooling/embed_utils.py

@@ -48,7 +48,12 @@ def correctness_test_embed_models(hf_runner,
     example_prompts = [str(s).strip() for s in example_prompts]
 
     vllm_extra_kwargs = vllm_extra_kwargs or {}
-    vllm_extra_kwargs["dtype"] = model_info.dtype
+
+    if isinstance(model_info.dtype, str):
+        vllm_extra_kwargs["dtype"] = model_info.dtype
+    else:
+        vllm_extra_kwargs["dtype"] = model_info.dtype.dtype
+        vllm_extra_kwargs["attn_dtype"] = model_info.dtype.attn_dtype
 
     with vllm_runner(model_info.name,
                      task="embed",

tests/models/language/pooling/mteb_utils.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# ruff: noqa: SIM117
 from collections.abc import Sequence
 
 import mteb
@@ -88,7 +89,13 @@ def mteb_test_embed_models(hf_runner,
         pytest.skip("Skipping test.")
 
     vllm_extra_kwargs = vllm_extra_kwargs or {}
-    vllm_extra_kwargs["dtype"] = model_info.dtype
+
+    if isinstance(model_info.dtype, str):
+        vllm_extra_kwargs["dtype"] = model_info.dtype
+
+    else:
+        vllm_extra_kwargs["dtype"] = model_info.dtype.dtype
+        vllm_extra_kwargs["attn_dtype"] = model_info.dtype.attn_dtype
 
     with vllm_runner(model_info.name,
                      task="embed",
@@ -102,18 +109,23 @@ def mteb_test_embed_models(hf_runner,
         vllm_main_score = run_mteb_embed_task(VllmMtebEncoder(vllm_model),
                                               MTEB_EMBED_TASKS)
         vllm_dtype = vllm_model.model.llm_engine.model_config.dtype
+        model_dtype = getattr(
+            vllm_model.model.llm_engine.model_config.hf_config, "torch_dtype",
+            vllm_dtype)
 
-    with set_default_torch_dtype(vllm_dtype) and hf_runner(
-            model_info.name, is_sentence_transformer=True,
-            dtype=vllm_dtype) as hf_model:
+    with set_default_torch_dtype(model_dtype):
+        with hf_runner(model_info.name,
+                       is_sentence_transformer=True,
+                       dtype=model_dtype) as hf_model:
 
-        if hf_model_callback is not None:
-            hf_model_callback(hf_model)
+            if hf_model_callback is not None:
+                hf_model_callback(hf_model)
 
-        st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
+            st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
+            st_dtype = next(hf_model.model.parameters()).dtype
 
-    print("VLLM:", vllm_main_score)
-    print("SentenceTransformers:", st_main_score)
+    print("VLLM:", vllm_dtype, vllm_main_score)
+    print("SentenceTransformers:", model_dtype, st_dtype, st_main_score)
     print("Difference:", st_main_score - vllm_main_score)
 
     assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_EMBED_TOL)

tests/models/language/pooling/test_gte.py

@@ -10,27 +10,24 @@
     ########## BertModel
     EmbedModelInfo("thenlper/gte-large",
                    architecture="BertModel",
-                   dtype="float32",
+                   dtype="hybrid",
                    enable_test=True),
     EmbedModelInfo("thenlper/gte-base",
                    architecture="BertModel",
-                   dtype="float32",
+                   dtype="hybrid",
                    enable_test=False),
     EmbedModelInfo("thenlper/gte-small",
                    architecture="BertModel",
-                   dtype="float32",
+                   dtype="hybrid",
                    enable_test=False),
     EmbedModelInfo("thenlper/gte-large-zh",
                    architecture="BertModel",
-                   dtype="float32",
                    enable_test=False),
     EmbedModelInfo("thenlper/gte-base-zh",
                    architecture="BertModel",
-                   dtype="float32",
                    enable_test=False),
     EmbedModelInfo("thenlper/gte-small-zh",
                    architecture="BertModel",
-                   dtype="float32",
                    enable_test=False),
     ########### NewModel
     EmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
@@ -45,7 +42,6 @@
     ########### Qwen2ForCausalLM
     EmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
                    architecture="Qwen2ForCausalLM",
-                   dtype="float32",
                    enable_test=True),
     ########## ModernBertModel
     EmbedModelInfo("Alibaba-NLP/gte-modernbert-base",

tests/models/language/pooling/test_intfloat.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+
+from .embed_utils import EmbedModelInfo, correctness_test_embed_models
+from .mteb_utils import mteb_test_embed_models
+
+MODELS = [
+    ########## BertModel
+    EmbedModelInfo("intfloat/e5-small",
+                   architecture="BertModel",
+                   enable_test=True),
+    EmbedModelInfo("intfloat/e5-base",
+                   architecture="BertModel",
+                   enable_test=False),
+    EmbedModelInfo("intfloat/e5-large",
+                   architecture="BertModel",
+                   enable_test=False),
+    EmbedModelInfo("intfloat/multilingual-e5-small",
+                   architecture="BertModel",
+                   enable_test=False),
+    ########## XLMRobertaModel
+    EmbedModelInfo("intfloat/multilingual-e5-base",
+                   architecture="XLMRobertaModel",
+                   enable_test=True),
+    EmbedModelInfo("intfloat/multilingual-e5-large",
+                   architecture="XLMRobertaModel",
+                   enable_test=False),
+    EmbedModelInfo("intfloat/multilingual-e5-large-instruct",
+                   architecture="XLMRobertaModel",
+                   dtype="hybrid",
+                   enable_test=False),
+]
+
+
+@pytest.mark.parametrize("model_info", MODELS)
+def test_embed_models_mteb(hf_runner, vllm_runner,
+                           model_info: EmbedModelInfo) -> None:
+    mteb_test_embed_models(hf_runner, vllm_runner, model_info)
+
+
+@pytest.mark.parametrize("model_info", MODELS)
+def test_embed_models_correctness(hf_runner, vllm_runner,
+                                  model_info: EmbedModelInfo,
+                                  example_prompts) -> None:
+    correctness_test_embed_models(hf_runner, vllm_runner, model_info,
+                                  example_prompts)

tests/models/language/pooling/test_jina.py

@@ -31,8 +31,8 @@
 EMBEDDING_MODELS = [
     EmbedModelInfo("jinaai/jina-embeddings-v3",
                    architecture="XLMRobertaModel",
-                   is_matryoshka=True,
-                   dtype="float32")
+                   dtype="hybrid",
+                   is_matryoshka=True)
 ]
 
 

tests/models/language/pooling/test_nomic.py

@@ -8,18 +8,15 @@
 MODELS = [
     EmbedModelInfo("nomic-ai/nomic-embed-text-v1",
                    architecture="NomicBertModel",
-                   dtype="float32",
                    enable_test=True),
     EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5",
                    architecture="NomicBertModel",
-                   dtype="float32",
                    enable_test=False),
     EmbedModelInfo("nomic-ai/CodeRankEmbed",
                    architecture="NomicBertModel",
                    enable_test=False),
     EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe",
                    architecture="NomicBertModel",
-                   dtype="float32",
                    enable_test=True)
 ]
 

tests/models/test_hybrid_dtype.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# ruff: noqa: SIM117
+import pytest
+import torch
+
+from tests.models.language.pooling.mteb_utils import mteb_test_embed_models
+from tests.models.utils import Dtype, EmbedModelInfo
+from vllm.config import _STR_DTYPE_TO_TORCH_DTYPE
+
+high_precision_data_types = [
+    Dtype(dtype="auto"),  # hybrid
+    Dtype(dtype="float32"),
+    Dtype(dtype="hybrid"),
+    Dtype(dtype="float32", attn_dtype="float16"),
+    Dtype(dtype="float32", attn_dtype="bfloat16")
+]
+low_precision_data_types = [Dtype(dtype="float16"), Dtype(dtype="bfloat16")]
+data_types = high_precision_data_types + low_precision_data_types
+embed_model = "intfloat/e5-small"
+generate_model = "EleutherAI/pythia-70m"
+
+
+@pytest.mark.parametrize("dtype", data_types)
+def test_dtype(vllm_runner, dtype: Dtype):
+    with vllm_runner(embed_model,
+                     dtype=dtype.dtype,
+                     max_model_len=None,
+                     attn_dtype=dtype.attn_dtype) as vllm_model:
+        model_config = vllm_model.model.llm_engine.model_config
+        if dtype.dtype == "hybrid" or dtype.dtype == "auto":
+            assert model_config.dtype == torch.float32
+            assert model_config.attn_dtype == torch.float16
+        elif dtype.attn_dtype == "auto":
+            assert model_config.dtype == model_config.attn_dtype
+        else:
+            assert model_config.dtype == _STR_DTYPE_TO_TORCH_DTYPE[dtype.dtype]
+            assert model_config.attn_dtype == _STR_DTYPE_TO_TORCH_DTYPE[
+                dtype.attn_dtype]
+
+
+@pytest.mark.parametrize("dtype", data_types)
+def test_embed_models_mteb(hf_runner, vllm_runner, dtype: Dtype):
+    model_info = EmbedModelInfo(embed_model,
+                                architecture="BertModel",
+                                dtype=dtype)
+
+    if model_info.dtype in high_precision_data_types:
+        mteb_test_embed_models(hf_runner, vllm_runner, model_info)
+    else:
+        with pytest.raises(AssertionError):
+            mteb_test_embed_models(hf_runner, vllm_runner, model_info)
+
+
+@pytest.mark.parametrize("model", [generate_model])
+@pytest.mark.parametrize("dtype", data_types)
+@pytest.mark.parametrize("max_tokens", [32])
+@pytest.mark.parametrize("num_logprobs", [4])
+def test_generate_models(hf_runner, vllm_runner, example_prompts, model: str,
+                         dtype: Dtype, max_tokens: int,
+                         num_logprobs: int) -> None:
+    if dtype.attn_dtype == "auto" and dtype.dtype != "hybrid":
+        with vllm_runner(model, dtype=dtype.dtype,
+                         attn_dtype=dtype.attn_dtype):
+            pass
+    else:
+        with pytest.raises(ValueError):
+            with vllm_runner(model,
+                             dtype=dtype.dtype,
+                             attn_dtype=dtype.attn_dtype):
+                pass

tests/models/utils.py

@@ -7,7 +7,7 @@
 import torch
 import torch.nn.functional as F
 
-from vllm.config import ModelConfig, TaskOption
+from vllm.config import AttnDType, ModelConfig, ModelDType, TaskOption
 from vllm.inputs import InputContext
 from vllm.sequence import Logprob, PromptLogprobs, SampleLogprobs
 
@@ -328,10 +328,15 @@ def matryoshka_fy(tensor: torch.Tensor, dimensions: int):
     return tensor
 
 
+class Dtype(NamedTuple):
+    dtype: ModelDType
+    attn_dtype: AttnDType = "auto"
+
+
 class EmbedModelInfo(NamedTuple):
     name: str
     is_matryoshka: bool = False
     matryoshka_dimensions: Optional[list[int]] = None
     architecture: str = ""
-    dtype: str = "auto"
+    dtype: Union[str, Dtype] = "auto"
     enable_test: bool = True

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py

@@ -6,6 +6,6 @@
 class DummyPlatform(CudaPlatform):
     device_name = "DummyDevice"
 
-    def get_attn_backend_cls(self, backend_name, head_size, dtype,
+    def get_attn_backend_cls(self, backend_name, head_size, attn_dtype,
                              kv_cache_dtype, block_size, use_v1, use_mla):
         return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend"  # noqa E501