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, 02:15 UTC
[Bugfix] Fix tensor parallel issue in Qwen3 reranker weight loading (#20682)
Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com>
yurhett committed
Jul 12, 2025 at 11:52 UTC
11c0198615802ea0b17d877ebdf2e61cd4cfca0c
3 files changed
+38
-7
tests/models/language/pooling/mteb_utils.py
+3
-2
@@ -268,7 +268,8 @@ def mteb_test_rerank_models(hf_runner,
268
model_info: RerankModelInfo,
269
vllm_extra_kwargs=None,
270
hf_model_callback=None,
271
- vllm_mteb_encoder=VllmMtebEncoder):
271
+ vllm_mteb_encoder=VllmMtebEncoder,
272
+ atol=MTEB_RERANK_TOL):
273
if not model_info.enable_test:
274
# A model family has many models with the same architecture,
275
# and we don't need to test each one.
@@ -301,4 +302,4 @@ def mteb_test_rerank_models(hf_runner,
302
print("SentenceTransformers:", st_dtype, st_main_score)
303
print("Difference:", st_main_score - vllm_main_score)
304
304
- assert st_main_score == pytest.approx(vllm_main_score, abs=MTEB_RERANK_TOL)
305
+ assert st_main_score == pytest.approx(vllm_main_score, abs=atol)
tests/models/language/pooling/test_qwen3_reranker.py
+27
@@ -6,6 +6,7 @@ import pytest
6
import torch
7
8
from tests.conftest import HfRunner
9
+from tests.utils import multi_gpu_test
10
11
from .mteb_utils import RerankModelInfo, mteb_test_rerank_models
12
@@ -87,3 +88,29 @@ def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
88
89
mteb_test_rerank_models(Qwen3RerankerHfRunner, vllm_runner, model_info,
90
vllm_extra_kwargs)
91
+
92
+
93
+@pytest.mark.parametrize("model_info", RERANK_MODELS)
94
+@multi_gpu_test(num_gpus=2)
95
+def test_rerank_models_mteb_tp(vllm_runner,
96
+ model_info: RerankModelInfo) -> None:
97
+
98
+ assert model_info.architecture == "Qwen3ForSequenceClassification"
99
+
100
+ vllm_extra_kwargs: dict[str, Any] = {
101
+ "hf_overrides": {
102
+ "architectures": ["Qwen3ForSequenceClassification"],
103
+ "classifier_from_token": ["no", "yes"],
104
+ "is_original_qwen3_reranker": True,
105
+ },
106
+ "tensor_parallel_size": 2,
107
+ }
108
+
109
+ if model_info.name == "Qwen/Qwen3-Reranker-4B":
110
+ vllm_extra_kwargs["max_num_seqs"] = 1
111
+
112
+ mteb_test_rerank_models(Qwen3RerankerHfRunner,
113
+ vllm_runner,
114
+ model_info,
115
+ vllm_extra_kwargs,
116
+ atol=1.2e-2)
vllm/model_executor/models/adapters.py
+8
-5
@@ -322,6 +322,8 @@ def load_weights_using_from_2_way_softmax(
322
# refer to https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
323
from vllm.model_executor.layers.vocab_parallel_embedding import (
324
ParallelLMHead)
325
+ from vllm.model_executor.model_loader.weight_utils import (
326
+ default_weight_loader)
327
from vllm.model_executor.models.utils import AutoWeightsLoader
328
329
model_config = model.vllm_config.model_config
@@ -329,8 +331,6 @@ def load_weights_using_from_2_way_softmax(
331
tokens = cast(list[int], tokens)
332
assert len(tokens) == 2
333
332
- device = model.score.weight.device
333
-
334
if model.config.tie_word_embeddings:
335
model.lm_head = model.model.embed_tokens
336
else:
@@ -349,10 +349,13 @@ def load_weights_using_from_2_way_softmax(
349
350
false_id = tokenizer.convert_tokens_to_ids(tokens[0])
351
true_id = tokenizer.convert_tokens_to_ids(tokens[1])
352
- weight = model.lm_head.weight.data[true_id].to(device).to(
353
- torch.float32) - model.lm_head.weight.data[false_id].to(device).to(
352
+ weight = model.lm_head.weight.data[[true_id]].to(
353
+ torch.float32) - model.lm_head.weight.data[[false_id]].to(
354
torch.float32)
355
- model.score.weight.data.copy_(weight)
355
+
356
+ param = model.score.weight
357
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
358
+ weight_loader(param, weight)
359
360
del model.lm_head
361
loaded_weights.add("score.weight")