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 6, 19:15 UTC
[Model] Add support for BERT-like Chinese ERNIE pooling models (#36385)
Signed-off-by: whyiug <whyiug@hotmail.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
whyiug committed
Mar 13, 2026 at 11:23 UTC
1ce13cf9926dda15cd3fe40296411ac721979401
7 files changed
+317
-3
docs/models/supported_models.md
+5
-1
@@ -514,6 +514,7 @@ These models primarily support the [`LLM.embed`](./pooling_models.md#llmembed) A
514
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
515
| `BertModel`<sup>C</sup> | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | |
516
| `BertSpladeSparseEmbeddingModel` | SPLADE | `naver/splade-v3` | | |
517
+| `ErnieModel` | BERT-like Chinese ERNIE | `shibing624/text2vec-base-chinese-sentence` | | |
518
| `Gemma2Model`<sup>C</sup> | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | ✅︎ |
519
| `Gemma3TextModel`<sup>C</sup> | Gemma 3-based | `google/embeddinggemma-300m`, etc. | ✅︎ | ✅︎ |
520
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
@@ -556,8 +557,9 @@ These models primarily support the [`LLM.classify`](./pooling_models.md#llmclass
557
558
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
559
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
559
-| `JambaForSequenceClassification` | Jamba | `ai21labs/Jamba-tiny-reward-dev`, etc. | ✅︎ | ✅︎ |
560
+| `ErnieForSequenceClassification` | BERT-like Chinese ERNIE | `Forrest20231206/ernie-3.0-base-zh-cls` | | |
561
| `GPT2ForSequenceClassification` | GPT2 | `nie3e/sentiment-polish-gpt2-small` | | |
562
+| `JambaForSequenceClassification` | Jamba | `ai21labs/Jamba-tiny-reward-dev`, etc. | ✅︎ | ✅︎ |
563
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
564
565
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./pooling_models.md#model-conversion))
@@ -574,6 +576,7 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
576
| Architecture | Models | Example HF Models | Score template (see note) | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
577
| ------------ | ------ | ----------------- | ------------------------- | --------------------------- | --------------------------------------- |
578
| `BertForSequenceClassification` | BERT-based | `cross-encoder/ms-marco-MiniLM-L-6-v2`, etc. | N/A | | |
579
+| `ErnieForSequenceClassification` | BERT-like Chinese ERNIE | `Forrest20231206/ernie-3.0-base-zh-cls` | N/A | | |
580
| `GemmaForSequenceClassification` | Gemma-based | `BAAI/bge-reranker-v2-gemma`(see note), etc. | [bge-reranker-v2-gemma.jinja](../../examples/pooling/score/template/bge-reranker-v2-gemma.jinja) | ✅︎ | ✅︎ |
581
| `GteNewForSequenceClassification` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-reranker-base`, etc. | N/A | | |
582
| `LlamaBidirectionalForSequenceClassification`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-rerank-1b-v2`, etc. | [nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja) | ✅︎ | ✅︎ |
@@ -639,6 +642,7 @@ These models primarily support the [`LLM.encode`](./pooling_models.md#llmencode)
642
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
643
| ------------ | ------ | ----------------- | --------------------------- | --------------------------------------- |
644
| `BertForTokenClassification` | bert-based | `boltuix/NeuroBERT-NER` (see note), etc. | | |
645
+| `ErnieForTokenClassification` | BERT-like Chinese ERNIE | `gyr66/Ernie-3.0-base-chinese-finetuned-ner` | | |
646
| `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | |
647
648
!!! note
tests/models/language/pooling/test_classification.py
+2
@@ -18,6 +18,7 @@ from vllm.platforms import current_platform
18
pytest.mark.slow_test,
19
],
20
),
21
+ pytest.param("Forrest20231206/ernie-3.0-base-zh-cls"),
22
],
23
)
24
@pytest.mark.parametrize("dtype", ["half"] if current_platform.is_rocm() else ["float"])
@@ -47,5 +48,6 @@ def test_models(
48
assert torch.allclose(
49
hf_output,
50
vllm_output,
51
+ atol=1e-3 if dtype == "float" else 1e-2,
52
rtol=2e-3 if dtype == "float" else 1e-2,
53
)
tests/models/language/pooling/test_token_classification.py
+8
-2
@@ -25,11 +25,17 @@ def seed_everything():
25
yield
26
27
28
-@pytest.mark.parametrize("model", ["boltuix/NeuroBERT-NER"])
28
+@pytest.mark.parametrize(
29
+ "model",
30
+ [
31
+ "boltuix/NeuroBERT-NER",
32
+ "gyr66/Ernie-3.0-base-chinese-finetuned-ner",
33
+ ],
34
+)
35
# The float32 is required for this tiny model to pass the test.
36
@pytest.mark.parametrize("dtype", ["float"])
37
@torch.inference_mode
32
-def test_bert_models(
38
+def test_bert_like_models(
39
hf_runner,
40
vllm_runner,
41
example_prompts,
tests/models/language/pooling_mteb_test/test_ernie.py
new
+45
@@ -0,0 +1,45 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from tests.models.language.pooling.embed_utils import correctness_test_embed_models
7
+from tests.models.utils import EmbedModelInfo
8
+
9
+from .mteb_embed_utils import mteb_test_embed_models
10
+
11
+MODELS = [
12
+ EmbedModelInfo(
13
+ "shibing624/text2vec-base-chinese-sentence",
14
+ architecture="ErnieModel",
15
+ mteb_score=0.536523112,
16
+ seq_pooling_type="MEAN",
17
+ attn_type="encoder_only",
18
+ is_prefix_caching_supported=False,
19
+ is_chunked_prefill_supported=False,
20
+ enable_test=True,
21
+ ),
22
+]
23
+
24
+
25
+@pytest.mark.parametrize("model_info", MODELS)
26
+def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
27
+ mteb_test_embed_models(
28
+ hf_runner,
29
+ vllm_runner,
30
+ model_info,
31
+ vllm_extra_kwargs={"gpu_memory_utilization": 0.2},
32
+ )
33
+
34
+
35
+@pytest.mark.parametrize("model_info", MODELS)
36
+def test_embed_models_correctness(
37
+ hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
38
+) -> None:
39
+ correctness_test_embed_models(
40
+ hf_runner,
41
+ vllm_runner,
42
+ model_info,
43
+ example_prompts,
44
+ vllm_extra_kwargs={"gpu_memory_utilization": 0.2},
45
+ )
tests/models/registry.py
+7
@@ -552,6 +552,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
552
_EMBEDDING_EXAMPLE_MODELS = {
553
# [Text-only]
554
"BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
555
+ "ErnieModel": _HfExamplesInfo("shibing624/text2vec-base-chinese-sentence"),
556
"BertSpladeSparseEmbeddingModel": _HfExamplesInfo(
557
"naver/splade-v3",
558
hf_overrides={"architectures": ["BertSpladeSparseEmbeddingModel"]},
@@ -666,6 +667,9 @@ _REWARD_EXAMPLE_MODELS = {
667
668
_TOKEN_CLASSIFICATION_EXAMPLE_MODELS = {
669
"BertForTokenClassification": _HfExamplesInfo("boltuix/NeuroBERT-NER"),
670
+ "ErnieForTokenClassification": _HfExamplesInfo(
671
+ "gyr66/Ernie-3.0-base-chinese-finetuned-ner"
672
+ ),
673
"ModernBertForTokenClassification": _HfExamplesInfo(
674
"disham993/electrical-ner-ModernBERT-base"
675
),
@@ -675,6 +679,9 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
679
"BertForSequenceClassification": _HfExamplesInfo(
680
"cross-encoder/ms-marco-MiniLM-L-6-v2"
681
),
682
+ "ErnieForSequenceClassification": _HfExamplesInfo(
683
+ "Forrest20231206/ernie-3.0-base-zh-cls",
684
+ ),
685
"GPT2ForSequenceClassification": _HfExamplesInfo(
686
"nie3e/sentiment-polish-gpt2-small"
687
),
vllm/model_executor/models/ernie.py
new
+247
@@ -0,0 +1,247 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Iterable
5
+
6
+import torch
7
+from torch import nn
8
+from transformers import BertConfig
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.model_executor.layers.pooler import DispatchPooler
12
+from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
13
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
14
+from vllm.sequence import IntermediateTensors
15
+
16
+from .bert import (
17
+ TOKEN_TYPE_SHIFT,
18
+ BertEmbedding,
19
+ BertEmbeddingModel,
20
+ BertModel,
21
+ BertPoolingModel,
22
+ _decode_token_type_ids,
23
+ _encode_token_type_ids,
24
+)
25
+from .interfaces import SupportsCrossEncoding, SupportsQuant
26
+from .interfaces_base import attn_type, default_pooling_type
27
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
28
+
29
+_LEGACY_SUFFIX_MAPPER = WeightsMapper(
30
+ orig_to_new_suffix={
31
+ ".gamma": ".weight",
32
+ ".beta": ".bias",
33
+ }
34
+)
35
+
36
+
37
+class ErnieEmbedding(BertEmbedding):
38
+ def __init__(self, config: BertConfig):
39
+ super().__init__(config)
40
+
41
+ task_type_vocab_size = max(1, getattr(config, "task_type_vocab_size", 1))
42
+ self.task_type_embeddings = VocabParallelEmbedding(
43
+ task_type_vocab_size, config.hidden_size
44
+ )
45
+
46
+ def forward(
47
+ self,
48
+ input_ids: torch.Tensor,
49
+ position_ids: torch.Tensor,
50
+ inputs_embeds: torch.Tensor | None = None,
51
+ ) -> torch.Tensor:
52
+ token_type_ids = _decode_token_type_ids(input_ids)
53
+ task_type_ids = torch.zeros_like(token_type_ids)
54
+
55
+ if inputs_embeds is None:
56
+ inputs_embeds = self.word_embeddings(input_ids)
57
+
58
+ position_embeddings = self.position_embeddings(position_ids)
59
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
60
+ task_type_embeddings = self.task_type_embeddings(task_type_ids)
61
+
62
+ embeddings = (
63
+ inputs_embeds
64
+ + token_type_embeddings
65
+ + task_type_embeddings
66
+ + position_embeddings
67
+ )
68
+ embeddings = self.LayerNorm(embeddings)
69
+ return embeddings
70
+
71
+
72
+@default_pooling_type(seq_pooling_type="CLS")
73
+class ErnieModel(BertModel):
74
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
75
+ super().__init__(
76
+ vllm_config=vllm_config,
77
+ prefix=prefix,
78
+ embedding_class=ErnieEmbedding,
79
+ )
80
+
81
+
82
+class ErniePoolingModel(BertPoolingModel):
83
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
84
+ super().__init__(
85
+ vllm_config=vllm_config,
86
+ prefix=prefix,
87
+ embedding_class=ErnieEmbedding,
88
+ )
89
+
90
+
91
+@default_pooling_type(seq_pooling_type="CLS")
92
+class ErnieEmbeddingModel(BertEmbeddingModel):
93
+ def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> ErnieModel:
94
+ return ErnieModel(vllm_config=vllm_config, prefix=prefix)
95
+
96
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
97
+ weights_list = list(weights)
98
+ has_model_prefix = any(name.startswith("model.") for name, _ in weights_list)
99
+ has_ernie_prefix = any(name.startswith("ernie.") for name, _ in weights_list)
100
+
101
+ mapper: WeightsMapper | None = None
102
+ if not has_model_prefix:
103
+ if has_ernie_prefix:
104
+ mapper = WeightsMapper(orig_to_new_prefix={"ernie.": "model."})
105
+ else:
106
+ mapper = WeightsMapper(orig_to_new_prefix={"": "model."})
107
+ if mapper is None:
108
+ mapper = _LEGACY_SUFFIX_MAPPER
109
+ else:
110
+ mapper = mapper | _LEGACY_SUFFIX_MAPPER
111
+
112
+ loader = AutoWeightsLoader(self, skip_prefixes=["lm_head.", "cls."])
113
+ return loader.load_weights(weights_list, mapper=mapper)
114
+
115
+
116
+@default_pooling_type(seq_pooling_type="CLS")
117
+class ErnieForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQuant):
118
+ is_pooling_model = True
119
+
120
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
121
+ super().__init__()
122
+ config = vllm_config.model_config.hf_config
123
+
124
+ self.num_labels = config.num_labels
125
+ self.ernie = ErniePoolingModel(
126
+ vllm_config=vllm_config,
127
+ prefix=maybe_prefix(prefix, "ernie"),
128
+ )
129
+ self.classifier = nn.Linear(
130
+ config.hidden_size,
131
+ config.num_labels,
132
+ dtype=vllm_config.model_config.head_dtype,
133
+ )
134
+
135
+ pooler_config = vllm_config.model_config.pooler_config
136
+ assert pooler_config is not None
137
+
138
+ self.pooler = DispatchPooler.for_seq_cls(
139
+ pooler_config,
140
+ pooling=self.ernie.pooler,
141
+ classifier=self.classifier,
142
+ )
143
+
144
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
145
+ return self.ernie.embed_input_ids(input_ids)
146
+
147
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
148
+ weights_list = list(weights)
149
+ has_ernie_prefix = any(name.startswith("ernie.") for name, _ in weights_list)
150
+ has_bert_prefix = any(name.startswith("bert.") for name, _ in weights_list)
151
+
152
+ mapper: WeightsMapper | None = None
153
+ if has_bert_prefix and not has_ernie_prefix:
154
+ mapper = WeightsMapper(orig_to_new_prefix={"bert.": "ernie."})
155
+ if mapper is None:
156
+ mapper = _LEGACY_SUFFIX_MAPPER
157
+ else:
158
+ mapper = mapper | _LEGACY_SUFFIX_MAPPER
159
+
160
+ loader = AutoWeightsLoader(self, skip_prefixes=["cls.", "lm_head."])
161
+ return loader.load_weights(weights_list, mapper=mapper)
162
+
163
+ def forward(
164
+ self,
165
+ input_ids: torch.Tensor | None,
166
+ positions: torch.Tensor,
167
+ intermediate_tensors: IntermediateTensors | None = None,
168
+ inputs_embeds: torch.Tensor | None = None,
169
+ token_type_ids: torch.Tensor | None = None,
170
+ ) -> torch.Tensor:
171
+ if token_type_ids is not None:
172
+ assert self.ernie.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)
173
+ assert input_ids is not None
174
+ _encode_token_type_ids(input_ids, token_type_ids)
175
+
176
+ return self.ernie(
177
+ input_ids=input_ids,
178
+ positions=positions,
179
+ inputs_embeds=inputs_embeds,
180
+ intermediate_tensors=intermediate_tensors,
181
+ )
182
+
183
+
184
+@attn_type("encoder_only")
185
+@default_pooling_type(tok_pooling_type="ALL")
186
+class ErnieForTokenClassification(nn.Module):
187
+ is_pooling_model = True
188
+
189
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
190
+ super().__init__()
191
+ config = vllm_config.model_config.hf_config
192
+ self.head_dtype = vllm_config.model_config.head_dtype
193
+ self.num_labels = config.num_labels
194
+ self.ernie = ErnieModel(
195
+ vllm_config=vllm_config,
196
+ prefix=maybe_prefix(prefix, "ernie"),
197
+ )
198
+ self.classifier = nn.Linear(
199
+ config.hidden_size, config.num_labels, dtype=self.head_dtype
200
+ )
201
+
202
+ pooler_config = vllm_config.model_config.pooler_config
203
+ assert pooler_config is not None
204
+
205
+ self.pooler = pooler_for_token_classify(pooler_config)
206
+
207
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
208
+ return self.ernie.embed_input_ids(input_ids)
209
+
210
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
211
+ weights_list = list(weights)
212
+ has_ernie_prefix = any(name.startswith("ernie.") for name, _ in weights_list)
213
+ has_bert_prefix = any(name.startswith("bert.") for name, _ in weights_list)
214
+
215
+ mapper: WeightsMapper | None = None
216
+ if has_bert_prefix and not has_ernie_prefix:
217
+ mapper = WeightsMapper(orig_to_new_prefix={"bert.": "ernie."})
218
+ if mapper is None:
219
+ mapper = _LEGACY_SUFFIX_MAPPER
220
+ else:
221
+ mapper = mapper | _LEGACY_SUFFIX_MAPPER
222
+
223
+ loader = AutoWeightsLoader(self, skip_prefixes=["cls.", "lm_head."])
224
+ return loader.load_weights(weights_list, mapper=mapper)
225
+
226
+ def forward(
227
+ self,
228
+ input_ids: torch.Tensor | None,
229
+ positions: torch.Tensor,
230
+ intermediate_tensors: IntermediateTensors | None = None,
231
+ inputs_embeds: torch.Tensor | None = None,
232
+ token_type_ids: torch.Tensor | None = None,
233
+ ) -> torch.Tensor:
234
+ if token_type_ids is not None:
235
+ assert self.ernie.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)
236
+ assert input_ids is not None
237
+ _encode_token_type_ids(input_ids, token_type_ids)
238
+
239
+ hidden_states = self.ernie(
240
+ input_ids=input_ids,
241
+ positions=positions,
242
+ inputs_embeds=inputs_embeds,
243
+ intermediate_tensors=intermediate_tensors,
244
+ )
245
+
246
+ hidden_states = hidden_states.to(self.head_dtype)
247
+ return self.classifier(hidden_states)
vllm/model_executor/models/registry.py
+3
@@ -214,6 +214,7 @@ _EMBEDDING_MODELS = {
214
# [Text-only]
215
"BertModel": ("bert", "BertEmbeddingModel"),
216
"BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),
217
+ "ErnieModel": ("ernie", "ErnieEmbeddingModel"),
218
"BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
219
"DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
220
"Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
@@ -286,6 +287,7 @@ _REWARD_MODELS = {
287
288
_TOKEN_CLASSIFICATION_MODELS = {
289
"BertForTokenClassification": ("bert", "BertForTokenClassification"),
290
+ "ErnieForTokenClassification": ("ernie", "ErnieForTokenClassification"),
291
"ModernBertForTokenClassification": (
292
"modernbert",
293
"ModernBertForTokenClassification",
@@ -295,6 +297,7 @@ _TOKEN_CLASSIFICATION_MODELS = {
297
_SEQUENCE_CLASSIFICATION_MODELS = {
298
"BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
299
"GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
300
+ "ErnieForSequenceClassification": ("ernie", "ErnieForSequenceClassification"),
301
"GteNewForSequenceClassification": (
302
"bert_with_rope",
303
"GteNewForSequenceClassification",