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 14, 05:30 UTC
[New Model]: Support GteNewModelForSequenceClassification (#23524)
Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Aug 28, 2025 at 15:36 UTC
11a7fafaa8807bfeea4b60466c576ec6a7031bfd
13 files changed
+157
-76
docs/models/supported_models.md
+4
@@ -497,6 +497,7 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
497
|--------------|--------|-------------------|----------------------|---------------------------|---------------------|
498
| `BertForSequenceClassification` | BERT-based | `cross-encoder/ms-marco-MiniLM-L-6-v2`, etc. | | | ✅︎ |
499
| `GemmaForSequenceClassification` | Gemma-based | `BAAI/bge-reranker-v2-gemma` (see note), etc. | ✅︎ | ✅︎ | ✅︎ |
500
+| `GteNewForSequenceClassification` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-reranker-base`, etc. | | | ✅︎ |
501
| `Qwen2ForSequenceClassification` | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2` (see note), etc. | ✅︎ | ✅︎ | ✅︎ |
502
| `Qwen3ForSequenceClassification` | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B` (see note), etc. | ✅︎ | ✅︎ | ✅︎ |
503
| `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | | | ✅︎ |
@@ -513,6 +514,9 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
514
vllm serve BAAI/bge-reranker-v2-gemma --hf_overrides '{"architectures": ["GemmaForSequenceClassification"],"classifier_from_token": ["Yes"],"method": "no_post_processing"}'
515
```
516
517
+!!! note
518
+ The second-generation GTE model (mGTE-TRM) is named `NewForSequenceClassification`. The name `NewForSequenceClassification` is too generic, you should set `--hf-overrides '{"architectures": ["GteNewForSequenceClassification"]}'` to specify the use of the `GteNewForSequenceClassification` architecture.
519
+
520
!!! note
521
Load the official original `mxbai-rerank-v2` by using the following command.
522
tests/conftest.py
+2
-3
@@ -456,11 +456,10 @@ class HfRunner:
456
# output is final logits
457
all_inputs = self.get_inputs(prompts)
458
outputs = []
459
+ problem_type = getattr(self.config, "problem_type", "")
460
+
461
for inputs in all_inputs:
462
output = self.model(**self.wrap_device(inputs))
461
-
462
- problem_type = getattr(self.config, "problem_type", "")
463
-
463
if problem_type == "regression":
464
logits = output.logits[0].tolist()
465
elif problem_type == "multi_label_classification":
tests/models/language/pooling/embed_utils.py
+3
@@ -51,6 +51,9 @@ def correctness_test_embed_models(hf_runner,
51
vllm_extra_kwargs = vllm_extra_kwargs or {}
52
vllm_extra_kwargs["dtype"] = model_info.dtype
53
54
+ if model_info.hf_overrides is not None:
55
+ vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides
56
+
57
with vllm_runner(model_info.name,
58
runner="pooling",
59
max_model_len=None,
tests/models/language/pooling/mteb_utils.py
+6
@@ -172,6 +172,9 @@ def mteb_test_embed_models(hf_runner,
172
vllm_extra_kwargs = vllm_extra_kwargs or {}
173
vllm_extra_kwargs["dtype"] = model_info.dtype
174
175
+ if model_info.hf_overrides is not None:
176
+ vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides
177
+
178
with vllm_runner(model_info.name,
179
runner="pooling",
180
max_model_len=None,
@@ -284,6 +287,9 @@ def mteb_test_rerank_models(hf_runner,
287
vllm_extra_kwargs = vllm_extra_kwargs or {}
288
vllm_extra_kwargs["dtype"] = model_info.dtype
289
290
+ if model_info.hf_overrides is not None:
291
+ vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides
292
+
293
with vllm_runner(model_info.name,
294
runner="pooling",
295
max_model_len=None,
tests/models/language/pooling/test_bge_reranker_v2_gemma.py
+9
-15
@@ -13,7 +13,14 @@ from .mteb_utils import VllmMtebEncoder, mteb_test_rerank_models
13
14
RERANK_MODELS = [
15
LASTPoolingRerankModelInfo("BAAI/bge-reranker-v2-gemma",
16
- architecture="GemmaForSequenceClassification"),
16
+ architecture="GemmaForSequenceClassification",
17
+ hf_overrides={
18
+ "architectures":
19
+ ["GemmaForSequenceClassification"],
20
+ "classifier_from_token": ["Yes"],
21
+ "method":
22
+ "no_post_processing",
23
+ }),
24
]
25
26
PROMPT = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'." # noqa: E501
@@ -119,22 +126,9 @@ class GemmaMtebEncoder(VllmMtebEncoder):
126
127
128
@pytest.mark.parametrize("model_info", RERANK_MODELS)
122
-def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo,
123
- monkeypatch) -> None:
124
- monkeypatch.setenv("VLLM_USE_V1", "0")
125
-
126
- assert model_info.architecture == "GemmaForSequenceClassification"
127
-
128
- vllm_extra_kwargs: dict[str, Any] = {
129
- "hf_overrides": {
130
- "architectures": ["GemmaForSequenceClassification"],
131
- "classifier_from_token": ["Yes"],
132
- "method": "no_post_processing",
133
- }
134
- }
129
+def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
130
131
mteb_test_rerank_models(GemmaRerankerHfRunner,
132
vllm_runner,
133
model_info,
139
- vllm_extra_kwargs,
134
vllm_mteb_encoder=GemmaMtebEncoder)
tests/models/language/pooling/test_gte.py
+11
-13
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any
3
4
import pytest
5
@@ -33,12 +32,15 @@ MODELS = [
32
########### NewModel
33
CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
34
architecture="GteNewModel",
35
+ hf_overrides={"architectures": ["GteNewModel"]},
36
enable_test=True),
37
CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-base-en-v1.5",
38
architecture="GteNewModel",
39
+ hf_overrides={"architectures": ["GteNewModel"]},
40
enable_test=True),
41
CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-large-en-v1.5",
42
architecture="GteNewModel",
43
+ hf_overrides={"architectures": ["GteNewModel"]},
44
enable_test=True),
45
########### Qwen2ForCausalLM
46
LASTPoolingEmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
@@ -60,11 +62,16 @@ MODELS = [
62
]
63
64
RERANK_MODELS = [
63
- # classifier_pooling: mean
65
CLSPoolingRerankModelInfo(
66
+ # classifier_pooling: mean
67
"Alibaba-NLP/gte-reranker-modernbert-base",
68
architecture="ModernBertForSequenceClassification",
69
enable_test=True),
70
+ CLSPoolingRerankModelInfo(
71
+ "Alibaba-NLP/gte-multilingual-reranker-base",
72
+ architecture="GteNewForSequenceClassification",
73
+ hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
74
+ enable_test=True),
75
]
76
77
@@ -75,12 +82,7 @@ def test_embed_models_mteb(hf_runner, vllm_runner,
82
check_transformers_version(model_info.name,
83
max_transformers_version="4.53.2")
84
78
- vllm_extra_kwargs: dict[str, Any] = {}
79
- if model_info.architecture == "GteNewModel":
80
- vllm_extra_kwargs["hf_overrides"] = {"architectures": ["GteNewModel"]}
81
-
82
- mteb_test_embed_models(hf_runner, vllm_runner, model_info,
83
- vllm_extra_kwargs)
85
+ mteb_test_embed_models(hf_runner, vllm_runner, model_info)
86
87
88
@pytest.mark.parametrize("model_info", MODELS)
@@ -91,12 +93,8 @@ def test_embed_models_correctness(hf_runner, vllm_runner,
93
check_transformers_version(model_info.name,
94
max_transformers_version="4.53.2")
95
94
- vllm_extra_kwargs: dict[str, Any] = {}
95
- if model_info.architecture == "GteNewModel":
96
- vllm_extra_kwargs["hf_overrides"] = {"architectures": ["GteNewModel"]}
97
-
96
correctness_test_embed_models(hf_runner, vllm_runner, model_info,
99
- example_prompts, vllm_extra_kwargs)
97
+ example_prompts)
98
99
100
@pytest.mark.parametrize("model_info", RERANK_MODELS)
tests/models/language/pooling/test_mxbai_rerank.py
+9
-10
@@ -10,12 +10,20 @@ from tests.conftest import HfRunner
10
from ...utils import LASTPoolingRerankModelInfo, RerankModelInfo
11
from .mteb_utils import mteb_test_rerank_models
12
13
+mxbai_rerank_hf_overrides = {
14
+ "architectures": ["Qwen2ForSequenceClassification"],
15
+ "classifier_from_token": ["0", "1"],
16
+ "method": "from_2_way_softmax",
17
+}
18
+
19
RERANK_MODELS = [
20
LASTPoolingRerankModelInfo("mixedbread-ai/mxbai-rerank-base-v2",
21
architecture="Qwen2ForSequenceClassification",
22
+ hf_overrides=mxbai_rerank_hf_overrides,
23
enable_test=True),
24
LASTPoolingRerankModelInfo("mixedbread-ai/mxbai-rerank-large-v2",
25
architecture="Qwen2ForSequenceClassification",
26
+ hf_overrides=mxbai_rerank_hf_overrides,
27
enable_test=False)
28
]
29
@@ -71,13 +79,4 @@ class MxbaiRerankerHfRunner(HfRunner):
79
80
@pytest.mark.parametrize("model_info", RERANK_MODELS)
81
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
74
- vllm_extra_kwargs: dict[str, Any] = {}
75
- if model_info.architecture == "Qwen2ForSequenceClassification":
76
- vllm_extra_kwargs["hf_overrides"] = {
77
- "architectures": ["Qwen2ForSequenceClassification"],
78
- "classifier_from_token": ["0", "1"],
79
- "method": "from_2_way_softmax",
80
- }
81
-
82
- mteb_test_rerank_models(MxbaiRerankerHfRunner, vllm_runner, model_info,
83
- vllm_extra_kwargs)
82
+ mteb_test_rerank_models(MxbaiRerankerHfRunner, vllm_runner, model_info)
tests/models/language/pooling/test_qwen3_reranker.py
+9
-17
@@ -11,12 +11,20 @@ from tests.utils import multi_gpu_test
11
from ...utils import LASTPoolingRerankModelInfo, RerankModelInfo
12
from .mteb_utils import mteb_test_rerank_models
13
14
+qwen3_reranker_hf_overrides = {
15
+ "architectures": ["Qwen3ForSequenceClassification"],
16
+ "classifier_from_token": ["no", "yes"],
17
+ "is_original_qwen3_reranker": True,
18
+}
19
+
20
RERANK_MODELS = [
21
LASTPoolingRerankModelInfo("Qwen/Qwen3-Reranker-0.6B",
22
architecture="Qwen3ForSequenceClassification",
23
+ hf_overrides=qwen3_reranker_hf_overrides,
24
enable_test=True),
25
LASTPoolingRerankModelInfo("Qwen/Qwen3-Reranker-4B",
26
architecture="Qwen3ForSequenceClassification",
27
+ hf_overrides=qwen3_reranker_hf_overrides,
28
enable_test=False)
29
]
30
@@ -74,18 +82,7 @@ class Qwen3RerankerHfRunner(HfRunner):
82
@pytest.mark.parametrize("model_info", RERANK_MODELS)
83
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
84
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
- mteb_test_rerank_models(Qwen3RerankerHfRunner, vllm_runner, model_info,
88
- vllm_extra_kwargs)
85
+ mteb_test_rerank_models(Qwen3RerankerHfRunner, vllm_runner, model_info)
86
87
88
@pytest.mark.parametrize("model_info", RERANK_MODELS)
@@ -96,11 +93,6 @@ def test_rerank_models_mteb_tp(vllm_runner,
93
assert model_info.architecture == "Qwen3ForSequenceClassification"
94
95
vllm_extra_kwargs: dict[str, Any] = {
99
- "hf_overrides": {
100
- "architectures": ["Qwen3ForSequenceClassification"],
101
- "classifier_from_token": ["no", "yes"],
102
- "is_original_qwen3_reranker": True,
103
- },
96
"tensor_parallel_size": 2,
97
}
98
tests/models/registry.py
+4
@@ -365,6 +365,10 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
365
366
# [Cross-encoder]
367
"BertForSequenceClassification": _HfExamplesInfo("cross-encoder/ms-marco-MiniLM-L-6-v2", v0_only=True), # noqa: E501
368
+ "GteNewForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-multilingual-reranker-base", # noqa: E501
369
+ trust_remote_code=True,
370
+ hf_overrides={
371
+ "architectures": ["GteNewForSequenceClassification"]}),# noqa: E501
372
"ModernBertForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-reranker-modernbert-base", v0_only=True), # noqa: E501
373
"RobertaForSequenceClassification": _HfExamplesInfo("cross-encoder/quora-roberta-base", v0_only=True), # noqa: E501
374
"XLMRobertaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-m3", v0_only=True), # noqa: E501
tests/models/utils.py
+18
-10
@@ -3,7 +3,8 @@
3
4
import warnings
5
from collections.abc import Sequence
6
-from typing import Any, NamedTuple, Optional, Union
6
+from dataclasses import dataclass
7
+from typing import Any, Optional, Union
8
9
import torch
10
import torch.nn.functional as F
@@ -339,36 +340,43 @@ def softmax(data):
340
return F.softmax(data, dim=-1)
341
342
342
-class EmbedModelInfo(NamedTuple):
343
+@dataclass
344
+class ModelInfo:
345
name: str
344
- is_matryoshka: bool = False
345
- matryoshka_dimensions: Optional[list[int]] = None
346
architecture: str = ""
347
dtype: str = "auto"
348
+ hf_overrides: Optional[dict[str, Any]] = None
349
default_pooling_type: str = ""
350
enable_test: bool = True
351
352
353
+@dataclass
354
+class EmbedModelInfo(ModelInfo):
355
+ is_matryoshka: bool = False
356
+ matryoshka_dimensions: Optional[list[int]] = None
357
+
358
+
359
+@dataclass
360
class CLSPoolingEmbedModelInfo(EmbedModelInfo):
361
default_pooling_type: str = "CLS"
362
363
364
+@dataclass
365
class LASTPoolingEmbedModelInfo(EmbedModelInfo):
366
default_pooling_type: str = "LAST"
367
368
360
-class RerankModelInfo(NamedTuple):
361
- name: str
362
- architecture: str = ""
363
- dtype: str = "auto"
364
- default_pooling_type: str = ""
365
- enable_test: bool = True
369
+@dataclass
370
+class RerankModelInfo(ModelInfo):
371
+ pass
372
373
374
+@dataclass
375
class CLSPoolingRerankModelInfo(RerankModelInfo):
376
default_pooling_type: str = "CLS"
377
378
379
+@dataclass
380
class LASTPoolingRerankModelInfo(RerankModelInfo):
381
default_pooling_type: str = "LAST"
382
vllm/model_executor/models/bert_with_rope.py
+77
-6
@@ -27,12 +27,15 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
27
from vllm.model_executor.layers.vocab_parallel_embedding import (
28
VocabParallelEmbedding)
29
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
30
-from vllm.model_executor.models.utils import WeightsMapper
30
+from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
31
+ maybe_prefix)
32
from vllm.model_executor.utils import set_weight_attrs
33
from vllm.platforms import current_platform
34
from vllm.sequence import IntermediateTensors
35
35
-from .interfaces import SupportsQuant
36
+from ..layers.pooler import ClassifierPooler, DispatchPooler, Pooler
37
+from .bert import BertPooler
38
+from .interfaces import SupportsCrossEncoding, SupportsQuant
39
from .interfaces_base import default_pooling_type
40
41
@@ -406,9 +409,14 @@ class BertWithRopeEncoder(nn.Module):
409
class BertWithRope(nn.Module, SupportsQuant):
410
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
411
409
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
412
+ def __init__(self,
413
+ *,
414
+ vllm_config: VllmConfig,
415
+ prefix: str = "",
416
+ add_pooling_layer: bool = False):
417
super().__init__()
418
self.vllm_config = vllm_config
419
+ self.add_pooling_layer = add_pooling_layer
420
self.config = vllm_config.model_config.hf_config
421
self.embeddings = BertWithRopeEmbedding(self.config)
422
self.encoder = BertWithRopeEncoder(
@@ -416,6 +424,7 @@ class BertWithRope(nn.Module, SupportsQuant):
424
bias=getattr(self.config, "bias", True),
425
rotary_kwargs=self.config.rotary_kwargs,
426
prefix=f"{prefix}.encoder")
427
+ self.pooler = BertPooler(self.config) if add_pooling_layer else None
428
429
def forward(
430
self,
@@ -448,7 +457,7 @@ class BertWithRope(nn.Module, SupportsQuant):
457
params_dict = dict(self.named_parameters())
458
loaded_params: set[str] = set()
459
for name, loaded_weight in weights:
451
- if "pooler" in name:
460
+ if not self.add_pooling_layer and "pooler" in name:
461
continue
462
for (param_name, weight_name, shard_id) in stacked_params_mapping:
463
if weight_name not in name:
@@ -508,8 +517,8 @@ class GteNewModel(BertWithRope):
517
"attention.o_proj": "attn.out_proj",
518
})
519
511
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
512
- super().__init__(vllm_config=vllm_config, prefix=prefix)
520
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs):
521
+ super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
522
523
# GteNewModel only gate_up_proj does not have bias.
524
# Hack method learned from vllm/model_executor/models/glm.py
@@ -614,3 +623,65 @@ class JinaRobertaModel(BertWithRope):
623
torch.Tensor]]) -> set[str]:
624
weights = self.jina_merge_lora_weights(weights)
625
return super().load_weights(weights)
626
+
627
+
628
+@default_pooling_type("CLS")
629
+class GteNewForSequenceClassification(nn.Module, SupportsCrossEncoding):
630
+ is_pooling_model = True
631
+
632
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
633
+ super().__init__()
634
+ config = vllm_config.model_config.hf_config
635
+ quant_config = vllm_config.quant_config
636
+
637
+ self.new = GteNewModel(vllm_config=vllm_config,
638
+ prefix=prefix,
639
+ add_pooling_layer=True)
640
+ self.classifier = RowParallelLinear(config.hidden_size,
641
+ config.num_labels,
642
+ input_is_parallel=False,
643
+ bias=True,
644
+ quant_config=quant_config,
645
+ prefix=maybe_prefix(
646
+ prefix, "classifier"),
647
+ return_bias=False)
648
+
649
+ pooler_config = vllm_config.model_config.pooler_config
650
+ assert pooler_config is not None
651
+
652
+ self.pooler = DispatchPooler({
653
+ "encode":
654
+ Pooler.for_encode(pooler_config),
655
+ "classify":
656
+ ClassifierPooler(
657
+ pooling=self.new.pooler,
658
+ classifier=self.classifier,
659
+ act_fn=ClassifierPooler.act_fn_for_seq_cls(
660
+ vllm_config.model_config),
661
+ ),
662
+ "score":
663
+ ClassifierPooler(
664
+ pooling=self.new.pooler,
665
+ classifier=self.classifier,
666
+ act_fn=ClassifierPooler.act_fn_for_cross_encoder(
667
+ vllm_config.model_config),
668
+ ),
669
+ })
670
+
671
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
672
+ loader = AutoWeightsLoader(self)
673
+ loaded_params = loader.load_weights(weights)
674
+ return loaded_params
675
+
676
+ def forward(
677
+ self,
678
+ input_ids: Optional[torch.Tensor],
679
+ positions: torch.Tensor,
680
+ intermediate_tensors: Optional[IntermediateTensors] = None,
681
+ inputs_embeds: Optional[torch.Tensor] = None,
682
+ ) -> torch.Tensor:
683
+
684
+ return self.new(input_ids=input_ids,
685
+ positions=positions,
686
+ inputs_embeds=inputs_embeds,
687
+ intermediate_tensors=intermediate_tensors)
vllm/model_executor/models/config.py
+1
@@ -406,6 +406,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
406
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
407
"GteModel": SnowflakeGteNewModelConfig,
408
"GteNewModel": GteNewModelConfig,
409
+ "GteNewForSequenceClassification": GteNewModelConfig,
410
"NomicBertModel": NomicBertModelConfig,
411
"Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
412
"Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
vllm/model_executor/models/registry.py
+4
-2
@@ -191,12 +191,14 @@ _EMBEDDING_MODELS = {
191
192
_CROSS_ENCODER_MODELS = {
193
"BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
194
+ "GteNewForSequenceClassification": ("bert_with_rope",
195
+ "GteNewForSequenceClassification"),
196
+ "ModernBertForSequenceClassification": ("modernbert",
197
+ "ModernBertForSequenceClassification"),
198
"RobertaForSequenceClassification": ("roberta",
199
"RobertaForSequenceClassification"),
200
"XLMRobertaForSequenceClassification": ("roberta",
201
"RobertaForSequenceClassification"),
198
- "ModernBertForSequenceClassification": ("modernbert",
199
- "ModernBertForSequenceClassification"),
202
# [Auto-converted (see adapters.py)]
203
"JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"), # noqa: E501,
204
}