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, 04:45 UTC
[Fix] Support cls pooling in ModernBertPooler (#20067)
Signed-off-by: shengzhe.li <shengzhe.li@sbintuitions.co.jp>
lsz05 committed
Jun 26, 2025 at 04:07 UTC
23a04e08958e75b5528dae310c1c5774dd03551d
1 file changed
+8
-1
vllm/model_executor/models/modernbert.py
+8
-1
@@ -258,6 +258,7 @@ class ModernBertPooler(nn.Module):
258
super().__init__()
259
self.dense = nn.Linear(config.hidden_size, config.hidden_size,
260
config.classifier_bias)
261
+ self.pooling_type = config.classifier_pooling
262
self.act = nn.GELU()
263
self.norm = nn.LayerNorm(config.hidden_size,
264
eps=config.norm_eps,
@@ -265,7 +266,13 @@ class ModernBertPooler(nn.Module):
266
267
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
268
pooled_output = hidden_states
268
- pooled_output = pooled_output.mean(dim=0, keepdim=False)
269
+ if self.pooling_type == "mean":
270
+ pooled_output = pooled_output.mean(dim=0, keepdim=False)
271
+ elif self.pooling_type == "cls":
272
+ pooled_output = pooled_output[0, :]
273
+ else:
274
+ raise ValueError("Pooling type should be either `cls` or `mean`, "
275
+ f"but got {self.pooling_type}")
276
pooled_output = self.norm(self.act(self.dense(pooled_output)))
277
return pooled_output
278