diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index 85ef42675b..7348ec9383 100644 --- a/vllm/model_executor/models/bert.py +++ b/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)), ) diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index df901a6566..3a8abee67e 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/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))), )