A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][CI] Stabilize ModernBERT token-classification parity against Hugging Face (#44040)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 9, 2026 at 10:52 UTC
b6971198008e37637b80b0cdc183bf4bad11376e
1 file changed
+21
-6
vllm/model_executor/models/modernbert.py
+21
-6
@@ -63,7 +63,11 @@ class ModernBertEmbeddings(nn.Module):
63
64
class ModernBertAttention(nn.Module):
65
def __init__(
66
- self, config: ModernBertConfig, layer_id: int | None = None, prefix: str = ""
66
+ self,
67
+ config: ModernBertConfig,
68
+ layer_id: int | None = None,
69
+ prefix: str = "",
70
+ dtype: torch.dtype | None = None,
71
):
72
super().__init__()
73
self.config = config
@@ -90,11 +94,13 @@ class ModernBertAttention(nn.Module):
94
rope_parameters = config.rope_parameters[layer_type]
95
sliding_window: int | None = None
96
if layer_type == "sliding_attention":
93
- sliding_window = config.local_attention // 2
97
+ # Treats the local attention boundary as inclusive
98
+ sliding_window = config.sliding_window + 1
99
else:
100
# Transformers v4
101
sliding_window = None
102
if layer_id % config.global_attn_every_n_layers != 0:
103
+ # ModernBertConfig does not expose sliding_window
104
sliding_window = config.local_attention // 2
105
rope_theta = (
106
config.local_rope_theta
@@ -109,7 +115,7 @@ class ModernBertAttention(nn.Module):
115
head_size=self.head_dim,
116
max_position=config.max_position_embeddings,
117
rope_parameters=rope_parameters,
112
- dtype=torch.float16,
118
+ dtype=dtype,
119
)
120
self.attn = EncoderOnlyAttention(
121
self.num_heads,
@@ -146,7 +152,7 @@ class ModernBertMLP(nn.Module):
152
self.Wi = nn.Linear(
153
config.hidden_size, int(config.intermediate_size) * 2, bias=config.mlp_bias
154
)
149
- self.act = nn.GELU()
155
+ self.act = ACT2FN[config.hidden_activation]
156
self.Wo = RowParallelLinear(
157
config.intermediate_size,
158
config.hidden_size,
@@ -161,7 +167,11 @@ class ModernBertMLP(nn.Module):
167
168
class ModernBertLayer(nn.Module):
169
def __init__(
164
- self, config: ModernBertConfig, prefix: str = "", layer_id: int | None = None
170
+ self,
171
+ config: ModernBertConfig,
172
+ prefix: str = "",
173
+ layer_id: int | None = None,
174
+ dtype: torch.dtype | None = None,
175
):
176
super().__init__()
177
self.config = config
@@ -172,7 +182,10 @@ class ModernBertLayer(nn.Module):
182
config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
183
)
184
self.attn = ModernBertAttention(
175
- config=config, layer_id=layer_id, prefix=f"{prefix}.attn"
185
+ config=config,
186
+ layer_id=layer_id,
187
+ prefix=f"{prefix}.attn",
188
+ dtype=dtype,
189
)
190
self.mlp_norm = nn.LayerNorm(
191
config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
@@ -197,12 +210,14 @@ class ModernBertEncoderLayer(nn.Module):
210
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
211
super().__init__()
212
config = vllm_config.model_config.hf_config
213
+ dtype = vllm_config.model_config.dtype
214
self.layers = nn.ModuleList(
215
[
216
ModernBertLayer(
217
config=config,
218
layer_id=layer_id,
219
prefix=f"{prefix}.layers.{layer_id}",
220
+ dtype=dtype,
221
)
222
for layer_id in range(config.num_hidden_layers)
223
]