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 10, 06:30 UTC
[Model] Standardize pooling heads (#32148)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Jan 13, 2026 at 01:01 UTC
8863c2b25c84573e968ba291a26648ad17778c43
9 files changed
+182
-149
vllm/model_executor/layers/pooler/common.py
+6
-1
@@ -2,12 +2,17 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Callable
4
from dataclasses import dataclass
5
+from typing import TypeVar
6
7
import torch
8
9
from vllm.pooling_params import PoolingParams
10
11
+_T = TypeVar("_T", bound=torch.Tensor | list[torch.Tensor])
12
+
13
+ProjectorFn = Callable[[torch.Tensor], torch.Tensor]
14
ClassifierFn = Callable[[torch.Tensor], torch.Tensor]
15
+ActivationFn = Callable[[_T], _T]
16
17
18
@dataclass(frozen=True)
@@ -24,4 +29,4 @@ class PoolingParamsUpdate:
29
params.requires_token_ids = self.requires_token_ids
30
31
27
-__all__ = ["ClassifierFn", "PoolingParamsUpdate"]
32
+__all__ = ["ActivationFn", "ClassifierFn", "ProjectorFn", "PoolingParamsUpdate"]
vllm/model_executor/layers/pooler/seqwise/heads.py
+41
-47
@@ -7,14 +7,7 @@ from typing import TypeAlias
7
import torch
8
import torch.nn as nn
9
10
-from vllm.config import get_current_vllm_config
11
-from vllm.model_executor.layers.pooler import ClassifierFn
12
-from vllm.model_executor.layers.pooler.activations import (
13
- PoolerActivation,
14
- PoolerNormalize,
15
- resolve_classifier_act_fn,
16
-)
17
-from vllm.model_executor.models.adapters import _load_st_projector
10
+from vllm.model_executor.layers.pooler import ActivationFn, ClassifierFn, ProjectorFn
11
from vllm.tasks import PoolingTask
12
from vllm.v1.pool.metadata import PoolingMetadata
13
@@ -38,17 +31,17 @@ class SequencePoolerHead(nn.Module, ABC):
31
32
33
class EmbeddingPoolerHead(SequencePoolerHead):
41
- def __init__(self) -> None:
34
+ def __init__(
35
+ self,
36
+ projector: ProjectorFn | None = None,
37
+ head_dtype: torch.dtype | str | None = None,
38
+ activation: ActivationFn | None = None,
39
+ ) -> None:
40
super().__init__()
41
44
- # Load ST projector if available
45
- vllm_config = get_current_vllm_config()
46
- model_config = vllm_config.model_config
47
-
48
- self.projector = _load_st_projector(model_config)
49
- self.head_dtype = model_config.head_dtype
50
-
51
- self.activation = PoolerNormalize()
42
+ self.projector = projector
43
+ self.head_dtype = head_dtype
44
+ self.activation = activation
45
46
def get_supported_tasks(self) -> Set[PoolingTask]:
47
return {"embed"}
@@ -65,7 +58,8 @@ class EmbeddingPoolerHead(SequencePoolerHead):
58
pooled_data = torch.stack(pooled_data)
59
# pooled_data shape: [batchsize, hidden_dimension]
60
68
- pooled_data = pooled_data.to(self.head_dtype)
61
+ if self.head_dtype is not None:
62
+ pooled_data = pooled_data.to(self.head_dtype)
63
64
# Apply ST projector
65
if self.projector is not None:
@@ -88,15 +82,16 @@ class EmbeddingPoolerHead(SequencePoolerHead):
82
]
83
84
# for normalize
91
- flags = [p.normalize for p in pooling_params]
92
- if len(set(flags)) == 1:
93
- if flags[0]:
94
- pooled_data = self.activation(pooled_data)
95
- else:
96
- pooled_data = [
97
- self.activation(vecs) if f else vecs
98
- for vecs, f in zip(pooled_data, flags)
99
- ]
85
+ if self.activation is not None:
86
+ flags = [p.normalize for p in pooling_params]
87
+ if len(set(flags)) == 1:
88
+ if flags[0]:
89
+ pooled_data = self.activation(pooled_data)
90
+ else:
91
+ pooled_data = [
92
+ self.activation(vecs) if f else vecs
93
+ for vecs, f in zip(pooled_data, flags)
94
+ ]
95
96
# pooled_data shape: [batchsize, embedding_dimension]
97
return pooled_data
@@ -106,20 +101,16 @@ class ClassifierPoolerHead(SequencePoolerHead):
101
def __init__(
102
self,
103
classifier: ClassifierFn | None = None,
109
- act_fn: PoolerActivation | str | None = None,
104
+ logit_bias: float | None = None,
105
+ head_dtype: torch.dtype | str | None = None,
106
+ activation: ActivationFn | None = None,
107
) -> None:
108
super().__init__()
109
113
- vllm_config = get_current_vllm_config()
114
- model_config = vllm_config.model_config
115
-
110
self.classifier = classifier
117
- self.logit_bias: float | None = model_config.pooler_config.logit_bias
118
- self.head_dtype = model_config.head_dtype
119
-
120
- self.act_fn = resolve_classifier_act_fn(
121
- model_config, static_num_labels=True, act_fn=act_fn
122
- )
111
+ self.logit_bias = logit_bias
112
+ self.head_dtype = head_dtype
113
+ self.activation = activation
114
115
def get_supported_tasks(self) -> Set[PoolingTask]:
116
return {"classify", "score"}
@@ -136,7 +127,8 @@ class ClassifierPoolerHead(SequencePoolerHead):
127
pooled_data = torch.stack(pooled_data)
128
# pooled_data shape: [batchsize, hidden_size]
129
139
- pooled_data = pooled_data.to(self.head_dtype)
130
+ if self.head_dtype is not None:
131
+ pooled_data = pooled_data.to(self.head_dtype)
132
133
if self.classifier is not None:
134
pooled_data = self.classifier(pooled_data)
@@ -145,13 +137,15 @@ class ClassifierPoolerHead(SequencePoolerHead):
137
if self.logit_bias is not None:
138
pooled_data -= self.logit_bias
139
148
- flags = [p.use_activation for p in pooling_params]
149
- if len(set(flags)) == 1:
150
- scores = self.act_fn(pooled_data) if flags[0] else pooled_data
151
- else:
152
- scores = [
153
- self.act_fn(vecs) if f else vecs for vecs, f in zip(pooled_data, flags)
154
- ]
140
+ if self.activation is not None:
141
+ flags = [p.use_activation for p in pooling_params]
142
+ if len(set(flags)) == 1:
143
+ pooled_data = self.activation(pooled_data) if flags[0] else pooled_data
144
+ else:
145
+ pooled_data = [
146
+ self.activation(vecs) if f else vecs
147
+ for vecs, f in zip(pooled_data, flags)
148
+ ]
149
156
- # scores shape: [batchsize, num_labels]
157
- return scores
150
+ # pooled_data shape: [batchsize, num_labels]
151
+ return pooled_data
vllm/model_executor/layers/pooler/seqwise/poolers.py
+25
-4
@@ -5,10 +5,15 @@ from typing import TypeAlias
5
6
import torch
7
8
-from vllm.config import PoolerConfig
8
+from vllm.config import PoolerConfig, get_current_vllm_config
9
from vllm.model_executor.layers.pooler import ClassifierFn, PoolingParamsUpdate
10
from vllm.model_executor.layers.pooler.abstract import Pooler
11
-from vllm.model_executor.layers.pooler.activations import PoolerActivation
11
+from vllm.model_executor.layers.pooler.activations import (
12
+ PoolerActivation,
13
+ PoolerNormalize,
14
+ resolve_classifier_act_fn,
15
+)
16
+from vllm.model_executor.models.adapters import _load_st_projector
17
from vllm.tasks import POOLING_TASKS, PoolingTask
18
from vllm.v1.pool.metadata import PoolingMetadata
19
@@ -86,7 +91,14 @@ class SequencePooler(Pooler):
91
92
def pooler_for_embed(pooler_config: PoolerConfig):
93
pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
89
- head = EmbeddingPoolerHead()
94
+
95
+ vllm_config = get_current_vllm_config()
96
+ model_config = vllm_config.model_config
97
+ head = EmbeddingPoolerHead(
98
+ projector=_load_st_projector(model_config),
99
+ head_dtype=model_config.head_dtype,
100
+ activation=PoolerNormalize(),
101
+ )
102
103
return SequencePooler(pooling=pooling, head=head)
104
@@ -101,6 +113,15 @@ def pooler_for_classify(
113
if pooling is None:
114
pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
115
104
- head = ClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
116
+ vllm_config = get_current_vllm_config()
117
+ model_config = vllm_config.model_config
118
+ head = ClassifierPoolerHead(
119
+ classifier=classifier,
120
+ logit_bias=model_config.pooler_config.logit_bias,
121
+ head_dtype=model_config.head_dtype,
122
+ activation=resolve_classifier_act_fn(
123
+ model_config, static_num_labels=True, act_fn=act_fn
124
+ ),
125
+ )
126
127
return SequencePooler(pooling=pooling, head=head)
vllm/model_executor/layers/pooler/tokwise/heads.py
+23
-32
@@ -7,14 +7,7 @@ from typing import TypeAlias
7
import torch
8
import torch.nn as nn
9
10
-from vllm.config import get_current_vllm_config
11
-from vllm.model_executor.layers.pooler import ClassifierFn
12
-from vllm.model_executor.layers.pooler.activations import (
13
- PoolerActivation,
14
- PoolerNormalize,
15
- resolve_classifier_act_fn,
16
-)
17
-from vllm.model_executor.models.adapters import _load_st_projector
10
+from vllm.model_executor.layers.pooler import ActivationFn, ClassifierFn, ProjectorFn
11
from vllm.pooling_params import PoolingParams
12
from vllm.tasks import PoolingTask
13
from vllm.v1.pool.metadata import PoolingMetadata
@@ -49,17 +42,17 @@ class TokenPoolerHead(nn.Module, ABC):
42
43
44
class TokenEmbeddingPoolerHead(TokenPoolerHead):
52
- def __init__(self) -> None:
45
+ def __init__(
46
+ self,
47
+ projector: ProjectorFn | None = None,
48
+ head_dtype: torch.dtype | str | None = None,
49
+ activation: ActivationFn | None = None,
50
+ ) -> None:
51
super().__init__()
52
55
- # Load ST projector if available
56
- vllm_config = get_current_vllm_config()
57
- model_config = vllm_config.model_config
58
-
59
- self.projector = _load_st_projector(model_config)
60
- self.head_dtype = model_config.head_dtype
61
-
62
- self.activation = PoolerNormalize()
53
+ self.projector = projector
54
+ self.head_dtype = head_dtype
55
+ self.activation = activation
56
57
def get_supported_tasks(self) -> Set[PoolingTask]:
58
return {"token_embed"}
@@ -73,7 +66,8 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
66
if pooled_data is None:
67
return None
68
76
- pooled_data = pooled_data.to(self.head_dtype)
69
+ if self.head_dtype is not None:
70
+ pooled_data = pooled_data.to(self.head_dtype)
71
# pooled_data shape: [n_tokens, hidden_dimension]
72
73
# Apply ST projector
@@ -85,7 +79,7 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
79
pooled_data = pooled_data[..., : pooling_param.dimensions]
80
81
# for normalize
88
- if pooling_param.normalize:
82
+ if self.activation is not None and pooling_param.normalize:
83
pooled_data = self.activation(pooled_data)
84
85
# pooled_data shape: [n_tokens, embedding_dimension]
@@ -96,20 +90,16 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
90
def __init__(
91
self,
92
classifier: ClassifierFn | None = None,
99
- act_fn: PoolerActivation | str | None = None,
93
+ logit_bias: float | None = None,
94
+ head_dtype: torch.dtype | str | None = None,
95
+ activation: ActivationFn | None = None,
96
) -> None:
97
super().__init__()
98
103
- vllm_config = get_current_vllm_config()
104
- model_config = vllm_config.model_config
105
-
99
self.classifier = classifier
107
- self.logit_bias: float | None = model_config.pooler_config.logit_bias
108
- self.head_dtype = model_config.head_dtype
109
-
110
- self.act_fn = resolve_classifier_act_fn(
111
- model_config, static_num_labels=False, act_fn=act_fn
112
- )
100
+ self.logit_bias = logit_bias
101
+ self.head_dtype = head_dtype
102
+ self.activation = activation
103
104
def get_supported_tasks(self) -> Set[PoolingTask]:
105
return {"token_classify"}
@@ -123,7 +113,8 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
113
if pooled_data is None:
114
return None
115
126
- pooled_data = pooled_data.to(self.head_dtype)
116
+ if self.head_dtype is not None:
117
+ pooled_data = pooled_data.to(self.head_dtype)
118
# hidden_states shape: [n_token, hidden_size]
119
120
if self.classifier is not None:
@@ -135,8 +126,8 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
126
if self.logit_bias is not None:
127
scores -= self.logit_bias
128
138
- if pooling_param.use_activation:
139
- scores = self.act_fn(scores)
129
+ if self.activation is not None and pooling_param.use_activation:
130
+ scores = self.activation(scores)
131
132
# scores shape: [n_token, num_labels]
133
return scores
vllm/model_executor/layers/pooler/tokwise/poolers.py
+25
-4
@@ -5,10 +5,15 @@ from typing import TypeAlias
5
6
import torch
7
8
-from vllm.config import PoolerConfig
8
+from vllm.config import PoolerConfig, get_current_vllm_config
9
from vllm.model_executor.layers.pooler import ClassifierFn, PoolingParamsUpdate
10
from vllm.model_executor.layers.pooler.abstract import Pooler
11
-from vllm.model_executor.layers.pooler.activations import PoolerActivation
11
+from vllm.model_executor.layers.pooler.activations import (
12
+ PoolerActivation,
13
+ PoolerNormalize,
14
+ resolve_classifier_act_fn,
15
+)
16
+from vllm.model_executor.models.adapters import _load_st_projector
17
from vllm.tasks import POOLING_TASKS, PoolingTask
18
from vllm.v1.pool.metadata import PoolingMetadata
19
@@ -86,7 +91,14 @@ class TokenPooler(Pooler):
91
92
def pooler_for_token_embed(pooler_config: PoolerConfig):
93
pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())
89
- head = TokenEmbeddingPoolerHead()
94
+
95
+ vllm_config = get_current_vllm_config()
96
+ model_config = vllm_config.model_config
97
+ head = TokenEmbeddingPoolerHead(
98
+ projector=_load_st_projector(model_config),
99
+ head_dtype=model_config.head_dtype,
100
+ activation=PoolerNormalize(),
101
+ )
102
103
return TokenPooler(pooling=pooling, head=head)
104
@@ -101,6 +113,15 @@ def pooler_for_token_classify(
113
if pooling is None:
114
pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())
115
104
- head = TokenClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
116
+ vllm_config = get_current_vllm_config()
117
+ model_config = vllm_config.model_config
118
+ head = TokenClassifierPoolerHead(
119
+ classifier=classifier,
120
+ logit_bias=model_config.pooler_config.logit_bias,
121
+ head_dtype=model_config.head_dtype,
122
+ activation=resolve_classifier_act_fn(
123
+ model_config, static_num_labels=False, act_fn=act_fn
124
+ ),
125
+ )
126
127
return TokenPooler(pooling=pooling, head=head)
vllm/model_executor/models/bert.py
+25
-24
@@ -8,7 +8,7 @@ from torch import nn
8
from transformers import BertConfig
9
10
from vllm.compilation.decorators import support_torch_compile
11
-from vllm.config import CacheConfig, PoolerConfig, VllmConfig
11
+from vllm.config import CacheConfig, ModelConfig, PoolerConfig, VllmConfig
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.activation import get_act_fn
14
from vllm.model_executor.layers.attention.encoder_only_attention import (
@@ -24,11 +24,11 @@ from vllm.model_executor.layers.pooler import (
24
Pooler,
25
PoolingParamsUpdate,
26
)
27
+from vllm.model_executor.layers.pooler.activations import LambdaPoolerActivation
28
from vllm.model_executor.layers.pooler.seqwise import (
29
+ EmbeddingPoolerHead,
30
SequencePooler,
29
- SequencePoolerHeadOutput,
31
SequencePoolerOutput,
31
- SequencePoolingMethodOutput,
32
get_seq_pooling_method,
33
)
34
from vllm.model_executor.layers.pooler.tokwise import (
@@ -94,26 +94,32 @@ class BertEmbedding(nn.Module):
94
95
96
class BertPooler(SequencePooler):
97
- def __init__(self, config: BertConfig, pooler_config: PoolerConfig):
97
+ def __init__(self, model_config: ModelConfig):
98
+ pooler_config = model_config.pooler_config
99
+ assert pooler_config is not None
100
+
101
+ config: BertConfig = model_config.hf_config
102
+
103
super().__init__(
104
pooling=get_seq_pooling_method(pooler_config.seq_pooling_type),
100
- head=self.head,
105
+ # We set this dummy to avoid adding parameters to nn.Module too early
106
+ head=nn.Identity(),
107
)
108
103
- self.dense = nn.Linear(config.hidden_size, config.hidden_size)
104
- self.activation = nn.Tanh()
105
-
106
- def head(
107
- self,
108
- pooled_data: SequencePoolingMethodOutput,
109
- pooling_metadata: PoolingMetadata,
110
- ) -> SequencePoolerHeadOutput:
111
- if isinstance(pooled_data, list):
112
- pooled_data = torch.stack(pooled_data)
109
+ head_dtype = model_config.head_dtype
110
+ self.dense = nn.Linear(
111
+ config.hidden_size,
112
+ config.hidden_size,
113
+ dtype=head_dtype,
114
+ )
115
+ self.act_fn = nn.Tanh()
116
114
- pooled_data = self.dense(pooled_data)
115
- pooled_data = self.activation(pooled_data)
116
- return pooled_data
117
+ # Use lambdas so that weights are not registered under `self.head`
118
+ self.head = EmbeddingPoolerHead(
119
+ projector=lambda x: self.dense(x),
120
+ head_dtype=head_dtype,
121
+ activation=LambdaPoolerActivation(self.act_fn),
122
+ )
123
124
125
class BertEncoder(nn.Module):
@@ -449,12 +455,7 @@ class BertPoolingModel(BertModel):
455
embedding_class=embedding_class,
456
)
457
452
- config = vllm_config.model_config.hf_config
453
-
454
- pooler_config = vllm_config.model_config.pooler_config
455
- assert pooler_config is not None
456
-
457
- self.pooler = BertPooler(config, pooler_config)
458
+ self.pooler = BertPooler(vllm_config.model_config)
459
460
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
461
other_weights, loaded_stacked_params = self._load_weights(weights)
vllm/model_executor/models/bert_with_rope.py
+1
-4
@@ -466,10 +466,7 @@ class BertWithRope(nn.Module, SupportsQuant):
466
)
467
468
if add_pooling_layer:
469
- pooler_config = vllm_config.model_config.pooler_config
470
- assert pooler_config is not None
471
-
472
- self.pooler = BertPooler(self.config, pooler_config)
469
+ self.pooler = BertPooler(vllm_config.model_config)
470
else:
471
self.pooler = None
472
vllm/model_executor/models/gritlm.py
+11
-14
@@ -5,7 +5,7 @@ from collections.abc import Set
5
import numpy as np
6
import torch
7
8
-from vllm.config import ModelConfig, PoolerConfig, VllmConfig
8
+from vllm.config import ModelConfig, VllmConfig
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.pooler import (
11
DispatchPooler,
@@ -13,8 +13,8 @@ from vllm.model_executor.layers.pooler import (
13
)
14
from vllm.model_executor.layers.pooler.activations import PoolerNormalize
15
from vllm.model_executor.layers.pooler.seqwise import (
16
+ EmbeddingPoolerHead,
17
SequencePooler,
17
- SequencePoolerHeadOutput,
18
SequencePoolingMethod,
19
SequencePoolingMethodOutput,
20
get_seq_pooling_method,
@@ -178,25 +178,22 @@ class GritLMMeanPool(SequencePoolingMethod):
178
179
180
class GritLMPooler(SequencePooler):
181
- def __init__(self, model_config: ModelConfig, pooler_config: PoolerConfig):
181
+ def __init__(self, model_config: ModelConfig):
182
+ pooler_config = model_config.pooler_config
183
+ assert pooler_config is not None
184
+
185
super().__init__(
186
pooling=(
187
GritLMMeanPool(model_config)
188
if pooler_config.seq_pooling_type == "MEAN"
189
else get_seq_pooling_method(pooler_config.seq_pooling_type)
190
),
188
- head=self.head,
191
+ head=EmbeddingPoolerHead(
192
+ head_dtype=model_config.head_dtype,
193
+ activation=PoolerNormalize(),
194
+ ),
195
)
196
191
- self.activation = PoolerNormalize()
192
-
193
- def head(
194
- self,
195
- pooled_data: SequencePoolingMethodOutput,
196
- pooling_metadata: PoolingMetadata,
197
- ) -> SequencePoolerHeadOutput:
198
- return self.activation(pooled_data)
199
-
197
198
@default_pooling_type(seq_pooling_type="MEAN")
199
class GritLM(LlamaForCausalLM):
@@ -240,6 +237,6 @@ class GritLM(LlamaForCausalLM):
237
self.pooler = DispatchPooler(
238
{
239
"token_embed": pooler_for_token_embed(pooler_config),
243
- "embed": GritLMPooler(vllm_config.model_config, pooler_config),
240
+ "embed": GritLMPooler(vllm_config.model_config),
241
}
242
)
vllm/model_executor/models/modernbert.py
+25
-19
@@ -8,17 +8,17 @@ from transformers import ModernBertConfig
8
from transformers.activations import ACT2FN
9
10
from vllm.compilation.decorators import support_torch_compile
11
-from vllm.config import PoolerConfig, VllmConfig
11
+from vllm.config import ModelConfig, VllmConfig
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.attention.encoder_only_attention import (
14
EncoderOnlyAttention,
15
)
16
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
17
from vllm.model_executor.layers.pooler import DispatchPooler
18
+from vllm.model_executor.layers.pooler.activations import LambdaPoolerActivation
19
from vllm.model_executor.layers.pooler.seqwise import (
20
+ EmbeddingPoolerHead,
21
SequencePooler,
20
- SequencePoolerHeadOutput,
21
- SequencePoolingMethodOutput,
22
get_seq_pooling_method,
23
)
24
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
@@ -26,7 +26,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
26
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
27
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28
from vllm.sequence import IntermediateTensors
29
-from vllm.v1.pool.metadata import PoolingMetadata
29
30
from .interfaces import SupportsCrossEncoding
31
from .interfaces_base import attn_type, default_pooling_type
@@ -282,7 +281,11 @@ class ModernBertModel(nn.Module):
281
282
283
class ModernBertPooler(SequencePooler):
285
- def __init__(self, config: ModernBertConfig, pooler_config: PoolerConfig):
284
+ def __init__(self, model_config: ModelConfig):
285
+ pooler_config = model_config.pooler_config
286
+ assert pooler_config is not None
287
+
288
+ config: ModernBertConfig = model_config.hf_config
289
hf_pooling_type = config.classifier_pooling.upper()
290
# vllm_pooling_type = pooler_config.seq_pooling_type
291
# Currently we don't have a way to see if the user set the pooling type
@@ -290,27 +293,30 @@ class ModernBertPooler(SequencePooler):
293
294
super().__init__(
295
pooling=get_seq_pooling_method(hf_pooling_type),
293
- head=self.head,
296
+ # We set this dummy to avoid adding parameters to nn.Module too early
297
+ head=nn.Identity(),
298
)
299
300
+ head_dtype = model_config.head_dtype
301
self.dense = nn.Linear(
297
- config.hidden_size, config.hidden_size, config.classifier_bias
302
+ config.hidden_size,
303
+ config.hidden_size,
304
+ config.classifier_bias,
305
+ dtype=head_dtype,
306
)
307
self.act = nn.GELU()
308
self.norm = nn.LayerNorm(
301
- config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
309
+ config.hidden_size,
310
+ eps=config.norm_eps,
311
+ bias=config.norm_bias,
312
)
313
304
- def head(
305
- self,
306
- pooled_data: SequencePoolingMethodOutput,
307
- pooling_metadata: PoolingMetadata,
308
- ) -> SequencePoolerHeadOutput:
309
- if isinstance(pooled_data, list):
310
- pooled_data = torch.stack(pooled_data)
311
-
312
- pooled_data = pooled_data.to(self.dense.weight.dtype)
313
- return self.norm(self.act(self.dense(pooled_data)))
314
+ # Use lambdas so that weights are not registered under `self.head`
315
+ self.head = EmbeddingPoolerHead(
316
+ projector=lambda x: self.dense(x),
317
+ head_dtype=head_dtype,
318
+ activation=LambdaPoolerActivation(lambda x: self.norm(self.act(x))),
319
+ )
320
321
322
@default_pooling_type(seq_pooling_type="CLS")
@@ -335,7 +341,7 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
341
pooler_config = vllm_config.model_config.pooler_config
342
assert pooler_config is not None
343
338
- self.pooling = ModernBertPooler(config, pooler_config)
344
+ self.pooling = ModernBertPooler(vllm_config.model_config)
345
346
self.pooler = DispatchPooler.for_seq_cls(
347
pooler_config,