[Pooling] Preserve BERT-family heads for raw logits (#57664)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Taneem Ibrahim committed
Sep 25, 2026 at 18:10 UTC
6ec2f84e1a42cdca59779ffae20aa62c8beda993
2 files changed
+3
-6
vllm/model_executor/models/bert.py
+2
-3
@@ -25,7 +25,6 @@ from vllm.model_executor.layers.pooler import (
25
Pooler,
26
PoolingParamsUpdate,
27
)
28
-from vllm.model_executor.layers.pooler.activations import LambdaPoolerActivation
28
from vllm.model_executor.layers.pooler.seqwise import (
29
EmbeddingPoolerHead,
30
SequencePooler,
@@ -116,11 +115,11 @@ class BertPooler(SequencePooler):
115
)
116
self.act_fn = nn.Tanh()
117
118
+ # Keep the model's head layers even when output activation is disabled.
119
# Use lambdas so that weights are not registered under `self.head`
120
self.head = EmbeddingPoolerHead(
121
head_dtype=head_dtype,
122
- projector=lambda x: self.dense(x),
123
- activation=LambdaPoolerActivation(self.act_fn),
122
+ projector=lambda x: self.act_fn(self.dense(x)),
123
)
124
125
vllm/model_executor/models/modernbert.py
+1
-3
@@ -19,7 +19,6 @@ from vllm.model_executor.layers.linear import (
19
RowParallelLinear,
20
)
21
from vllm.model_executor.layers.pooler import DispatchPooler
22
-from vllm.model_executor.layers.pooler.activations import LambdaPoolerActivation
22
from vllm.model_executor.layers.pooler.seqwise import (
23
EmbeddingPoolerHead,
24
SequencePooler,
@@ -367,8 +366,7 @@ class ModernBertPooler(SequencePooler):
366
# Use lambdas so that weights are not registered under `self.head`
367
self.head = EmbeddingPoolerHead(
368
head_dtype=head_dtype,
370
- projector=lambda x: self.dense(x),
371
- activation=LambdaPoolerActivation(lambda x: self.norm(self.act(x))),
369
+ projector=lambda x: self.norm(self.act(self.dense(x))),
370
)
371
372