[Model] Add FP8 quantization support for ModernBERT (#53101)
Signed-off-by: thunguo <tew@apache.org> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
ThunGuo committed
Aug 24, 2026 at 11:04 UTC
e8888b2d68bd7c6cce0aada7f0e214e55020e20d
3 files changed
+96
-6
tests/models/language/pooling_mteb_test/mteb_embed_utils.py
+4
@@ -163,6 +163,7 @@ def mteb_test_embed_models(
163
hf_model_callback=None,
164
atol=MTEB_EMBED_TOL,
165
prompt_prefix: str | None = None,
166
+ vllm_model_callback=None,
167
):
168
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
169
@@ -178,6 +179,9 @@ def mteb_test_embed_models(
179
) as vllm_model:
180
model_config = vllm_model.llm.llm_engine.model_config
181
182
+ if vllm_model_callback is not None:
183
+ vllm_model_callback(vllm_model)
184
+
185
# Confirm whether vllm is using the correct architecture
186
if model_info.architecture:
187
assert model_info.architecture in model_config.architectures
tests/models/language/pooling_mteb_test/test_modernbert_fp8.py
new
+59
@@ -0,0 +1,59 @@
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.utils import EmbedModelInfo
7
+from tests.quantization.utils import is_quant_method_supported
8
+
9
+from .mteb_embed_utils import mteb_test_embed_models
10
+
11
+MODEL_INFO = EmbedModelInfo(
12
+ "Alibaba-NLP/gte-modernbert-base",
13
+ mteb_score=0.748193353,
14
+ architecture="ModernBertModel",
15
+ seq_pooling_type="CLS",
16
+ attn_type="encoder_only",
17
+ is_prefix_caching_supported=False,
18
+ is_chunked_prefill_supported=False,
19
+ enable_test=True,
20
+)
21
+
22
+
23
+def _assert_modernbert_online_fp8(model) -> None:
24
+ from vllm.model_executor.layers.quantization.online.fp8 import (
25
+ Fp8PerTensorOnlineLinearMethod,
26
+ )
27
+
28
+ layer = model.encoder_layer.layers[0]
29
+ linears = {
30
+ "attn.Wqkv": layer.attn.Wqkv,
31
+ "attn.Wo": layer.attn.Wo,
32
+ "mlp.Wi": layer.mlp.Wi,
33
+ "mlp.Wo": layer.mlp.Wo,
34
+ }
35
+ for name, linear in linears.items():
36
+ assert isinstance(linear.quant_method, Fp8PerTensorOnlineLinearMethod), (
37
+ f"{name} was not quantized with online FP8"
38
+ )
39
+
40
+
41
+@pytest.mark.skipif(
42
+ not is_quant_method_supported("fp8"),
43
+ reason="FP8 is not supported on this GPU type.",
44
+)
45
+def test_modernbert_online_fp8_mteb(hf_runner, vllm_runner, monkeypatch) -> None:
46
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
47
+
48
+ def check_online_fp8(vllm_model) -> None:
49
+ vllm_model.apply_model(_assert_modernbert_online_fp8)
50
+
51
+ mteb_test_embed_models(
52
+ hf_runner,
53
+ vllm_runner,
54
+ MODEL_INFO,
55
+ vllm_extra_kwargs={"quantization": "fp8_per_tensor"},
56
+ # Account for Online FP8 kernel variance across GPU architectures.
57
+ atol=2e-3,
58
+ vllm_model_callback=check_online_fp8,
59
+ )
vllm/model_executor/models/modernbert.py
+33
-6
@@ -13,7 +13,11 @@ from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.attention import (
14
EncoderOnlyAttention,
15
)
16
-from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
16
+from vllm.model_executor.layers.linear import (
17
+ MergedColumnParallelLinear,
18
+ QKVParallelLinear,
19
+ RowParallelLinear,
20
+)
21
from vllm.model_executor.layers.pooler import DispatchPooler
22
from vllm.model_executor.layers.pooler.activations import LambdaPoolerActivation
23
from vllm.model_executor.layers.pooler.seqwise import (
@@ -22,6 +26,7 @@ from vllm.model_executor.layers.pooler.seqwise import (
26
get_seq_pooling_method,
27
)
28
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
29
+from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.model_executor.layers.rotary_embedding import get_rope
31
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
32
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -68,6 +73,7 @@ class ModernBertAttention(nn.Module):
73
layer_id: int | None = None,
74
prefix: str = "",
75
dtype: torch.dtype | None = None,
76
+ quant_config: QuantizationConfig | None = None,
77
):
78
super().__init__()
79
self.config = config
@@ -85,6 +91,7 @@ class ModernBertAttention(nn.Module):
91
self.head_dim,
92
self.num_heads,
93
bias=config.attention_bias,
94
+ quant_config=quant_config,
95
prefix=f"{prefix}.Wqkv",
96
)
97
@@ -128,6 +135,7 @@ class ModernBertAttention(nn.Module):
135
config.hidden_size,
136
config.hidden_size,
137
bias=config.attention_bias,
138
+ quant_config=quant_config,
139
prefix=f"{prefix}.Wo",
140
)
141
@@ -146,22 +154,33 @@ class ModernBertAttention(nn.Module):
154
155
156
class ModernBertMLP(nn.Module):
149
- def __init__(self, config: ModernBertConfig, prefix: str = ""):
157
+ def __init__(
158
+ self,
159
+ config: ModernBertConfig,
160
+ prefix: str = "",
161
+ quant_config: QuantizationConfig | None = None,
162
+ ):
163
super().__init__()
164
self.config = config
152
- self.Wi = nn.Linear(
153
- config.hidden_size, int(config.intermediate_size) * 2, bias=config.mlp_bias
165
+ self.Wi = MergedColumnParallelLinear(
166
+ config.hidden_size,
167
+ [int(config.intermediate_size)] * 2,
168
+ bias=config.mlp_bias,
169
+ quant_config=quant_config,
170
+ prefix=f"{prefix}.Wi",
171
)
172
self.act = ACT2FN[config.hidden_activation]
173
self.Wo = RowParallelLinear(
174
config.intermediate_size,
175
config.hidden_size,
176
bias=config.mlp_bias,
177
+ quant_config=quant_config,
178
prefix=f"{prefix}.Wo",
179
)
180
181
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
164
- input, gate = self.Wi(hidden_states).chunk(2, dim=-1)
182
+ input_gate, _ = self.Wi(hidden_states)
183
+ input, gate = input_gate.chunk(2, dim=-1)
184
return self.Wo(self.act(input) * gate)[0]
185
186
@@ -172,6 +191,7 @@ class ModernBertLayer(nn.Module):
191
prefix: str = "",
192
layer_id: int | None = None,
193
dtype: torch.dtype | None = None,
194
+ quant_config: QuantizationConfig | None = None,
195
):
196
super().__init__()
197
self.config = config
@@ -186,11 +206,16 @@ class ModernBertLayer(nn.Module):
206
layer_id=layer_id,
207
prefix=f"{prefix}.attn",
208
dtype=dtype,
209
+ quant_config=quant_config,
210
)
211
self.mlp_norm = nn.LayerNorm(
212
config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
213
)
193
- self.mlp = ModernBertMLP(config, prefix=f"{prefix}.mlp")
214
+ self.mlp = ModernBertMLP(
215
+ config,
216
+ prefix=f"{prefix}.mlp",
217
+ quant_config=quant_config,
218
+ )
219
220
def forward(
221
self,
@@ -211,6 +236,7 @@ class ModernBertEncoderLayer(nn.Module):
236
super().__init__()
237
config = vllm_config.model_config.hf_config
238
dtype = vllm_config.model_config.dtype
239
+ quant_config = vllm_config.quant_config
240
self.layers = nn.ModuleList(
241
[
242
ModernBertLayer(
@@ -218,6 +244,7 @@ class ModernBertEncoderLayer(nn.Module):
244
layer_id=layer_id,
245
prefix=f"{prefix}.layers.{layer_id}",
246
dtype=dtype,
247
+ quant_config=quant_config,
248
)
249
for layer_id in range(config.num_hidden_layers)
250
]