[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>

vllm/model_executor/models/bert.py

@@ -25,7 +25,6 @@ from vllm.model_executor.layers.pooler import (
     Pooler,
     PoolingParamsUpdate,
 )
-from vllm.model_executor.layers.pooler.activations import LambdaPoolerActivation
 from vllm.model_executor.layers.pooler.seqwise import (
     EmbeddingPoolerHead,
     SequencePooler,
@@ -116,11 +115,11 @@ class BertPooler(SequencePooler):
         )
         self.act_fn = nn.Tanh()
 
+        # Keep the model's head layers even when output activation is disabled.
         # Use lambdas so that weights are not registered under `self.head`
         self.head = EmbeddingPoolerHead(
             head_dtype=head_dtype,
-            projector=lambda x: self.dense(x),
-            activation=LambdaPoolerActivation(self.act_fn),
+            projector=lambda x: self.act_fn(self.dense(x)),
         )
 
 

vllm/model_executor/models/modernbert.py

@@ -19,7 +19,6 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.pooler import DispatchPooler
-from vllm.model_executor.layers.pooler.activations import LambdaPoolerActivation
 from vllm.model_executor.layers.pooler.seqwise import (
     EmbeddingPoolerHead,
     SequencePooler,
@@ -367,8 +366,7 @@ class ModernBertPooler(SequencePooler):
         # Use lambdas so that weights are not registered under `self.head`
         self.head = EmbeddingPoolerHead(
             head_dtype=head_dtype,
-            projector=lambda x: self.dense(x),
-            activation=LambdaPoolerActivation(lambda x: self.norm(self.act(x))),
+            projector=lambda x: self.norm(self.act(self.dense(x))),
         )