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, 08:30 UTC
[Bugfix] Fix the failing gte embedding test (#18720)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
May 29, 2025 at 22:39 UTC
c9479b292086b9ae0623f619af15cf54d58842c1
4 files changed
+20
-13
tests/conftest.py
+7
-6
@@ -311,6 +311,7 @@ class HfRunner:
311
dtype: str = "auto",
312
*,
313
model_kwargs: Optional[dict[str, Any]] = None,
314
+ trust_remote_code: bool = True,
315
is_sentence_transformer: bool = False,
316
is_cross_encoder: bool = False,
317
skip_tokenizer_init: bool = False,
@@ -320,7 +321,7 @@ class HfRunner:
321
322
self.config = AutoConfig.from_pretrained(
323
model_name,
323
- trust_remote_code=True,
324
+ trust_remote_code=trust_remote_code,
325
)
326
self.device = self.get_default_device()
327
self.dtype = torch_dtype = _get_and_verify_dtype(self.config, dtype)
@@ -336,7 +337,7 @@ class HfRunner:
337
model_name,
338
device=self.device,
339
model_kwargs=model_kwargs,
339
- trust_remote_code=True,
340
+ trust_remote_code=trust_remote_code,
341
)
342
elif is_cross_encoder:
343
# Lazy init required for AMD CI
@@ -346,12 +347,12 @@ class HfRunner:
347
model_name,
348
device=self.device,
349
automodel_args=model_kwargs,
349
- trust_remote_code=True,
350
+ trust_remote_code=trust_remote_code,
351
)
352
else:
353
model = auto_cls.from_pretrained(
354
model_name,
354
- trust_remote_code=True,
355
+ trust_remote_code=trust_remote_code,
356
**model_kwargs,
357
)
358
@@ -372,7 +373,7 @@ class HfRunner:
373
self.tokenizer = AutoTokenizer.from_pretrained(
374
model_name,
375
torch_dtype=torch_dtype,
375
- trust_remote_code=True,
376
+ trust_remote_code=trust_remote_code,
377
)
378
379
# don't put this import at the top level
@@ -381,7 +382,7 @@ class HfRunner:
382
self.processor = AutoProcessor.from_pretrained(
383
model_name,
384
torch_dtype=torch_dtype,
384
- trust_remote_code=True,
385
+ trust_remote_code=trust_remote_code,
386
)
387
if skip_tokenizer_init:
388
self.tokenizer = self.processor.tokenizer
tests/models/language/pooling/test_embedding.py
+11
-7
@@ -10,18 +10,22 @@ from ...utils import check_embeddings_close
10
@pytest.mark.parametrize(
11
"model",
12
[
13
- # [Encoder-only]
14
- pytest.param("BAAI/bge-base-en-v1.5",
15
- marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
16
- pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
17
- pytest.param("intfloat/multilingual-e5-small"),
18
- pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct"),
13
+ # Be careful of the order of models, decoder-only models should be
14
+ # placed before encoder-only models, otherwise `Qwen2.5-0.5B-Instruct`
15
+ # case won't pass because gte-Qwen2-1.5B-instruct will cache custom
16
+ # model code with bidirectional attention.
17
# [Decoder-only]
18
pytest.param("BAAI/bge-multilingual-gemma2",
19
marks=[pytest.mark.core_model]),
20
pytest.param("intfloat/e5-mistral-7b-instruct",
21
marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
22
pytest.param("ssmits/Qwen2-7B-Instruct-embed-base"),
23
+ # [Encoder-only]
24
+ pytest.param("BAAI/bge-base-en-v1.5",
25
+ marks=[pytest.mark.core_model, pytest.mark.cpu_model]),
26
+ pytest.param("sentence-transformers/all-MiniLM-L12-v2"),
27
+ pytest.param("intfloat/multilingual-e5-small"),
28
+ pytest.param("Alibaba-NLP/gte-Qwen2-1.5B-instruct"),
29
# [Cross-Encoder]
30
pytest.param("sentence-transformers/stsb-roberta-base-v2"),
31
],
@@ -44,7 +48,7 @@ def test_models(
48
vllm_extra_kwargs = {}
49
if model == "ssmits/Qwen2-7B-Instruct-embed-base":
50
vllm_extra_kwargs["override_pooler_config"] = \
47
- PoolerConfig(pooling_type="MEAN")
51
+ PoolerConfig(pooling_type="MEAN", normalize=False)
52
53
# The example_prompts has ending "\n", for example:
54
# "Write a short story about a robot that dreams for the first time.\n"
tests/models/language/pooling/test_gte.py
+1
@@ -45,6 +45,7 @@ MODELS = [
45
########### Qwen2ForCausalLM
46
EmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
47
architecture="Qwen2ForCausalLM",
48
+ dtype="float32",
49
enable_test=True),
50
########## ModernBertModel
51
EmbedModelInfo("Alibaba-NLP/gte-modernbert-base",
tests/models/utils.py
+1
@@ -314,6 +314,7 @@ def check_embeddings_close(
314
dim=0)
315
316
fail_msg = (f"Test{prompt_idx}:"
317
+ f"\nCosine similarity: \t{sim:.4f}"
318
f"\n{name_0}:\t{embeddings_0[:16]!r}"
319
f"\n{name_1}:\t{embeddings_1[:16]!r}")
320