A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
Improve the output precision of embedding models (#19092)
wang.yuqi committed
Jun 4, 2025 at 19:48 UTC
35cf32df304770b9dd3878438544b3a1a1cc79a5
8 files changed
+69
-28
tests/models/language/pooling/embed_utils.py
+1
-5
@@ -56,14 +56,10 @@ def correctness_test_embed_models(hf_runner,
56
max_model_len=None,
57
**vllm_extra_kwargs) as vllm_model:
58
vllm_outputs = vllm_model.encode(example_prompts)
59
- vllm_dtype = vllm_model.model.llm_engine.model_config.dtype
60
- model_dtype = getattr(
61
- vllm_model.model.llm_engine.model_config.hf_config, "torch_dtype",
62
- vllm_dtype)
59
60
with hf_runner(
61
model_info.name,
66
- dtype=model_dtype,
62
+ dtype="float32",
63
is_sentence_transformer=True,
64
) as hf_model:
65
tests/models/language/pooling/mteb_utils.py
+6
-6
@@ -7,7 +7,6 @@ import numpy as np
7
import pytest
8
9
from tests.models.utils import EmbedModelInfo
10
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
10
11
# Most models on the STS12 task (See #17175):
12
# - Model implementation and minor changes in tensor dtype
@@ -104,17 +103,18 @@ def mteb_test_embed_models(hf_runner,
103
MTEB_EMBED_TASKS)
104
vllm_dtype = vllm_model.model.llm_engine.model_config.dtype
105
107
- with set_default_torch_dtype(vllm_dtype) and hf_runner(
108
- model_info.name, is_sentence_transformer=True,
109
- dtype=vllm_dtype) as hf_model:
106
+ with hf_runner(model_info.name,
107
+ is_sentence_transformer=True,
108
+ dtype="float32") as hf_model:
109
110
if hf_model_callback is not None:
111
hf_model_callback(hf_model)
112
113
st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
114
+ st_dtype = next(hf_model.model.parameters()).dtype
115
116
- print("VLLM:", vllm_main_score)
117
- print("SentenceTransformers:", st_main_score)
116
+ print("VLLM:", vllm_dtype, vllm_main_score)
117
+ print("SentenceTransformers:", st_dtype, st_main_score)
118
print("Difference:", st_main_score - vllm_main_score)
119
120
assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_EMBED_TOL)
tests/models/language/pooling/test_gte.py
-7
@@ -11,27 +11,21 @@ MODELS = [
11
########## BertModel
12
EmbedModelInfo("thenlper/gte-large",
13
architecture="BertModel",
14
- dtype="float32",
14
enable_test=True),
15
EmbedModelInfo("thenlper/gte-base",
16
architecture="BertModel",
18
- dtype="float32",
17
enable_test=False),
18
EmbedModelInfo("thenlper/gte-small",
19
architecture="BertModel",
22
- dtype="float32",
20
enable_test=False),
21
EmbedModelInfo("thenlper/gte-large-zh",
22
architecture="BertModel",
26
- dtype="float32",
23
enable_test=False),
24
EmbedModelInfo("thenlper/gte-base-zh",
25
architecture="BertModel",
30
- dtype="float32",
26
enable_test=False),
27
EmbedModelInfo("thenlper/gte-small-zh",
28
architecture="BertModel",
34
- dtype="float32",
29
enable_test=False),
30
########### NewModel
31
EmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
@@ -46,7 +40,6 @@ MODELS = [
40
########### Qwen2ForCausalLM
41
EmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
42
architecture="Qwen2ForCausalLM",
49
- dtype="float32",
43
enable_test=True),
44
########## ModernBertModel
45
EmbedModelInfo("Alibaba-NLP/gte-modernbert-base",
tests/models/language/pooling/test_intfloat.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import pytest
3
+
4
+from ...utils import EmbedModelInfo
5
+from .embed_utils import correctness_test_embed_models
6
+from .mteb_utils import mteb_test_embed_models
7
+
8
+MODELS = [
9
+ ########## BertModel
10
+ EmbedModelInfo("intfloat/e5-small",
11
+ architecture="BertModel",
12
+ enable_test=True),
13
+ EmbedModelInfo("intfloat/e5-base",
14
+ architecture="BertModel",
15
+ enable_test=False),
16
+ EmbedModelInfo("intfloat/e5-large",
17
+ architecture="BertModel",
18
+ enable_test=False),
19
+ EmbedModelInfo("intfloat/multilingual-e5-small",
20
+ architecture="BertModel",
21
+ enable_test=False),
22
+ ########## XLMRobertaModel
23
+ EmbedModelInfo("intfloat/multilingual-e5-base",
24
+ architecture="XLMRobertaModel",
25
+ enable_test=True),
26
+ EmbedModelInfo("intfloat/multilingual-e5-large",
27
+ architecture="XLMRobertaModel",
28
+ enable_test=False),
29
+ EmbedModelInfo("intfloat/multilingual-e5-large-instruct",
30
+ architecture="XLMRobertaModel",
31
+ enable_test=False),
32
+]
33
+
34
+
35
+@pytest.mark.parametrize("model_info", MODELS)
36
+def test_embed_models_mteb(hf_runner, vllm_runner,
37
+ model_info: EmbedModelInfo) -> None:
38
+ mteb_test_embed_models(hf_runner, vllm_runner, model_info)
39
+
40
+
41
+@pytest.mark.parametrize("model_info", MODELS)
42
+def test_embed_models_correctness(hf_runner, vllm_runner,
43
+ model_info: EmbedModelInfo,
44
+ example_prompts) -> None:
45
+ correctness_test_embed_models(hf_runner, vllm_runner, model_info,
46
+ example_prompts)
tests/models/language/pooling/test_jina.py
+1
-2
@@ -32,8 +32,7 @@ TEXTS_2 = [
32
EMBEDDING_MODELS = [
33
EmbedModelInfo("jinaai/jina-embeddings-v3",
34
architecture="XLMRobertaModel",
35
- is_matryoshka=True,
36
- dtype="float32")
35
+ is_matryoshka=True)
36
]
37
38
tests/models/language/pooling/test_nomic.py
-3
@@ -9,18 +9,15 @@ from .mteb_utils import mteb_test_embed_models
9
MODELS = [
10
EmbedModelInfo("nomic-ai/nomic-embed-text-v1",
11
architecture="NomicBertModel",
12
- dtype="float32",
12
enable_test=True),
13
EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5",
14
architecture="NomicBertModel",
16
- dtype="float32",
15
enable_test=False),
16
EmbedModelInfo("nomic-ai/CodeRankEmbed",
17
architecture="NomicBertModel",
18
enable_test=False),
19
EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe",
20
architecture="NomicBertModel",
23
- dtype="float32",
21
enable_test=True)
22
]
23
vllm/model_executor/models/bert.py
+9
-4
@@ -414,10 +414,15 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
414
intermediate_tensors: Optional[IntermediateTensors] = None,
415
inputs_embeds: Optional[torch.Tensor] = None,
416
) -> torch.Tensor:
417
- return self.model(input_ids=input_ids,
418
- position_ids=positions,
419
- inputs_embeds=inputs_embeds,
420
- intermediate_tensors=intermediate_tensors)
417
+ hidden_states = self.model(input_ids=input_ids,
418
+ position_ids=positions,
419
+ inputs_embeds=inputs_embeds,
420
+ intermediate_tensors=intermediate_tensors)
421
+
422
+ # convert the embedding output to float32,
423
+ # otherwise precision will be lost significantly
424
+ hidden_states = hidden_states.to(torch.float32)
425
+ return hidden_states
426
427
def pooler(
428
self,
vllm/model_executor/models/bert_with_rope.py
+6
-1
@@ -432,7 +432,12 @@ class BertWithRope(nn.Module, SupportsV0Only, SupportsQuant):
432
else:
433
hidden_states = self.embeddings(input_ids=input_ids,
434
token_type_ids=token_type_ids)
435
- return self.encoder(positions, hidden_states)
435
+ hidden_states = self.encoder(positions, hidden_states)
436
+
437
+ # convert the embedding output to float32,
438
+ # otherwise precision will be lost significantly
439
+ hidden_states = hidden_states.to(torch.float32)
440
+ return hidden_states
441
442
def load_weights(self, weights: Iterable[tuple[str,
443
torch.Tensor]]) -> set[str]: