@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f40f763f12
Fork 0

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, 06:00 UTC

[CI] Add mteb testing for rerank models (#19344)

wang.yuqi committed Jun 16, 2025 at 16:36 UTC f40f763f12c0ffbbc781b0ea23f4b9868db70d21
15 files changed +428 -256
requirements/test.in
+1 -1
@@ -33,7 +33,7 @@ num2words # required for smolvlm test
33 opencv-python-headless >= 4.11.0 # required for video test
34 datamodel_code_generator # required for minicpm3 test
35 lm-eval[api]==0.4.8 # required for model evaluation test
36 -mteb>=1.38.11, <2 # required for mteb test
36 +mteb[bm25s]>=1.38.11, <2 # required for mteb test
37 transformers==4.52.4
38 tokenizers==0.21.1
39 huggingface-hub[hf_xet]>=0.30.0 # Required for Xet downloads.
requirements/test.txt
+6
@@ -51,6 +51,8 @@ black==24.10.0
51 # via datamodel-code-generator
52 blobfile==3.0.0
53 # via -r requirements/test.in
54 +bm25s==0.2.13
55 + # via mteb
56 boto3==1.35.57
57 # via tensorizer
58 botocore==1.35.57
@@ -344,6 +346,7 @@ numpy==1.26.4
346 # -r requirements/test.in
347 # accelerate
348 # bitsandbytes
349 + # bm25s
350 # contourpy
351 # cupy-cuda12x
352 # datasets
@@ -534,6 +537,8 @@ pyparsing==3.2.0
537 # via matplotlib
538 pyrate-limiter==3.7.0
539 # via schemathesis
540 +pystemmer==3.0.0
541 + # via mteb
542 pytablewriter==1.2.0
543 # via lm-eval
544 pytest==8.3.3
@@ -668,6 +673,7 @@ scikit-learn==1.5.2
673 # sentence-transformers
674 scipy==1.13.1
675 # via
676 + # bm25s
677 # librosa
678 # mteb
679 # scikit-learn
tests/conftest.py
+9 -3
@@ -727,8 +727,12 @@ class HfRunner:
727 **kwargs) -> list[list[torch.Tensor]]:
728 return self.model.encode(prompts, *args, **kwargs)
729
730 - def predict(self, prompts: list[list[str]]) -> torch.Tensor:
731 - return self.model.predict(prompts, convert_to_tensor=True)
730 + def predict(self, prompts: list[list[str]], *args,
731 + **kwargs) -> torch.Tensor:
732 + return self.model.predict(prompts,
733 + *args,
734 + convert_to_tensor=True,
735 + **kwargs)
736
737 def __enter__(self):
738 return self
@@ -1037,8 +1041,10 @@ class VllmRunner:
1041 self,
1042 text_1: Union[str, list[str]],
1043 text_2: Union[str, list[str]],
1044 + *args,
1045 + **kwargs,
1046 ) -> list[float]:
1041 - req_outputs = self.model.score(text_1, text_2)
1047 + req_outputs = self.model.score(text_1, text_2, *args, **kwargs)
1048 return [req_output.outputs.score for req_output in req_outputs]
1049
1050 def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
tests/entrypoints/openai/correctness/test_mteb_embed.py renamed
+6 -10
@@ -7,34 +7,30 @@ import pytest
7 from tests.models.language.pooling.mteb_utils import (MTEB_EMBED_TASKS,
8 MTEB_EMBED_TOL,
9 OpenAIClientMtebEncoder,
10 - run_mteb_embed_task,
11 - run_mteb_embed_task_st)
10 + run_mteb_embed_task)
11 from tests.utils import RemoteOpenAIServer
12
13 os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
14
16 -MODEL_NAME = "BAAI/bge-m3"
17 -DTYPE = "float16"
18 -MAIN_SCORE = 0.7873427091972599
15 +MODEL_NAME = "intfloat/e5-small"
16 +MAIN_SCORE = 0.7422994752439667
17
18
19 @pytest.fixture(scope="module")
20 def server():
21 args = [
24 - "--task", "embed", "--dtype", DTYPE, "--enforce-eager",
25 - "--max-model-len", "512"
22 + "--task", "embed", "--enforce-eager", "--disable-uvicorn-access-log"
23 ]
24
25 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
26 yield remote_server
27
28
32 -def test_mteb(server):
29 +def test_mteb_embed(server):
30 client = server.get_client()
31 encoder = OpenAIClientMtebEncoder(MODEL_NAME, client)
32 vllm_main_score = run_mteb_embed_task(encoder, MTEB_EMBED_TASKS)
36 - st_main_score = MAIN_SCORE or run_mteb_embed_task_st(
37 - MODEL_NAME, MTEB_EMBED_TASKS)
33 + st_main_score = MAIN_SCORE
34
35 print("VLLM main score: ", vllm_main_score)
36 print("SentenceTransformer main score: ", st_main_score)
tests/entrypoints/openai/correctness/test_mteb_score.py new
+59
@@ -0,0 +1,59 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import os
4 +
5 +import pytest
6 +
7 +# yapf conflicts with isort for this block
8 +# yapf: disable
9 +from tests.models.language.pooling.mteb_utils import (MTEB_RERANK_LANGS,
10 + MTEB_RERANK_TASKS,
11 + MTEB_RERANK_TOL,
12 + RerankClientMtebEncoder,
13 + ScoreClientMtebEncoder,
14 + run_mteb_rerank)
15 +# yapf: enable
16 +from tests.utils import RemoteOpenAIServer
17 +
18 +os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
19 +
20 +MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2"
21 +MAIN_SCORE = 0.33437
22 +
23 +
24 +@pytest.fixture(scope="module")
25 +def server():
26 + args = [
27 + "--task", "score", "--enforce-eager", "--disable-uvicorn-access-log"
28 + ]
29 +
30 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
31 + yield remote_server
32 +
33 +
34 +def test_mteb_score(server):
35 + url = server.url_for("score")
36 + encoder = ScoreClientMtebEncoder(MODEL_NAME, url)
37 + vllm_main_score = run_mteb_rerank(encoder, MTEB_RERANK_TASKS,
38 + MTEB_RERANK_LANGS)
39 + st_main_score = MAIN_SCORE
40 +
41 + print("VLLM main score: ", vllm_main_score)
42 + print("SentenceTransformer main score: ", st_main_score)
43 + print("Difference: ", st_main_score - vllm_main_score)
44 +
45 + assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_RERANK_TOL)
46 +
47 +
48 +def test_mteb_rerank(server):
49 + url = server.url_for("rerank")
50 + encoder = RerankClientMtebEncoder(MODEL_NAME, url)
51 + vllm_main_score = run_mteb_rerank(encoder, MTEB_RERANK_TASKS,
52 + MTEB_RERANK_LANGS)
53 + st_main_score = MAIN_SCORE
54 +
55 + print("VLLM main score: ", vllm_main_score)
56 + print("SentenceTransformer main score: ", st_main_score)
57 + print("Difference: ", st_main_score - vllm_main_score)
58 +
59 + assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_RERANK_TOL)
tests/models/language/pooling/mteb_utils.py
+183 -9
@@ -1,14 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import tempfile
5 from collections.abc import Sequence
6 +from typing import Optional
7
8 import mteb
9 import numpy as np
10 import pytest
11 +import requests
12
9 -from tests.models.utils import EmbedModelInfo
13 +from tests.models.utils import EmbedModelInfo, RerankModelInfo
14
11 -# Most models on the STS12 task (See #17175):
15 +# Most embedding models on the STS12 task (See #17175):
16 # - Model implementation and minor changes in tensor dtype
17 # results in differences less than 1e-4
18 # - Different model results in differences more than 1e-3
@@ -16,6 +20,11 @@ from tests.models.utils import EmbedModelInfo
20 MTEB_EMBED_TASKS = ["STS12"]
21 MTEB_EMBED_TOL = 1e-4
22
23 +# See #19344
24 +MTEB_RERANK_TASKS = ["NFCorpus"]
25 +MTEB_RERANK_LANGS = ["en"]
26 +MTEB_RERANK_TOL = 1e-3
27 +
28
29 class VllmMtebEncoder(mteb.Encoder):
30
@@ -39,6 +48,27 @@ class VllmMtebEncoder(mteb.Encoder):
48 embeds = embeds[np.argsort(r)]
49 return embeds
50
51 + def predict(
52 + self,
53 + sentences: list[tuple[str, str,
54 + Optional[str]]], # query, corpus, prompt
55 + *args,
56 + **kwargs,
57 + ) -> np.ndarray:
58 + r = self.rng.permutation(len(sentences))
59 + sentences = [sentences[i] for i in r]
60 +
61 + queries = [s[0] for s in sentences]
62 + corpus = [s[1] for s in sentences]
63 +
64 + outputs = self.model.score(queries,
65 + corpus,
66 + truncate_prompt_tokens=-1,
67 + use_tqdm=False)
68 + scores = np.array(outputs)
69 + scores = scores[np.argsort(r)]
70 + return scores
71 +
72
73 class OpenAIClientMtebEncoder(mteb.Encoder):
74
@@ -62,21 +92,72 @@ class OpenAIClientMtebEncoder(mteb.Encoder):
92 return embeds
93
94
95 +class ScoreClientMtebEncoder(mteb.Encoder):
96 +
97 + def __init__(self, model_name: str, url):
98 + super().__init__()
99 + self.model_name = model_name
100 + self.url = url
101 + self.rng = np.random.default_rng(seed=42)
102 +
103 + def predict(
104 + self,
105 + sentences: list[tuple[str, str,
106 + Optional[str]]], # query, corpus, prompt
107 + *args,
108 + **kwargs,
109 + ) -> np.ndarray:
110 + r = self.rng.permutation(len(sentences))
111 + sentences = [sentences[i] for i in r]
112 +
113 + outputs = []
114 + for query, corpus, prompt in sentences:
115 + outputs.append(self.get_score(query, corpus))
116 +
117 + scores = np.array(outputs)
118 + scores = scores[np.argsort(r)]
119 + return scores
120 +
121 + def get_score(self, query, corpus):
122 + response = requests.post(self.url,
123 + json={
124 + "model": self.model_name,
125 + "text_1": query,
126 + "text_2": corpus,
127 + "truncate_prompt_tokens": -1,
128 + }).json()
129 + return response['data'][0]["score"]
130 +
131 +
132 +class RerankClientMtebEncoder(ScoreClientMtebEncoder):
133 +
134 + def get_score(self, query, corpus):
135 + response = requests.post(self.url,
136 + json={
137 + "model": self.model_name,
138 + "query": query,
139 + "documents": [corpus],
140 + "truncate_prompt_tokens": -1,
141 + }).json()
142 + return response['results'][0]["relevance_score"]
143 +
144 +
145 def run_mteb_embed_task(encoder, tasks):
146 tasks = mteb.get_tasks(tasks=tasks)
147 evaluation = mteb.MTEB(tasks=tasks)
68 - results = evaluation.run(encoder, verbosity=0, output_folder=None)
148 + results = evaluation.run(
149 + encoder,
150 + verbosity=0,
151 + output_folder=None,
152 + encode_kwargs={
153 + "show_progress_bar": False,
154 + },
155 + )
156
157 main_score = results[0].scores["test"][0]["main_score"]
158 return main_score
159
160
74 -def run_mteb_embed_task_st(model_name, tasks):
75 - from sentence_transformers import SentenceTransformer
76 - model = SentenceTransformer(model_name)
77 - return run_mteb_embed_task(model, tasks)
78 -
79 -
161 def mteb_test_embed_models(hf_runner,
162 vllm_runner,
163 model_info: EmbedModelInfo,
@@ -118,3 +199,96 @@ def mteb_test_embed_models(hf_runner,
199 print("Difference:", st_main_score - vllm_main_score)
200
201 assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_EMBED_TOL)
202 +
203 +
204 +def run_mteb_rerank(cross_encoder, tasks, languages):
205 + with tempfile.TemporaryDirectory() as results_folder:
206 + bm25s = mteb.get_model("bm25s")
207 + tasks = mteb.get_tasks(tasks=tasks, languages=languages)
208 +
209 + subset = "default"
210 + eval_splits = ["test"]
211 +
212 + evaluation = mteb.MTEB(tasks=tasks)
213 + evaluation.run(
214 + bm25s,
215 + verbosity=0,
216 + eval_splits=eval_splits,
217 + save_predictions=True,
218 + output_folder=f"{results_folder}/stage1",
219 + encode_kwargs={"show_progress_bar": False},
220 + )
221 +
222 + results = evaluation.run(
223 + cross_encoder,
224 + verbosity=0,
225 + eval_splits=eval_splits,
226 + top_k=10,
227 + save_predictions=True,
228 + output_folder=f"{results_folder}/stage2",
229 + previous_results=
230 + f"{results_folder}/stage1/NFCorpus_{subset}_predictions.json",
231 + encode_kwargs={"show_progress_bar": False},
232 + )
233 + main_score = results[0].scores["test"][0]["main_score"]
234 + return main_score
235 +
236 +
237 +def mteb_test_rerank_models(hf_runner,
238 + vllm_runner,
239 + model_info: RerankModelInfo,
240 + vllm_extra_kwargs=None,
241 + hf_model_callback=None):
242 + if not model_info.enable_test:
243 + # A model family has many models with the same architecture,
244 + # and we don't need to test each one.
245 + pytest.skip("Skipping test.")
246 +
247 + vllm_extra_kwargs = vllm_extra_kwargs or {}
248 + vllm_extra_kwargs["dtype"] = model_info.dtype
249 +
250 + with vllm_runner(model_info.name,
251 + task="score",
252 + max_model_len=None,
253 + **vllm_extra_kwargs) as vllm_model:
254 +
255 + if model_info.architecture:
256 + assert (model_info.architecture
257 + in vllm_model.model.llm_engine.model_config.architectures)
258 +
259 + vllm_main_score = run_mteb_rerank(VllmMtebEncoder(vllm_model),
260 + tasks=MTEB_RERANK_TASKS,
261 + languages=MTEB_RERANK_LANGS)
262 + vllm_dtype = vllm_model.model.llm_engine.model_config.dtype
263 +
264 + with hf_runner(model_info.name, is_cross_encoder=True,
265 + dtype="float32") as hf_model:
266 +
267 + original_predict = hf_model.predict
268 +
269 + def _predict(
270 + sentences: list[tuple[str, str,
271 + Optional[str]]], # query, corpus, prompt
272 + *args,
273 + **kwargs,
274 + ):
275 + # vllm and st both remove the prompt, fair comparison.
276 + prompts = [(s[0], s[1]) for s in sentences]
277 + return original_predict(prompts, *args, **kwargs, batch_size=8)
278 +
279 + hf_model.predict = _predict
280 + hf_model.original_predict = original_predict
281 +
282 + if hf_model_callback is not None:
283 + hf_model_callback(hf_model)
284 +
285 + st_main_score = run_mteb_rerank(hf_model,
286 + tasks=MTEB_RERANK_TASKS,
287 + languages=MTEB_RERANK_LANGS)
288 + st_dtype = next(hf_model.model.model.parameters()).dtype
289 +
290 + print("VLLM:", vllm_dtype, vllm_main_score)
291 + print("SentenceTransformers:", st_dtype, st_main_score)
292 + print("Difference:", st_main_score - vllm_main_score)
293 +
294 + assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_RERANK_TOL)
tests/models/language/pooling/test_baai.py
+23 -2
@@ -2,8 +2,9 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4
5 -from .embed_utils import EmbedModelInfo, correctness_test_embed_models
6 -from .mteb_utils import mteb_test_embed_models
5 +from ...utils import EmbedModelInfo, RerankModelInfo
6 +from .embed_utils import correctness_test_embed_models
7 +from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
8
9 MODELS = [
10 ########## BertModel
@@ -57,6 +58,20 @@ MODELS = [
58 enable_test=True),
59 ]
60
61 +RERANK_MODELS = [
62 + ########## XLMRobertaForSequenceClassification
63 + RerankModelInfo("BAAI/bge-reranker-base",
64 + architecture="XLMRobertaForSequenceClassification",
65 + enable_test=True),
66 + RerankModelInfo("BAAI/bge-reranker-large",
67 + architecture="XLMRobertaForSequenceClassification",
68 + enable_test=False),
69 + RerankModelInfo("BAAI/bge-reranker-v2-m3",
70 + architecture="XLMRobertaForSequenceClassification",
71 + dtype="float32",
72 + enable_test=False)
73 +]
74 +
75
76 @pytest.mark.parametrize("model_info", MODELS)
77 def test_embed_models_mteb(hf_runner, vllm_runner,
@@ -70,3 +85,9 @@ def test_embed_models_correctness(hf_runner, vllm_runner,
85 example_prompts) -> None:
86 correctness_test_embed_models(hf_runner, vllm_runner, model_info,
87 example_prompts)
88 +
89 +
90 +@pytest.mark.parametrize("model_info", RERANK_MODELS)
91 +def test_rerank_models_mteb(hf_runner, vllm_runner,
92 + model_info: RerankModelInfo) -> None:
93 + mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
tests/models/language/pooling/test_cross_encoder.py new
+18
@@ -0,0 +1,18 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +
5 +from .mteb_utils import RerankModelInfo, mteb_test_rerank_models
6 +
7 +RERANK_MODELS = [
8 + RerankModelInfo("cross-encoder/ms-marco-TinyBERT-L-2-v2",
9 + architecture="BertForSequenceClassification"),
10 + RerankModelInfo("tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
11 + architecture="Qwen3ForSequenceClassification")
12 +]
13 +
14 +
15 +@pytest.mark.parametrize("model_info", RERANK_MODELS)
16 +def test_rerank_models_mteb(hf_runner, vllm_runner,
17 + model_info: RerankModelInfo) -> None:
18 + mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
tests/models/language/pooling/test_jina.py
+16 -62
@@ -6,28 +6,10 @@ import pytest
6
7 from vllm import PoolingParams
8
9 -from .embed_utils import (EmbedModelInfo, check_embeddings_close,
9 +from ...utils import EmbedModelInfo, RerankModelInfo
10 +from .embed_utils import (check_embeddings_close,
11 correctness_test_embed_models, matryoshka_fy)
11 -from .mteb_utils import mteb_test_embed_models
12 -
13 -SCORING_MODELS = [
14 - "jinaai/jina-reranker-v2-base-multilingual", # Roberta
15 -]
16 -
17 -TEXTS_1 = ["Organic skincare products for sensitive skin"]
18 -
19 -TEXTS_2 = [
20 - "Organic skincare for sensitive skin with aloe vera and chamomile.",
21 - "New makeup trends focus on bold colors and innovative techniques",
22 - "Bio-Hautpflege für empfindliche Haut mit Aloe Vera und Kamille",
23 - "Neue Make-up-Trends setzen auf kräftige Farben und innovative Techniken", # noqa: E501
24 - "Cuidado de la piel orgánico para piel sensible con aloe vera y manzanilla", # noqa: E501
25 - "Las nuevas tendencias de maquillaje se centran en colores vivos y técnicas innovadoras", # noqa: E501
26 - "针对敏感肌专门设计的天然有机护肤产品",
27 - "新的化妆趋势注重鲜艳的颜色和创新的技巧",
28 - "敏感肌のために特別に設計された天然有機スキンケア製品",
29 - "新しいメイクのトレンドは鮮やかな色と革新的な技術に焦点を当てています",
30 -]
12 +from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
13
14 EMBEDDING_MODELS = [
15 EmbedModelInfo("jinaai/jina-embeddings-v3",
@@ -35,47 +17,13 @@ EMBEDDING_MODELS = [
17 is_matryoshka=True)
18 ]
19
38 -
39 -@pytest.fixture(scope="module", params=SCORING_MODELS)
40 -def model_name(request):
41 - yield request.param
42 -
43 -
44 -@pytest.mark.parametrize("dtype", ["half"])
45 -def test_llm_1_to_1(vllm_runner, hf_runner, model_name, dtype: str):
46 -
47 - text_pair = [TEXTS_1[0], TEXTS_2[0]]
48 -
49 - with hf_runner(model_name, dtype=dtype, is_cross_encoder=True) as hf_model:
50 - hf_outputs = hf_model.predict([text_pair]).tolist()
51 -
52 - with vllm_runner(model_name, task="score", dtype=dtype,
53 - max_model_len=None) as vllm_model:
54 - vllm_outputs = vllm_model.score(text_pair[0], text_pair[1])
55 -
56 - assert len(vllm_outputs) == 1
57 - assert len(hf_outputs) == 1
58 -
59 - assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.01)
60 -
61 -
62 -@pytest.mark.parametrize("dtype", ["half"])
63 -def test_llm_1_to_N(vllm_runner, hf_runner, model_name, dtype: str):
64 -
65 - text_pairs = [[TEXTS_1[0], text] for text in TEXTS_2]
66 -
67 - with hf_runner(model_name, dtype=dtype, is_cross_encoder=True) as hf_model:
68 - hf_outputs = hf_model.predict(text_pairs).tolist()
69 -
70 - with vllm_runner(model_name, task="score", dtype=dtype,
71 - max_model_len=None) as vllm_model:
72 - vllm_outputs = vllm_model.score(TEXTS_1[0], TEXTS_2)
73 -
74 - assert len(vllm_outputs) == 10
75 - assert len(hf_outputs) == 10
76 -
77 - assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.01)
78 - assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.01)
20 +RERANK_MODELS = [
21 + RerankModelInfo(
22 + "jinaai/jina-reranker-v2-base-multilingual",
23 + architecture="XLMRobertaForSequenceClassification",
24 + dtype="float32",
25 + )
26 +]
27
28
29 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
@@ -106,6 +54,12 @@ def test_embed_models_correctness(hf_runner, vllm_runner,
54 hf_model_callback=hf_model_callback)
55
56
57 +@pytest.mark.parametrize("model_info", RERANK_MODELS)
58 +def test_rerank_models_mteb(hf_runner, vllm_runner,
59 + model_info: RerankModelInfo) -> None:
60 + mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
61 +
62 +
63 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
64 @pytest.mark.parametrize("dtype", ["half"])
65 @pytest.mark.parametrize("dimensions", [16, 32])
tests/models/language/pooling/test_qwen3_reranker.py
+84 -80
@@ -1,87 +1,91 @@
1 # SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 import pytest
6 +import torch
7 +
8 +from tests.conftest import HfRunner
9
4 -model_name = "Qwen/Qwen3-Reranker-4B"
10 +from .mteb_utils import RerankModelInfo, mteb_test_rerank_models
11
6 -text_1 = "What is the capital of France?"
7 -texts_2 = [
8 - "The capital of Brazil is Brasilia.",
9 - "The capital of France is Paris.",
12 +RERANK_MODELS = [
13 + RerankModelInfo("Qwen/Qwen3-Reranker-0.6B",
14 + architecture="Qwen3ForSequenceClassification",
15 + dtype="float32",
16 + enable_test=True),
17 + RerankModelInfo("Qwen/Qwen3-Reranker-4B",
18 + architecture="Qwen3ForSequenceClassification",
19 + dtype="float32",
20 + enable_test=False)
21 ]
22
23
13 -def vllm_reranker(model_name):
14 - from vllm import LLM
15 -
16 - model = LLM(model=model_name,
17 - task="score",
18 - hf_overrides={
19 - "architectures": ["Qwen3ForSequenceClassification"],
20 - "classifier_from_token": ["no", "yes"],
21 - "is_original_qwen3_reranker": True,
22 - },
23 - dtype="float32")
24 -
25 - text_1 = "What is the capital of France?"
26 - texts_2 = [
27 - "The capital of Brazil is Brasilia.",
28 - "The capital of France is Paris.",
29 - ]
30 -
31 - outputs = model.score(text_1, texts_2)
32 -
33 - return [output.outputs.score for output in outputs]
34 -
35 -
36 -def hf_reranker(model_name):
37 - import torch
38 - from transformers import AutoModelForCausalLM, AutoTokenizer
39 -
40 - tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
41 - model = AutoModelForCausalLM.from_pretrained(model_name).eval()
42 -
43 - token_false_id = tokenizer.convert_tokens_to_ids("no")
44 - token_true_id = tokenizer.convert_tokens_to_ids("yes")
45 -
46 - max_length = 8192
47 -
48 - def process_inputs(pairs):
49 - inputs = tokenizer(pairs,
50 - padding=False,
51 - truncation='longest_first',
52 - return_attention_mask=False,
53 - max_length=max_length)
54 - for i, ele in enumerate(inputs['input_ids']):
55 - inputs['input_ids'][i] = ele
56 - inputs = tokenizer.pad(inputs,
57 - padding=True,
58 - return_tensors="pt",
59 - max_length=max_length)
60 - for key in inputs:
61 - inputs[key] = inputs[key].to(model.device)
62 - return inputs
63 -
64 - @torch.no_grad()
65 - def compute_logits(inputs, **kwargs):
66 - batch_scores = model(**inputs).logits[:, -1, :]
67 - true_vector = batch_scores[:, token_true_id]
68 - false_vector = batch_scores[:, token_false_id]
69 - batch_scores = torch.stack([false_vector, true_vector], dim=1)
70 - batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
71 - scores = batch_scores[:, 1].exp().tolist()
72 - return scores
73 -
74 - pairs = [(text_1, texts_2[0]), (text_1, texts_2[1])]
75 - inputs = process_inputs(pairs)
76 - scores = compute_logits(inputs)
77 -
78 - return scores
79 -
80 -
81 -@pytest.mark.parametrize("model_name", [model_name])
82 -def test_model(model_name):
83 - hf_outputs = hf_reranker(model_name)
84 - vllm_outputs = vllm_reranker(model_name)
85 -
86 - assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.01)
87 - assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.01)
24 +class Qwen3RerankerHfRunner(HfRunner):
25 +
26 + def __init__(self,
27 + model_name: str,
28 + dtype: str = "auto",
29 + *args: Any,
30 + **kwargs: Any) -> None:
31 + from transformers import AutoModelForCausalLM, AutoTokenizer
32 + super().__init__(model_name, dtype, auto_cls=AutoModelForCausalLM)
33 +
34 + self.tokenizer = AutoTokenizer.from_pretrained(model_name,
35 + padding_side='left')
36 + self.token_false_id = self.tokenizer.convert_tokens_to_ids("no")
37 + self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
38 +
39 + def predict(self, prompts: list[list[str]], *args,
40 + **kwargs) -> torch.Tensor:
41 +
42 + def process_inputs(pairs):
43 + inputs = self.tokenizer(pairs,
44 + padding=False,
45 + truncation='longest_first',
46 + return_attention_mask=False)
47 + for i, ele in enumerate(inputs['input_ids']):
48 + inputs['input_ids'][i] = ele
49 + inputs = self.tokenizer.pad(inputs,
50 + padding=True,
51 + return_tensors="pt")
52 + for key in inputs:
53 + inputs[key] = inputs[key].to(self.model.device)
54 + return inputs
55 +
56 + @torch.no_grad()
57 + def compute_logits(inputs):
58 + batch_scores = self.model(**inputs).logits[:, -1, :]
59 + true_vector = batch_scores[:, self.token_true_id]
60 + false_vector = batch_scores[:, self.token_false_id]
61 + batch_scores = torch.stack([false_vector, true_vector], dim=1)
62 + batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
63 + scores = batch_scores[:, 1].exp()
64 + return scores
65 +
66 + scores = []
67 + for prompt in prompts:
68 + inputs = process_inputs([prompt])
69 + score = compute_logits(inputs)
70 + scores.append(score[0].item())
71 + return torch.Tensor(scores)
72 +
73 +
74 +@pytest.mark.parametrize("model_info", RERANK_MODELS)
75 +def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
76 +
77 + assert model_info.architecture == "Qwen3ForSequenceClassification"
78 +
79 + vllm_extra_kwargs: dict[str, Any] = {
80 + "hf_overrides": {
81 + "architectures": ["Qwen3ForSequenceClassification"],
82 + "classifier_from_token": ["no", "yes"],
83 + "is_original_qwen3_reranker": True,
84 + }
85 + }
86 +
87 + if model_info.name == "Qwen/Qwen3-Reranker-4B":
88 + vllm_extra_kwargs["max_num_seqs"] = 1
89 +
90 + mteb_test_rerank_models(Qwen3RerankerHfRunner, vllm_runner, model_info,
91 + vllm_extra_kwargs)
tests/models/language/pooling/test_qwen3_reranker_seq_cls.py deleted
-73
@@ -1,73 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -import pytest
3 -
4 -model_name = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
5 -
6 -text_1 = "What is the capital of France?"
7 -texts_2 = [
8 - "The capital of Brazil is Brasilia.",
9 - "The capital of France is Paris.",
10 -]
11 -
12 -
13 -def vllm_reranker(model_name):
14 - from vllm import LLM
15 -
16 - model = LLM(model=model_name, task="score")
17 - outputs = model.score(text_1, texts_2)
18 -
19 - return [output.outputs.score for output in outputs]
20 -
21 -
22 -def hf_reranker(model_name):
23 - import torch
24 - from transformers import AutoModelForCausalLM, AutoTokenizer
25 -
26 - tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
27 - model = AutoModelForCausalLM.from_pretrained(model_name).eval()
28 -
29 - token_false_id = tokenizer.convert_tokens_to_ids("no")
30 - token_true_id = tokenizer.convert_tokens_to_ids("yes")
31 -
32 - max_length = 8192
33 -
34 - def process_inputs(pairs):
35 - inputs = tokenizer(pairs,
36 - padding=False,
37 - truncation='longest_first',
38 - return_attention_mask=False,
39 - max_length=max_length)
40 - for i, ele in enumerate(inputs['input_ids']):
41 - inputs['input_ids'][i] = ele
42 - inputs = tokenizer.pad(inputs,
43 - padding=True,
44 - return_tensors="pt",
45 - max_length=max_length)
46 - for key in inputs:
47 - inputs[key] = inputs[key].to(model.device)
48 - return inputs
49 -
50 - @torch.no_grad()
51 - def compute_logits(inputs, **kwargs):
52 - batch_scores = model(**inputs).logits[:, -1, :]
53 - true_vector = batch_scores[:, token_true_id]
54 - false_vector = batch_scores[:, token_false_id]
55 - batch_scores = torch.stack([false_vector, true_vector], dim=1)
56 - batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
57 - scores = batch_scores[:, 1].exp().tolist()
58 - return scores
59 -
60 - pairs = [(text_1, texts_2[0]), (text_1, texts_2[1])]
61 - inputs = process_inputs(pairs)
62 - scores = compute_logits(inputs)
63 -
64 - return scores
65 -
66 -
67 -@pytest.mark.parametrize("model_name", [model_name])
68 -def test_model(model_name):
69 - hf_outputs = hf_reranker(model_name)
70 - vllm_outputs = vllm_reranker(model_name)
71 -
72 - assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.01)
73 - assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.01)
tests/models/utils.py
+7
@@ -336,3 +336,10 @@ class EmbedModelInfo(NamedTuple):
336 architecture: str = ""
337 dtype: str = "auto"
338 enable_test: bool = True
339 +
340 +
341 +class RerankModelInfo(NamedTuple):
342 + name: str
343 + architecture: str = ""
344 + dtype: str = "auto"
345 + enable_test: bool = True
vllm/model_executor/layers/pooler.py
+11 -1
@@ -156,7 +156,10 @@ class MeanPool(SimplePooler):
156 ) -> Union[list[torch.Tensor], torch.Tensor]:
157 prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
158
159 - cumsum = torch.cumsum(hidden_states, dim=0)
159 + # Use float32 for torch.cumsum in MeanPool,
160 + # otherwise precision will be lost significantly.
161 + cumsum = torch.cumsum(hidden_states, dim=0, dtype=torch.float32)
162 +
163 start_indices = torch.cat([
164 torch.tensor([0], device=hidden_states.device),
165 torch.cumsum(prompt_lens[:-1], dim=0)
@@ -220,6 +223,13 @@ class PoolerHead(nn.Module):
223 def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
224 pooling_metadata: PoolingMetadata):
225
226 + # Using float32 in PoolerHead
227 + if isinstance(pooled_data, list):
228 + for i in range(len(pooled_data)):
229 + pooled_data[i] = pooled_data[i].to(torch.float32)
230 + else:
231 + pooled_data = pooled_data.to(torch.float32)
232 +
233 dimensions_list = [
234 pooling_param.dimensions
235 for _, pooling_param in pooling_metadata.seq_groups
vllm/model_executor/models/bert.py
+4 -9
@@ -414,15 +414,10 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
414 intermediate_tensors: Optional[IntermediateTensors] = None,
415 inputs_embeds: Optional[torch.Tensor] = None,
416 ) -> torch.Tensor:
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
417 + return self.model(input_ids=input_ids,
418 + position_ids=positions,
419 + inputs_embeds=inputs_embeds,
420 + intermediate_tensors=intermediate_tensors)
421
422 def pooler(
423 self,
vllm/model_executor/models/bert_with_rope.py
+1 -6
@@ -432,12 +432,7 @@ 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 - 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
435 + return self.encoder(positions, hidden_states)
436
437 def load_weights(self, weights: Iterable[tuple[str,
438 torch.Tensor]]) -> set[str]: