@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1c3198b6c4
Fork 0

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, 01:00 UTC

[Model] Consolidate pooler implementations (#20927)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jul 16, 2025 at 21:39 UTC 1c3198b6c4e534f10b97a8f5e0d958f4bb634bc6
9 files changed +553 -367
vllm/model_executor/layers/pooler.py
+434 -247
@@ -1,22 +1,21 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 +from abc import ABC, abstractmethod
4 +from dataclasses import dataclass
5 from enum import IntEnum
5 -from typing import Optional, Union
6 +from typing import Callable, Optional, TypeVar, Union
7
8 import torch
9 import torch.nn as nn
10 import torch.nn.functional as F
10 -from typing_extensions import assert_never
11 +from transformers import PretrainedConfig
12
13 from vllm.config import ModelConfig, PoolerConfig
14 from vllm.model_executor.pooling_metadata import ( # noqa: E501
15 PoolingMetadata as V0PoolingMetadata)
16 from vllm.model_executor.pooling_metadata import PoolingTensors
17 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
17 -from vllm.transformers_utils.config import (
18 - get_classification_activation_function,
19 - get_cross_encoder_activation_function)
18 +from vllm.utils import resolve_obj_by_qualname
19 from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata
20
21 PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]
@@ -31,140 +30,202 @@ class PoolingType(IntEnum):
30 MEAN = 4
31
32
34 -class SimplePooler(nn.Module):
35 - """A layer that pools specific information from hidden states.
33 +@dataclass(frozen=True)
34 +class ResolvedPoolingConfig:
35 + pooling_type: PoolingType
36
37 - This layer does the following:
38 - 1. Extracts specific tokens or aggregates data based on pooling method.
39 - 2. Normalizes output if specified.
40 - 3. Returns structured results as `PoolerOutput`.
41 -
42 - Attributes:
43 - pooling_type: The type of pooling to use.
44 - normalize: Whether to normalize the pooled data.
45 - """
37 + normalize: bool
38 + softmax: bool
39 + step_tag_id: Optional[int]
40 + returned_token_ids: Optional[list[int]]
41
47 - @staticmethod
48 - def from_pooling_type(
42 + @classmethod
43 + def from_config_with_defaults(
44 + cls,
45 + pooler_config: PoolerConfig,
46 pooling_type: PoolingType,
50 - *,
47 normalize: bool,
48 softmax: bool,
49 step_tag_id: Optional[int] = None,
50 returned_token_ids: Optional[list[int]] = None,
55 - ) -> "SimplePooler":
56 - if pooling_type == PoolingType.LAST:
57 - assert step_tag_id is None and returned_token_ids is None
58 - return LastPool(normalize=normalize, softmax=softmax)
59 - if pooling_type == PoolingType.ALL:
60 - assert step_tag_id is None and returned_token_ids is None
61 - return AllPool(normalize=normalize, softmax=softmax)
62 - if pooling_type == PoolingType.CLS:
63 - assert step_tag_id is None and returned_token_ids is None
64 - return CLSPool(normalize=normalize, softmax=softmax)
65 - if pooling_type == PoolingType.MEAN:
66 - assert step_tag_id is None and returned_token_ids is None
67 - return MeanPool(normalize=normalize, softmax=softmax)
68 - if pooling_type == PoolingType.STEP:
69 - return StepPool(normalize=normalize,
70 - softmax=softmax,
71 - step_tag_id=step_tag_id,
72 - returned_token_ids=returned_token_ids)
51 + ) -> "ResolvedPoolingConfig":
52 + return cls(
53 + pooling_type=PoolingType[pooler_config.pooling_type]
54 + if pooler_config.pooling_type is not None else pooling_type,
55 + normalize=pooler_config.normalize
56 + if pooler_config.normalize is not None else normalize,
57 + softmax=pooler_config.softmax
58 + if pooler_config.softmax is not None else softmax,
59 + step_tag_id=pooler_config.step_tag_id
60 + if pooler_config.step_tag_id is not None else step_tag_id,
61 + returned_token_ids=pooler_config.returned_token_ids
62 + if pooler_config.returned_token_ids is not None else
63 + returned_token_ids,
64 + )
65
74 - assert_never(pooling_type)
66
76 - def __init__(self, *, normalize: bool, softmax: bool) -> None:
77 - super().__init__()
67 +def get_prompt_lens(
68 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
69 + pooling_metadata: PoolingMetadata,
70 +) -> torch.Tensor:
71 + if isinstance(pooling_metadata, V1PoolingMetadata):
72 + return pooling_metadata.prompt_lens
73 +
74 + assert isinstance(hidden_states, torch.Tensor)
75 + return PoolingTensors.from_pooling_metadata(
76 + pooling_metadata, hidden_states.device).prompt_lens
77 +
78 +
79 +def get_classification_activation_function(config: PretrainedConfig):
80 + return PoolerClassify()
81 +
82 +
83 +def get_cross_encoder_activation_function(config: PretrainedConfig):
84 + function_name: Optional[str] = None
85 + if (hasattr(config, "sentence_transformers")
86 + and "activation_fn" in config.sentence_transformers):
87 + function_name = config.sentence_transformers["activation_fn"]
88 + elif (hasattr(config, "sbert_ce_default_activation_function")
89 + and config.sbert_ce_default_activation_function is not None):
90 + function_name = config.sbert_ce_default_activation_function
91 +
92 + if function_name is not None:
93 + assert function_name.startswith("torch.nn.modules."), (
94 + "Loading of activation functions is restricted to "
95 + "torch.nn.modules for security reasons")
96 + fn = resolve_obj_by_qualname(function_name)()
97 + return PoolerActivation.wraps(fn)
98
79 - self.head = PoolerHead(normalize=normalize, softmax=softmax)
99 + return PoolerScore()
100
81 - def get_prompt_lens(
101 +
102 +def build_output(all_data: torch.Tensor) -> PoolerOutput:
103 + all_outputs = [PoolingSequenceGroupOutput(data) for data in all_data]
104 + return PoolerOutput(outputs=all_outputs)
105 +
106 +
107 +class BasePooler(nn.Module):
108 +
109 + @abstractmethod
110 + def forward(
111 self,
112 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
113 pooling_metadata: PoolingMetadata,
114 + ) -> PoolerOutput:
115 + raise NotImplementedError
116 +
117 +
118 +class PoolingMethod(nn.Module, ABC):
119 +
120 + @staticmethod
121 + def from_pooling_type(pooling_type: PoolingType) -> "PoolingMethod":
122 + if pooling_type == PoolingType.LAST:
123 + return LastPool()
124 + if pooling_type == PoolingType.ALL:
125 + return AllPool()
126 + if pooling_type == PoolingType.CLS:
127 + return CLSPool()
128 + if pooling_type == PoolingType.MEAN:
129 + return MeanPool()
130 +
131 + raise NotImplementedError(f"Unsupported method: {pooling_type}")
132 +
133 + @abstractmethod
134 + def forward_one(
135 + self,
136 + hidden_states: torch.Tensor,
137 + prompt_len: Optional[torch.Tensor] = None,
138 ) -> torch.Tensor:
86 - if isinstance(pooling_metadata, V1PoolingMetadata):
87 - return pooling_metadata.prompt_lens
88 - assert isinstance(hidden_states, torch.Tensor)
89 - return PoolingTensors.from_pooling_metadata(
90 - pooling_metadata, hidden_states.device).prompt_lens
139 + """
140 + Note:
141 + `prompt_len=None` means `prompt_len=len(hidden_states)`.
142 + """
143 + raise NotImplementedError
144
92 - def extract_states(
145 + @abstractmethod
146 + def forward_all(
147 self,
94 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
95 - pooling_metadata: PoolingMetadata,
148 + hidden_states: torch.Tensor,
149 + prompt_lens: torch.Tensor,
150 ) -> Union[list[torch.Tensor], torch.Tensor]:
151 raise NotImplementedError
152
99 - def build_output(self, data: torch.Tensor) -> PoolingSequenceGroupOutput:
100 - return PoolingSequenceGroupOutput(data)
101 -
153 def forward(
154 self,
155 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
156 pooling_metadata: PoolingMetadata,
106 - ) -> PoolerOutput:
107 - pooled_data = self.extract_states(hidden_states, pooling_metadata)
108 - pooled_data = self.head(pooled_data, pooling_metadata)
109 - pooled_outputs = [self.build_output(data) for data in pooled_data]
110 - return PoolerOutput(outputs=pooled_outputs)
157 + ) -> Union[list[torch.Tensor], torch.Tensor]:
158 + prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)
159 +
160 + if isinstance(hidden_states, list):
161 + return [
162 + self.forward_one(h, prompt_len)
163 + for h, prompt_len in zip(hidden_states, prompt_lens)
164 + ]
165
166 + return self.forward_all(hidden_states, prompt_lens)
167
113 -class CLSPool(SimplePooler):
168
115 - def extract_states(
169 +class CLSPool(PoolingMethod):
170 +
171 + def forward_one(
172 self,
117 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
118 - pooling_metadata: PoolingMetadata,
119 - ) -> Union[list[torch.Tensor], torch.Tensor]:
120 - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
173 + hidden_states: torch.Tensor,
174 + prompt_len: Optional[torch.Tensor] = None,
175 + ) -> torch.Tensor:
176 + assert prompt_len is None or prompt_len == hidden_states.shape[0], \
177 + "partial prefill not supported with CLS pooling"
178
122 - if isinstance(hidden_states, list):
123 - result = []
124 - for req_state, prompt_len in zip(hidden_states, prompt_lens):
125 - assert prompt_len == req_state.shape[0], \
126 - "partial prefill not supported with CLS pooling"
127 - result.append(req_state[0])
128 - return result
179 + return hidden_states[0]
180
181 + def forward_all(
182 + self,
183 + hidden_states: torch.Tensor,
184 + prompt_lens: torch.Tensor,
185 + ) -> Union[list[torch.Tensor], torch.Tensor]:
186 first_token_flat_indices = torch.zeros_like(prompt_lens)
187 first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1]
188 return hidden_states[first_token_flat_indices]
189
190
135 -class LastPool(SimplePooler):
191 +class LastPool(PoolingMethod):
192
137 - def extract_states(
193 + def forward_one(
194 self,
139 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
140 - pooling_metadata: PoolingMetadata,
141 - ) -> Union[list[torch.Tensor], torch.Tensor]:
142 - if isinstance(hidden_states, list):
143 - return [h[-1] for h in hidden_states]
144 -
145 - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
195 + hidden_states: torch.Tensor,
196 + prompt_len: Optional[torch.Tensor] = None,
197 + ) -> torch.Tensor:
198 + return hidden_states[-1]
199
200 + def forward_all(
201 + self,
202 + hidden_states: torch.Tensor,
203 + prompt_lens: torch.Tensor,
204 + ) -> Union[list[torch.Tensor], torch.Tensor]:
205 last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1
206 return hidden_states[last_token_flat_indices]
207
208
151 -class AllPool(SimplePooler):
209 +class AllPool(PoolingMethod):
210
153 - def extract_states(
211 + def forward_one(
212 self,
155 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
156 - pooling_metadata: PoolingMetadata,
157 - ) -> Union[list[torch.Tensor], torch.Tensor]:
158 - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
213 + hidden_states: torch.Tensor,
214 + prompt_len: Optional[torch.Tensor] = None,
215 + ) -> torch.Tensor:
216 + assert prompt_len is None or prompt_len == hidden_states.shape[0], \
217 + "partial prefill not supported with ALL pooling"
218
160 - if isinstance(hidden_states, list):
161 - for req_state, prompt_len in zip(hidden_states, prompt_lens):
162 - assert prompt_len == req_state.shape[0], \
163 - "partial prefill not supported with ALL pooling"
164 - return hidden_states
219 + return hidden_states
220
221 + def forward_all(
222 + self,
223 + hidden_states: torch.Tensor,
224 + prompt_lens: torch.Tensor,
225 + ) -> Union[list[torch.Tensor], torch.Tensor]:
226 offset = 0
227 pooled_data = list[torch.Tensor]()
228 +
229 for prompt_len in prompt_lens:
230 pooled_data.append(hidden_states[offset:offset + prompt_len])
231 offset += prompt_len
@@ -172,24 +233,23 @@ class AllPool(SimplePooler):
233 return pooled_data
234
235
175 -class MeanPool(SimplePooler):
236 +class MeanPool(PoolingMethod):
237
177 - def extract_states(
238 + def forward_one(
239 self,
179 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
180 - pooling_metadata: PoolingMetadata,
181 - ) -> Union[list[torch.Tensor], torch.Tensor]:
182 - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
240 + hidden_states: torch.Tensor,
241 + prompt_len: Optional[torch.Tensor] = None,
242 + ) -> torch.Tensor:
243 + assert prompt_len is None or prompt_len == hidden_states.shape[0], \
244 + "partial prefill not supported with MEAN pooling"
245
184 - if isinstance(hidden_states, list):
185 - result = []
186 - for req_state, prompt_len in zip(hidden_states, prompt_lens):
187 - assert prompt_len == req_state.shape[0], \
188 - "partial prefill not supported with mean pooling"
189 - result.append(torch.mean(req_state, dim=0,
190 - dtype=torch.float32))
191 - return result
246 + return hidden_states.mean(dim=0, dtype=torch.float32)
247
248 + def forward_all(
249 + self,
250 + hidden_states: torch.Tensor,
251 + prompt_lens: torch.Tensor,
252 + ) -> Union[list[torch.Tensor], torch.Tensor]:
253 # Use float32 for torch.cumsum in MeanPool,
254 # otherwise precision will be lost significantly.
255 cumsum = torch.cumsum(hidden_states, dim=0, dtype=torch.float32)
@@ -203,78 +263,127 @@ class MeanPool(SimplePooler):
263 hidden_states[start_indices]) / prompt_lens.unsqueeze(1)
264
265
206 -class StepPool(SimplePooler):
266 +_T = TypeVar("_T", torch.Tensor, list[torch.Tensor])
267
208 - def __init__(
209 - self,
210 - *,
211 - normalize: bool,
212 - softmax: bool,
213 - step_tag_id: Optional[int] = None,
214 - returned_token_ids: Optional[list[int]] = None,
215 - ):
216 - super().__init__(normalize=normalize, softmax=softmax)
268
218 - self.step_tag_id = step_tag_id
219 - self.returned_token_ids = returned_token_ids
269 +class BasePoolerActivation(nn.Module, ABC):
270
221 - def get_prompt_token_ids(
222 - self,
223 - pooling_metadata: PoolingMetadata,
224 - ) -> list[torch.Tensor]:
225 - if isinstance(pooling_metadata, V1PoolingMetadata):
226 - return [
227 - pooling_metadata.prompt_token_ids[i, :num]
228 - for i, num in enumerate(pooling_metadata.prompt_lens)
229 - ]
230 - return [
231 - torch.tensor(seq_data_i.prompt_token_ids)
232 - for seq_data_i in pooling_metadata.seq_data.values()
233 - ]
271 + @abstractmethod
272 + def forward(self, pooled_data: _T) -> _T:
273 + # shape:
274 + # classify (& score) -> (batch_size, num_classes)
275 + # embed -> (batch_size, embedding_dim) or list(embedding_dim)
276 + # (batch_size, dimensions) or list(dimensions) if using MRL
277 + raise NotImplementedError
278
235 - def extract_states(
236 - self,
237 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
238 - pooling_metadata: PoolingMetadata,
239 - ) -> Union[list[torch.Tensor], torch.Tensor]:
240 - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
241 - prompt_token_ids = self.get_prompt_token_ids(pooling_metadata)
279
243 - pooled_data_lst = list[torch.Tensor]()
244 - if isinstance(hidden_states, list):
245 - for req_state, prompt_len in zip(hidden_states, prompt_lens):
246 - assert prompt_len == req_state.shape[0], \
247 - "partial prefill not supported with step pooling"
248 - pooled_data_lst = hidden_states
249 - else:
250 - offset = 0
251 - for prompt_len in prompt_lens:
252 - pooled_data_i = hidden_states[offset:offset + prompt_len]
253 - offset += prompt_len
254 - pooled_data_lst.append(pooled_data_i)
280 +class PoolerActivation(BasePoolerActivation):
281
256 - pooled_data = list[torch.Tensor]()
257 - returned_token_ids = self.returned_token_ids
258 - step_tag_id = self.step_tag_id
282 + @staticmethod
283 + def wraps(module: nn.Module):
284 + if isinstance(module, nn.Identity):
285 + return PoolerIdentity()
286 + if isinstance(module, (nn.Sigmoid, nn.Softmax)):
287 + return PoolerClassify()
288
260 - for data, token_id in zip(pooled_data_lst, prompt_token_ids):
261 - if returned_token_ids is not None and len(returned_token_ids) > 0:
262 - data = data[:, returned_token_ids]
289 + return LambdaPoolerActivation(module)
290 +
291 + @abstractmethod
292 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
293 + raise NotImplementedError
294 +
295 + def forward(self, pooled_data: _T) -> _T:
296 + if isinstance(pooled_data, list):
297 + return [self.forward_chunk(data) for data in pooled_data]
298 +
299 + return self.forward_chunk(pooled_data)
300
264 - if step_tag_id is not None:
265 - data = data[token_id == step_tag_id]
266 - pooled_data.append(data)
301
302 +class PoolerIdentity(PoolerActivation):
303 +
304 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
305 return pooled_data
306
307
308 +class PoolerNormalize(PoolerActivation):
309 +
310 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
311 + x = F.normalize(pooled_data.float(), p=2, dim=-1)
312 + return x.to(pooled_data.dtype)
313 +
314 +
315 +class PoolerClassify(PoolerActivation):
316 +
317 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
318 + num_labels = pooled_data.shape[-1]
319 + if num_labels < 2:
320 + return F.sigmoid(pooled_data.float()).to(pooled_data.dtype)
321 +
322 + return F.softmax(pooled_data.float(), dim=-1).to(pooled_data.dtype)
323 +
324 +
325 +class PoolerScore(PoolerActivation):
326 +
327 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
328 + num_labels = pooled_data.shape[-1]
329 + if num_labels < 2:
330 + return F.sigmoid(pooled_data.float()).to(pooled_data.dtype)
331 +
332 + return pooled_data
333 +
334 +
335 +class LambdaPoolerActivation(PoolerActivation):
336 +
337 + def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]):
338 + super().__init__()
339 +
340 + self.fn = fn
341 +
342 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
343 + return self.fn(pooled_data)
344 +
345 +
346 class PoolerHead(nn.Module):
347
273 - def __init__(self, *, normalize: bool, softmax: bool) -> None:
348 + @classmethod
349 + def from_config_with_defaults(
350 + cls,
351 + pooler_config: PoolerConfig,
352 + pooling_type: PoolingType,
353 + normalize: bool,
354 + softmax: bool,
355 + ) -> "PoolerHead":
356 + resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
357 + pooler_config=pooler_config,
358 + pooling_type=pooling_type,
359 + normalize=normalize,
360 + softmax=softmax,
361 + step_tag_id=None,
362 + returned_token_ids=None,
363 + )
364 +
365 + return cls.from_config(resolved_config)
366 +
367 + @classmethod
368 + def from_config(cls, pooler_config: ResolvedPoolingConfig) -> "PoolerHead":
369 + if pooler_config.normalize and pooler_config.softmax:
370 + raise ValueError("`normalize=True` and `softmax=True` should not "
371 + "be set together")
372 +
373 + activation: PoolerActivation
374 + if pooler_config.normalize:
375 + activation = PoolerNormalize()
376 + elif pooler_config.softmax:
377 + activation = PoolerClassify()
378 + else:
379 + activation = PoolerIdentity()
380 +
381 + return cls(activation)
382 +
383 + def __init__(self, activation: PoolerActivation) -> None:
384 super().__init__()
385
276 - self.normalize = normalize
277 - self.softmax = softmax
386 + self.activation = activation
387
388 def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
389 pooling_metadata: PoolingMetadata):
@@ -312,35 +421,21 @@ class PoolerHead(nn.Module):
421 for vecs, d in zip(pooled_data, dimensions_list)
422 ]
423
315 - if self.normalize:
316 - if isinstance(pooled_data, list):
317 - pooled_data = [
318 - F.normalize(data, p=2, dim=-1) for data in pooled_data
319 - ]
320 - else:
321 - pooled_data = F.normalize(pooled_data, p=2, dim=-1)
424 + return self.activation(pooled_data)
425
323 - if self.softmax:
324 - if isinstance(pooled_data, list):
325 - pooled_data = [
326 - F.softmax(data, dim=-1)
327 - if data.shape[-1] >= 2 else F.sigmoid(data)
328 - for data in pooled_data
329 - ]
330 - else:
331 - if pooled_data.shape[-1] >= 2:
332 - pooled_data = F.softmax(pooled_data, dim=-1)
333 - else:
334 - pooled_data = F.sigmoid(pooled_data)
426
336 - # shape:
337 - # classify (& score) -> (batch_size, num_classes)
338 - # embed -> (batch_size, embedding_dim) or list(embedding_dim)
339 - # (batch_size, dimensions) or list(dimensions) if using MRL
340 - return pooled_data
427 +class SimplePooler(BasePooler):
428 + """A layer that pools specific information from hidden states.
429
430 + This layer does the following:
431 + 1. Extracts specific tokens or aggregates data based on pooling method.
432 + 2. Normalizes output if specified.
433 + 3. Returns structured results as `PoolerOutput`.
434
343 -class Pooler(nn.Module):
435 + Attributes:
436 + pooling_type: The type of pooling to use.
437 + normalize: Whether to normalize the pooled data.
438 + """
439
440 @classmethod
441 def from_config_with_defaults(
@@ -349,23 +444,146 @@ class Pooler(nn.Module):
444 pooling_type: PoolingType,
445 normalize: bool,
446 softmax: bool,
447 + ) -> "SimplePooler":
448 + resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
449 + pooler_config=pooler_config,
450 + pooling_type=pooling_type,
451 + normalize=normalize,
452 + softmax=softmax,
453 + )
454 + assert resolved_config.pooling_type != PoolingType.STEP
455 +
456 + return cls.from_config(resolved_config)
457 +
458 + @classmethod
459 + def from_config(
460 + cls,
461 + pooler_config: ResolvedPoolingConfig,
462 + ) -> "SimplePooler":
463 + pooling = PoolingMethod.from_pooling_type(pooler_config.pooling_type)
464 + head = PoolerHead.from_config(pooler_config)
465 +
466 + return cls(pooling, head)
467 +
468 + def __init__(self, pooling: PoolingMethod, head: PoolerHead) -> None:
469 + super().__init__()
470 +
471 + self.pooling = pooling
472 + self.head = head
473 +
474 + def forward(
475 + self,
476 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
477 + pooling_metadata: PoolingMetadata,
478 + ) -> PoolerOutput:
479 + pooled_data = self.pooling(hidden_states, pooling_metadata)
480 + pooled_data = self.head(pooled_data, pooling_metadata)
481 + return build_output(pooled_data)
482 +
483 +
484 +class StepPooler(BasePooler):
485 +
486 + @classmethod
487 + def from_config(cls, pooler_config: ResolvedPoolingConfig) -> "StepPooler":
488 + assert pooler_config.pooling_type == PoolingType.STEP
489 +
490 + return cls(
491 + PoolerHead.from_config(pooler_config),
492 + step_tag_id=pooler_config.step_tag_id,
493 + returned_token_ids=pooler_config.returned_token_ids,
494 + )
495 +
496 + def __init__(
497 + self,
498 + head: PoolerHead,
499 + *,
500 step_tag_id: Optional[int] = None,
501 returned_token_ids: Optional[list[int]] = None,
354 - ) -> SimplePooler:
355 - return SimplePooler.from_pooling_type(
356 - pooling_type=PoolingType[pooler_config.pooling_type]
357 - if pooler_config.pooling_type is not None else pooling_type,
358 - normalize=pooler_config.normalize
359 - if pooler_config.normalize is not None else normalize,
360 - softmax=pooler_config.softmax
361 - if pooler_config.softmax is not None else softmax,
362 - step_tag_id=pooler_config.step_tag_id
363 - if pooler_config.step_tag_id is not None else step_tag_id,
364 - returned_token_ids=pooler_config.returned_token_ids
365 - if pooler_config.returned_token_ids is not None else
366 - returned_token_ids,
502 + ) -> None:
503 + super().__init__()
504 +
505 + self.pooling = AllPool()
506 + self.head = head
507 + self.step_tag_id = step_tag_id
508 + self.returned_token_ids = returned_token_ids
509 +
510 + def get_prompt_token_ids(
511 + self,
512 + pooling_metadata: PoolingMetadata,
513 + ) -> list[torch.Tensor]:
514 + if isinstance(pooling_metadata, V1PoolingMetadata):
515 + return [
516 + pooling_metadata.prompt_token_ids[i, :num]
517 + for i, num in enumerate(pooling_metadata.prompt_lens)
518 + ]
519 + return [
520 + torch.tensor(seq_data_i.prompt_token_ids)
521 + for seq_data_i in pooling_metadata.seq_data.values()
522 + ]
523 +
524 + def extract_states(
525 + self,
526 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
527 + pooling_metadata: PoolingMetadata,
528 + ) -> Union[list[torch.Tensor], torch.Tensor]:
529 + pooled_data_lst = self.pooling(hidden_states, pooling_metadata)
530 + prompt_token_ids = self.get_prompt_token_ids(pooling_metadata)
531 +
532 + pooled_data = list[torch.Tensor]()
533 + returned_token_ids = self.returned_token_ids
534 + step_tag_id = self.step_tag_id
535 +
536 + for data, token_id in zip(pooled_data_lst, prompt_token_ids):
537 + if returned_token_ids is not None and len(returned_token_ids) > 0:
538 + data = data[:, returned_token_ids]
539 +
540 + if step_tag_id is not None:
541 + data = data[token_id == step_tag_id]
542 + pooled_data.append(data)
543 +
544 + return pooled_data
545 +
546 + def forward(
547 + self,
548 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
549 + pooling_metadata: PoolingMetadata,
550 + ) -> PoolerOutput:
551 + pooled_data = self.extract_states(hidden_states, pooling_metadata)
552 + pooled_data = self.head(pooled_data, pooling_metadata)
553 + return build_output(pooled_data)
554 +
555 +
556 +class Pooler(nn.Module):
557 +
558 + @staticmethod
559 + def from_config_with_defaults(
560 + pooler_config: PoolerConfig,
561 + pooling_type: PoolingType,
562 + normalize: bool,
563 + softmax: bool,
564 + step_tag_id: Optional[int] = None,
565 + returned_token_ids: Optional[list[int]] = None,
566 + ) -> BasePooler:
567 + resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
568 + pooler_config=pooler_config,
569 + pooling_type=pooling_type,
570 + normalize=normalize,
571 + softmax=softmax,
572 + step_tag_id=step_tag_id,
573 + returned_token_ids=returned_token_ids,
574 )
575
576 + if pooling_type == PoolingType.STEP:
577 + return StepPooler.from_config(resolved_config)
578 +
579 + return SimplePooler.from_config(resolved_config)
580 +
581 +
582 +PoolingFn = Callable[
583 + [Union[torch.Tensor, list[torch.Tensor]], PoolingMetadata],
584 + Union[torch.Tensor, list[torch.Tensor]]]
585 +ClassifierFn = Callable[[torch.Tensor], torch.Tensor]
586 +
587
588 class ClassifierPooler(nn.Module):
589 """A pooling layer for classification tasks.
@@ -382,69 +600,39 @@ class ClassifierPooler(nn.Module):
600 def __init__(
601 self,
602 config: ModelConfig,
385 - classifier: nn.Module,
386 - pooler: Optional[nn.Module] = None,
387 - ):
603 + pooling: PoolingFn,
604 + classifier: ClassifierFn,
605 + act_fn: Optional[PoolerActivation] = None,
606 + ) -> None:
607 super().__init__()
608 +
609 + self.pooling = pooling
610 self.classifier = classifier
390 - self.pooler = pooler
611
612 self.classification_act_fn = get_classification_activation_function(
393 - config.hf_config)
613 + config.hf_config) if act_fn is None else act_fn
614 self.cross_encoder_act_fn = get_cross_encoder_activation_function(
395 - config.hf_config)
615 + config.hf_config) if act_fn is None else act_fn
616
617 def _get_act_fn(self, use_cross_encoder: bool):
618 return (self.cross_encoder_act_fn
619 if use_cross_encoder else self.classification_act_fn)
620
401 - def get_prompt_lens(
402 - self,
403 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
404 - pooling_metadata: PoolingMetadata,
405 - ) -> torch.Tensor:
406 - if isinstance(pooling_metadata, V1PoolingMetadata):
407 - return pooling_metadata.prompt_lens
408 - assert isinstance(hidden_states, torch.Tensor)
409 - return PoolingTensors.from_pooling_metadata(
410 - pooling_metadata, hidden_states.device).prompt_lens
411 -
621 def forward(
622 self,
623 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
624 pooling_metadata: PoolingMetadata,
625 ) -> PoolerOutput:
626 """Pools sentence pair scores from the hidden_states."""
418 - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
627 + pooled_data = self.pooling(hidden_states, pooling_metadata)
628
420 - pooled_data = list[torch.Tensor]()
421 - if isinstance(hidden_states, list):
422 - for req_state, prompt_len in zip(hidden_states, prompt_lens):
423 - assert prompt_len == req_state.shape[0], \
424 - "partial prefill not supported with classifier"
425 - pooled_data = hidden_states
629 + # apply classifier once on the full batch if possible
630 + if isinstance(pooled_data, torch.Tensor):
631 + pooled_output = self.classifier(pooled_data)
632 + elif len({data.shape for data in pooled_data}) <= 1:
633 + pooled_output = self.classifier(torch.stack(pooled_data))
634 else:
427 - offset = 0
428 - for prompt_len in prompt_lens:
429 - pooled_data_i = hidden_states[offset:offset + prompt_len]
430 - offset += prompt_len
431 - pooled_data.append(pooled_data_i)
432 -
433 - pooled_data_lst = []
434 - for pooled_data_i in pooled_data:
435 -
436 - if self.pooler is not None:
437 - final_shape_tensor = self.pooler(pooled_data_i)
438 - else:
439 - final_shape_tensor = self.classifier(pooled_data_i)
440 -
441 - pooled_data_lst.append(final_shape_tensor)
442 -
443 - pooled_output = torch.stack(pooled_data_lst)
444 -
445 - if self.pooler is not None:
446 - # apply classifier once on the full batch if possible
447 - pooled_output = self.classifier(pooled_output)
635 + pooled_output = [self.classifier(data) for data in pooled_data]
636
637 if isinstance(pooling_metadata, V0PoolingMetadata):
638 use_cross_encoder_list = [
@@ -469,5 +657,4 @@ class ClassifierPooler(nn.Module):
657 pooled_output)
658 ])
659
472 - pooled_outputs = [PoolingSequenceGroupOutput(data) for data in scores]
473 - return PoolerOutput(outputs=pooled_outputs)
660 + return build_output(scores)
vllm/model_executor/models/adapters.py
+48 -51
@@ -58,22 +58,27 @@ def _create_pooling_model_cls(
58 ) -> None:
59 super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
60
61 + self.vllm_config = vllm_config
62 +
63 # These are not used in pooling models
64 for attr in ("lm_head", "logits_processor"):
65 if hasattr(self, attr):
66 delattr(self, attr)
67
68 + # If the model already defines a pooler instance, don't overwrite it
69 + if not getattr(self, "_pooler", None):
70 + self._init_pooler(vllm_config, prefix=prefix)
71 +
72 + def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
73 pooler_config = vllm_config.model_config.pooler_config
74 assert pooler_config is not None
75
69 - # If the model already defines a pooler instance, don't overwrite it
70 - if not getattr(self, "_pooler", None):
71 - self._pooler = Pooler.from_config_with_defaults(
72 - pooler_config,
73 - pooling_type=default_pooling_type,
74 - normalize=default_normalize,
75 - softmax=default_softmax,
76 - )
76 + self._pooler = Pooler.from_config_with_defaults(
77 + pooler_config,
78 + pooling_type=default_pooling_type,
79 + normalize=default_normalize,
80 + softmax=default_softmax,
81 + )
82
83 def pooler(
84 self,
@@ -165,7 +170,9 @@ def as_seq_cls_model(cls: _T) -> _T:
170
171 # Lazy import
172 from vllm.model_executor.layers.linear import RowParallelLinear
168 - from vllm.model_executor.layers.pooler import PoolerOutput, PoolingType
173 + from vllm.model_executor.layers.pooler import (ClassifierPooler,
174 + PoolerOutput, PoolingType,
175 + SimplePooler)
176 from vllm.model_executor.models.interfaces import SupportsCrossEncoding
177 from vllm.model_executor.pooling_metadata import PoolingMetadata
178 from vllm.sequence import IntermediateTensors
@@ -182,30 +189,40 @@ def as_seq_cls_model(cls: _T) -> _T:
189 class ModelForSequenceClassification(ModelForPooling,
190 SupportsCrossEncoding):
191
185 - def __init__(
186 - self,
187 - *,
188 - vllm_config: "VllmConfig",
189 - prefix: str = "",
190 - **kwargs: Any,
191 - ) -> None:
192 - super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
193 -
192 + def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
193 config = vllm_config.model_config.hf_config
194 quant_config = vllm_config.quant_config
195
197 - self.vllm_config = vllm_config
198 - self.task = vllm_config.model_config.task
199 - self.pooling_type = (
200 - vllm_config.model_config.pooler_config.pooling_type)
201 -
202 - self.score = RowParallelLinear(config.hidden_size,
203 - config.num_labels,
204 - quant_config=quant_config,
205 - input_is_parallel=False,
206 - bias=False,
207 - prefix=maybe_prefix(
208 - prefix, "score"))
196 + self.score = RowParallelLinear(
197 + config.hidden_size,
198 + config.num_labels,
199 + input_is_parallel=False,
200 + bias=False,
201 + params_dtype=torch.float32,
202 + quant_config=quant_config,
203 + prefix=maybe_prefix(prefix, "score"),
204 + )
205 +
206 + pooler_config = vllm_config.model_config.pooler_config
207 + assert pooler_config is not None
208 +
209 + pooler = SimplePooler.from_config_with_defaults(
210 + pooler_config,
211 + pooling_type=PoolingType.LAST,
212 + normalize=False,
213 + softmax=True,
214 + )
215 +
216 + self._pooler = ClassifierPooler(
217 + vllm_config.model_config,
218 + pooling=pooler.pooling,
219 + classifier=self._classifier,
220 + act_fn=pooler.head.activation,
221 + )
222 +
223 + def _classifier(self, x: torch.Tensor):
224 + x, _ = self.score(x.float())
225 + return x
226
227 def forward(
228 self,
@@ -222,27 +239,7 @@ def as_seq_cls_model(cls: _T) -> _T:
239 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
240 pooling_metadata: PoolingMetadata,
241 ) -> PoolerOutput:
225 -
226 - def get_logits(hidden_states):
227 - if isinstance(hidden_states, list):
228 - logits = [self.score(state)[0] for state in hidden_states]
229 - else:
230 - logits, _ = self.score(hidden_states)
231 - return logits
232 -
233 - if self.pooling_type == PoolingType.ALL:
234 - logits = get_logits(hidden_states)
235 - return self._pooler(logits, pooling_metadata)
236 - else:
237 - hidden_states = self._pooler.extract_states(
238 - hidden_states, pooling_metadata)
239 - logits = get_logits(hidden_states)
240 - pooled_data = self._pooler.head(logits, pooling_metadata)
241 -
242 - pooled_outputs = [
243 - self._pooler.build_output(data) for data in pooled_data
244 - ]
245 - return PoolerOutput(outputs=pooled_outputs)
242 + return self._pooler(hidden_states, pooling_metadata)
243
244 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
245 tokens = getattr(self.config, "classifier_from_token", None)
vllm/model_executor/models/bert.py
+16 -9
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Iterable
5 -from typing import Optional
5 +from typing import Optional, Union
6
7 import torch
8 from torch import nn
@@ -18,7 +18,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
18 QKVParallelLinear,
19 RowParallelLinear)
20 from vllm.model_executor.layers.pooler import (ClassifierPooler, Pooler,
21 - PoolingType)
21 + PoolingMethod, PoolingType)
22 from vllm.model_executor.layers.quantization import QuantizationConfig
23 from vllm.model_executor.layers.vocab_parallel_embedding import (
24 VocabParallelEmbedding)
@@ -84,14 +84,18 @@ class BertPooler(nn.Module):
84
85 def __init__(self, config: BertConfig):
86 super().__init__()
87 +
88 + self.pooling = PoolingMethod.from_pooling_type(PoolingType.CLS)
89 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
90 self.activation = nn.Tanh()
91
90 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
91 - # We "pool" the model by simply taking the hidden state corresponding
92 - # to the first token.
93 - first_token_tensor = hidden_states[0, :]
94 - pooled_output = self.dense(first_token_tensor)
92 + def forward(
93 + self,
94 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
95 + pooling_metadata: PoolingMetadata,
96 + ) -> Union[torch.Tensor, list[torch.Tensor]]:
97 + pooled_output = self.pooling(hidden_states, pooling_metadata)
98 + pooled_output = self.dense(pooled_output)
99 pooled_output = self.activation(pooled_output)
100 return pooled_output
101
@@ -472,8 +476,11 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,
476 embedding_class=BertEmbedding,
477 add_pooling_layer=True)
478 self.classifier = nn.Linear(config.hidden_size, config.num_labels)
475 - self._pooler = ClassifierPooler(vllm_config.model_config,
476 - self.classifier, self.bert.pooler)
479 + self._pooler = ClassifierPooler(
480 + vllm_config.model_config,
481 + pooling=self.bert.pooler,
482 + classifier=self.classifier,
483 + )
484
485 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
486 loader = AutoWeightsLoader(self)
vllm/model_executor/models/gritlm.py
+2 -2
@@ -9,7 +9,7 @@ import torch.nn as nn
9
10 from vllm.config import ModelConfig, VllmConfig
11 from vllm.logger import init_logger
12 -from vllm.model_executor.layers.pooler import PoolerHead
12 +from vllm.model_executor.layers.pooler import PoolerHead, PoolerNormalize
13 from vllm.model_executor.models.llama import LlamaForCausalLM
14 from vllm.model_executor.pooling_metadata import (PoolingMetadata,
15 PoolingTensors)
@@ -49,7 +49,7 @@ class GritLMPooler(nn.Module):
49 self.embed_pattern_ids = tokens_to_ids(
50 ["▁<", "|", "embed", "|", ">", "<0x0A>"])
51
52 - self.head = PoolerHead(normalize=True, softmax=False)
52 + self.head = PoolerHead(PoolerNormalize())
53
54 def _find_array(self, arr: array, target: array, start_idx: int) -> int:
55 """
vllm/model_executor/models/interfaces.py
+1 -1
@@ -659,7 +659,7 @@ def supports_cross_encoding(
659 def has_step_pooler(model: Union[type[object], object]) -> bool:
660 """Check if the model uses step pooler."""
661 return is_pooling_model(model) and any(
662 - type(module).__name__ == "StepPool" for module in model.modules())
662 + type(module).__name__ == "StepPooler" for module in model.modules())
663
664
665 class SupportsQuant:
vllm/model_executor/models/jamba.py
+26 -13
@@ -19,7 +19,8 @@ from vllm.model_executor.layers.linear import (QKVParallelLinear,
19 RowParallelLinear)
20 from vllm.model_executor.layers.logits_processor import LogitsProcessor
21 from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
22 -from vllm.model_executor.layers.pooler import Pooler, PoolingType
22 +from vllm.model_executor.layers.pooler import (ClassifierPooler, PoolingType,
23 + SimplePooler)
24 from vllm.model_executor.layers.quantization import QuantizationConfig
25 from vllm.model_executor.layers.vocab_parallel_embedding import (
26 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
@@ -564,29 +565,41 @@ class JambaForSequenceClassification(JambaForCausalLM):
565
566 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
567 super().__init__(vllm_config=vllm_config, prefix=prefix)
568 +
569 config = vllm_config.model_config.hf_config
570 num_labels: int = config.num_labels
571 score_bias: bool = getattr(config, 'score_bias', False)
570 - self.score = nn.Linear(config.hidden_size, num_labels, bias=score_bias)
572 +
573 + # TODO: The original reward weights have float32 accuracy data, we
574 + # would like to load them in fp32 to get that extra precision.
575 + # Currently weight_loader passes the weight which is already in bf16
576 + self.score = nn.Linear(
577 + config.hidden_size,
578 + num_labels,
579 + bias=score_bias,
580 + dtype=torch.float32,
581 + )
582
583 pooler_config = vllm_config.model_config.pooler_config
573 - self._pooler = Pooler.from_config_with_defaults(
584 + assert pooler_config is not None
585 +
586 + pooler = SimplePooler.from_config_with_defaults(
587 pooler_config,
588 pooling_type=PoolingType.LAST,
589 normalize=False,
577 - softmax=False)
590 + softmax=False,
591 + )
592 +
593 + self._pooler = ClassifierPooler(
594 + vllm_config.model_config,
595 + pooling=pooler.pooling,
596 + classifier=self.score,
597 + act_fn=pooler.head.activation,
598 + )
599
600 def pooler(
601 self,
602 hidden_states: torch.Tensor,
603 pooling_metadata: PoolingMetadata,
604 ) -> Optional[PoolerOutput]:
584 - hidden_states = hidden_states.float()
585 - logits = self.score(hidden_states)
586 - return self._pooler(logits, pooling_metadata)
587 -
588 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
589 - # TODO: The reward weights themselves have float32 accuracy data, we
590 - # would like to load them in fp32 to get that extra precision.
591 - super().load_weights(weights)
592 - self.score = self.score.float()
605 + return self._pooler(hidden_states, pooling_metadata)
vllm/model_executor/models/modernbert.py
+18 -15
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Iterable
4 -from typing import Optional
4 +from typing import Optional, Union
5
6 import torch
7 from torch import nn
@@ -13,7 +13,8 @@ from vllm.config import VllmConfig
13 from vllm.distributed import get_tensor_model_parallel_world_size
14 from vllm.model_executor.layers.linear import (QKVParallelLinear,
15 RowParallelLinear)
16 -from vllm.model_executor.layers.pooler import ClassifierPooler
16 +from vllm.model_executor.layers.pooler import (BasePooler, ClassifierPooler,
17 + PoolingMethod, PoolingType)
18 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
19 from vllm.model_executor.layers.vocab_parallel_embedding import (
20 VocabParallelEmbedding)
@@ -252,10 +253,13 @@ class ModernBertModel(nn.Module):
253 return norm_outputs
254
255
255 -class ModernBertPooler(nn.Module):
256 +class ModernBertPooler(BasePooler):
257
258 def __init__(self, config: ModernBertConfig):
259 super().__init__()
260 +
261 + pooling_type = PoolingType[config.classifier_pooling.upper()]
262 + self.pooling = PoolingMethod.from_pooling_type(pooling_type)
263 self.dense = nn.Linear(config.hidden_size, config.hidden_size,
264 config.classifier_bias)
265 self.pooling_type = config.classifier_pooling
@@ -264,15 +268,12 @@ class ModernBertPooler(nn.Module):
268 eps=config.norm_eps,
269 bias=config.norm_bias)
270
267 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
268 - pooled_output = hidden_states
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}")
271 + def forward(
272 + self,
273 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
274 + pooling_metadata: PoolingMetadata,
275 + ) -> Union[torch.Tensor, list[torch.Tensor]]:
276 + pooled_output = self.pooling(hidden_states, pooling_metadata)
277 pooled_output = self.norm(self.act(self.dense(pooled_output)))
278 return pooled_output
279
@@ -287,9 +288,11 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,
288 self.model = ModernBertModel(vllm_config=vllm_config,
289 prefix=maybe_prefix(prefix, "modernbert"))
290 self.classifier = nn.Linear(config.hidden_size, config.num_labels)
290 - self._pooler = ClassifierPooler(vllm_config.model_config,
291 - self.classifier,
292 - ModernBertPooler(config))
291 + self._pooler = ClassifierPooler(
292 + vllm_config.model_config,
293 + pooling=ModernBertPooler(config),
294 + classifier=self.classifier,
295 + )
296
297 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
298
vllm/model_executor/models/roberta.py
+8 -5
@@ -9,7 +9,7 @@ from torch import nn
9 from transformers import RobertaConfig
10
11 from vllm.config import VllmConfig
12 -from vllm.model_executor.layers.pooler import ClassifierPooler
12 +from vllm.model_executor.layers.pooler import ClassifierPooler, CLSPool
13 from vllm.model_executor.layers.vocab_parallel_embedding import (
14 VocabParallelEmbedding)
15 from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel
@@ -106,8 +106,8 @@ class RobertaClassificationHead(nn.Module):
106 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
107 self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
108
109 - def forward(self, features, **kwargs):
110 - x = features[0, :] # take <s> token (equiv. to [CLS])
109 + def forward(self, x: torch.Tensor) -> torch.Tensor:
110 + # CLSPool has already been applied in `pooling`
111 x = self.dense(x)
112 x = torch.tanh(x)
113 x = self.out_proj(x)
@@ -188,8 +188,11 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
188 add_pooling_layer=False)
189 self.classifier = RobertaClassificationHead(config)
190
191 - self._pooler = ClassifierPooler(vllm_config.model_config,
192 - self.classifier)
191 + self._pooler = ClassifierPooler(
192 + vllm_config.model_config,
193 + pooling=CLSPool(),
194 + classifier=self.classifier,
195 + )
196
197 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
198 loader = AutoWeightsLoader(self)
vllm/transformers_utils/config.py
-24
@@ -17,7 +17,6 @@ from huggingface_hub.utils import (EntryNotFoundError, HfHubHTTPError,
17 HFValidationError, LocalEntryNotFoundError,
18 RepositoryNotFoundError,
19 RevisionNotFoundError)
20 -from torch import nn
20 from transformers import GenerationConfig, PretrainedConfig
21 from transformers.models.auto.image_processing_auto import (
22 get_image_processor_config)
@@ -44,7 +43,6 @@ from vllm.transformers_utils.configs import (ChatGLMConfig, Cohere2Config,
43 # yapf: enable
44 from vllm.transformers_utils.configs.mistral import adapt_config_dict
45 from vllm.transformers_utils.utils import check_gguf_file
47 -from vllm.utils import resolve_obj_by_qualname
46
47 if envs.VLLM_USE_MODELSCOPE:
48 from modelscope import AutoConfig
@@ -775,28 +773,6 @@ def try_get_generation_config(
773 return None
774
775
778 -def get_classification_activation_function(config: PretrainedConfig):
779 - return nn.Sigmoid() if config.num_labels == 1 else nn.Softmax()
780 -
781 -
782 -def get_cross_encoder_activation_function(config: PretrainedConfig):
783 - function_name: Optional[str] = None
784 - if (hasattr(config, "sentence_transformers")
785 - and "activation_fn" in config.sentence_transformers):
786 - function_name = config.sentence_transformers["activation_fn"]
787 - elif (hasattr(config, "sbert_ce_default_activation_function")
788 - and config.sbert_ce_default_activation_function is not None):
789 - function_name = config.sbert_ce_default_activation_function
790 -
791 - if function_name is not None:
792 - assert function_name.startswith("torch.nn.modules."), (
793 - "Loading of activation functions is restricted to "
794 - "torch.nn.modules for security reasons")
795 - return resolve_obj_by_qualname(function_name)()
796 -
797 - return nn.Sigmoid() if config.num_labels == 1 else nn.Identity()
798 -
799 -
776 def try_get_safetensors_metadata(
777 model: str,
778 *,