@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 042af0c8d3
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][1/N] Support multiple poolers at model level (#21227)

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

Cyrus Leung committed Jul 21, 2025 at 17:22 UTC 042af0c8d3f0b8b5319f34e4cb9b690981bb5da4
22 files changed +550 -414
docs/models/pooling_models.md
+39 -14
@@ -11,26 +11,51 @@ before returning them.
11 As shown in the [Compatibility Matrix](../features/compatibility_matrix.md), most vLLM features are not applicable to
12 pooling models as they only work on the generation or decode stage, so performance may not improve as much.
13
14 -For pooling models, we support the following `--task` options.
15 -The selected option sets the default pooler used to extract the final hidden states:
14 +If the model doesn't implement this interface, you can set `--task` which tells vLLM
15 +to convert the model into a pooling model.
16
17 -| Task | Pooling Type | Normalization | Softmax |
18 -|---------------------------------|----------------|-----------------|-----------|
19 -| Embedding (`embed`) | `LAST` | ✅︎ | ❌ |
20 -| Classification (`classify`) | `LAST` | ❌ | ✅︎ |
21 -| Sentence Pair Scoring (`score`) | \* | \* | \* |
17 +| `--task` | Model type | Supported pooling tasks |
18 +|------------|----------------------|-------------------------------|
19 +| `embed` | Embedding model | `encode`, `embed` |
20 +| `classify` | Classification model | `encode`, `classify`, `score` |
21 +| `reward` | Reward model | `encode` |
22
23 -\*The default pooler is always defined by the model.
23 +## Pooling Tasks
24
25 -!!! note
26 - If the model's implementation in vLLM defines its own pooler, the default pooler is set to that instead of the one specified in this table.
25 +In vLLM, we define the following pooling tasks and corresponding APIs:
26 +
27 +| Task | APIs |
28 +|------------|--------------------|
29 +| `encode` | `encode` |
30 +| `embed` | `embed`, `score`\* |
31 +| `classify` | `classify` |
32 +| `score` | `score` |
33 +
34 +\*The `score` API falls back to `embed` task if the model does not support `score` task.
35 +
36 +Each pooling model in vLLM supports one or more of these tasks according to [Pooler.get_supported_tasks][vllm.model_executor.layers.Pooler.get_supported_tasks].
37 +
38 +By default, the pooler assigned to each task has the following attributes:
39 +
40 +| Task | Pooling Type | Normalization | Softmax |
41 +|------------|----------------|---------------|---------|
42 +| `encode` | `ALL` | ❌ | ❌ |
43 +| `embed` | `LAST` | ✅︎ | ❌ |
44 +| `classify` | `LAST` | ❌ | ✅︎ |
45 +
46 +These defaults may be overridden by the model's implementation in vLLM.
47
48 When loading [Sentence Transformers](https://huggingface.co/sentence-transformers) models,
29 -we attempt to override the default pooler based on its Sentence Transformers configuration file (`modules.json`).
49 +we attempt to override the defaults based on its Sentence Transformers configuration file (`modules.json`),
50 +which takes priority over the model's defaults.
51 +
52 +You can further customize this via the `--override-pooler-config` option,
53 +which takes priority over both the model's and Sentence Transformers's defaults.
54 +
55 +!!! note
56
31 -!!! tip
32 - You can customize the model's pooling method via the `--override-pooler-config` option,
33 - which takes priority over both the model's and Sentence Transformers's defaults.
57 + The above configuration may be disregarded if the model's implementation in vLLM defines its own pooler
58 + that is not based on [PoolerConfig][vllm.config.PoolerConfig].
59
60 ## Offline Inference
61
tests/models/test_transformers.py
+1 -1
@@ -144,7 +144,7 @@ def test_quantization(
144 "model",
145 ["jason9693/Qwen2.5-1.5B-apeach"],
146 )
147 -@pytest.mark.parametrize("dtype", ["half"])
147 +@pytest.mark.parametrize("dtype", ["float"])
148 def test_classify(
149 hf_runner,
150 vllm_runner,
tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py
+8 -7
@@ -8,7 +8,7 @@ import torch
8 import torch.nn as nn
9
10 from vllm.config import VllmConfig
11 -from vllm.model_executor.layers.pooler import Pooler, PoolingType
11 +from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
12 from vllm.model_executor.models.gemma2 import Gemma2Model
13 from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
14 from vllm.sequence import IntermediateTensors
@@ -26,12 +26,13 @@ class MyGemma2Embedding(nn.Module):
26 self.model = Gemma2Model(vllm_config=vllm_config,
27 prefix=maybe_prefix(prefix, "model"))
28
29 - self.pooler = Pooler.from_config_with_defaults(
30 - vllm_config.model_config.pooler_config,
31 - pooling_type=PoolingType.LAST,
32 - normalize=True,
33 - softmax=False,
34 - )
29 + pooler_config = vllm_config.model_config.pooler_config
30 + assert pooler_config is not None
31 +
32 + self.pooler = DispatchPooler({
33 + "encode": Pooler.for_encode(pooler_config),
34 + "embed": Pooler.for_embed(pooler_config),
35 + })
36
37 self.make_empty_intermediate_tensors = (
38 self.model.make_empty_intermediate_tensors)
vllm/config.py
+4 -4
@@ -94,7 +94,7 @@ ConfigT = TypeVar("ConfigT", bound=ConfigType)
94 TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",
95 "score", "reward", "transcription", "draft"]
96
97 -_ResolvedTask = Literal["generate", "transcription", "pooling", "embed",
97 +_ResolvedTask = Literal["generate", "transcription", "encode", "embed",
98 "classify", "reward", "draft"]
99
100 RunnerOption = Literal["auto", "generate", "pooling", "draft"]
@@ -103,7 +103,7 @@ RunnerType = Literal["generate", "pooling", "draft"]
103
104 _RUNNER_TASKS: dict[RunnerType, list[_ResolvedTask]] = {
105 "generate": ["generate", "transcription"],
106 - "pooling": ["pooling", "embed", "classify", "reward"],
106 + "pooling": ["encode", "embed", "classify", "reward"],
107 "draft": [],
108 }
109
@@ -579,7 +579,7 @@ class ModelConfig:
579 # user-selected task
580 if runner_type == "pooling" and self.task == "auto":
581 selected_task = all_supported_tasks[runner_type][-1]
582 - assert selected_task != "pooling"
582 + assert selected_task != "encode"
583 self.task = selected_task
584 self.supported_runner_types = supported_runner_types
585 self.runner_type = runner_type
@@ -884,7 +884,7 @@ class ModelConfig:
884
885 supported_tasks = list[_ResolvedTask]()
886 if registry.is_pooling_model(architectures):
887 - supported_tasks.append("pooling")
887 + supported_tasks.append("encode")
888
889 # For now, users must specify the task (other than "pooling")
890 # to use for pooling models
vllm/entrypoints/openai/api_server.py
+1 -1
@@ -1668,7 +1668,7 @@ async def init_app_state(
1668 request_logger=request_logger,
1669 chat_template=resolved_chat_template,
1670 chat_template_content_format=args.chat_template_content_format,
1671 - ) if "pooling" in model_config.supported_tasks else None
1671 + ) if "encode" in model_config.supported_tasks else None
1672 state.openai_serving_embedding = OpenAIServingEmbedding(
1673 engine_client,
1674 model_config,
vllm/model_executor/layers/pooler.py
+175 -171
@@ -1,15 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 +from collections.abc import Mapping, Set
5 from dataclasses import dataclass
6 from enum import IntEnum
7 +from itertools import groupby
8 from typing import Callable, Optional, TypeVar, Union
9
10 import torch
11 import torch.nn as nn
12 import torch.nn.functional as F
13 from transformers import PretrainedConfig
12 -from typing_extensions import assert_never
14
15 from vllm.config import ModelConfig, PoolerConfig
16 from vllm.model_executor.pooling_metadata import ( # noqa: E501
@@ -21,6 +22,10 @@ from vllm.utils import resolve_obj_by_qualname
22 from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata
23
24 PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]
25 +PoolingFn = Callable[
26 + [Union[torch.Tensor, list[torch.Tensor]], PoolingMetadata],
27 + Union[torch.Tensor, list[torch.Tensor]]]
28 +ClassifierFn = Callable[[torch.Tensor], torch.Tensor]
29
30
31 class PoolingType(IntEnum):
@@ -79,37 +84,81 @@ class Pooler(nn.Module, ABC):
84 """The interface required for all poolers used in pooling models in vLLM."""
85
86 @staticmethod
82 - def from_config_with_defaults(
87 + def for_encode(
88 pooler_config: PoolerConfig,
84 - pooling_type: PoolingType,
85 - normalize: bool,
86 - softmax: bool,
87 - step_tag_id: Optional[int] = None,
88 - returned_token_ids: Optional[list[int]] = None,
89 - ) -> "Pooler":
89 + *,
90 + default_pooling_type: PoolingType = PoolingType.ALL,
91 + default_normalize: bool = False,
92 + default_softmax: bool = False,
93 + default_step_tag_id: Optional[int] = None,
94 + default_returned_token_ids: Optional[list[int]] = None,
95 + ):
96 resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
97 pooler_config=pooler_config,
92 - pooling_type=pooling_type,
93 - normalize=normalize,
94 - softmax=softmax,
95 - step_tag_id=step_tag_id,
96 - returned_token_ids=returned_token_ids,
98 + pooling_type=default_pooling_type,
99 + normalize=default_normalize,
100 + softmax=default_softmax,
101 + step_tag_id=default_step_tag_id,
102 + returned_token_ids=default_returned_token_ids,
103 )
104
99 - if pooling_type == PoolingType.STEP:
105 + if resolved_config.pooling_type == PoolingType.STEP:
106 return StepPooler.from_config(resolved_config)
107
108 return SimplePooler.from_config(resolved_config)
109
104 - def get_pooling_updates(
105 - self,
106 - task: PoolingTask,
107 - ) -> Optional[PoolingParamsUpdate]:
110 + @staticmethod
111 + def for_embed(
112 + pooler_config: PoolerConfig,
113 + *,
114 + default_pooling_type: PoolingType = PoolingType.LAST,
115 + default_normalize: bool = True,
116 + default_softmax: bool = False,
117 + ):
118 + resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
119 + pooler_config=pooler_config,
120 + pooling_type=default_pooling_type,
121 + normalize=default_normalize,
122 + softmax=default_softmax,
123 + )
124 +
125 + return SimplePooler.from_config(resolved_config)
126 +
127 + @staticmethod
128 + def for_classify(
129 + pooler_config: PoolerConfig,
130 + classifier: Optional[ClassifierFn],
131 + *,
132 + default_pooling_type: PoolingType = PoolingType.LAST,
133 + default_normalize: bool = False,
134 + default_softmax: bool = True,
135 + ):
136 + resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
137 + pooler_config=pooler_config,
138 + pooling_type=default_pooling_type,
139 + normalize=default_normalize,
140 + softmax=default_softmax,
141 + )
142 + base_pooler = SimplePooler.from_config(resolved_config)
143 + if classifier is None:
144 + return base_pooler
145 +
146 + return ClassifierPooler(
147 + pooling=base_pooler.pooling,
148 + classifier=classifier,
149 + act_fn=base_pooler.head.activation,
150 + )
151 +
152 + @abstractmethod
153 + def get_supported_tasks(self) -> Set[PoolingTask]:
154 + """Determine which pooling tasks are supported."""
155 + raise NotImplementedError
156 +
157 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
158 """
109 - Construct the pooling parameters to use for a task,
110 - or `None` if the task is not supported.
159 + Construct the updated pooling parameters to use for a supported task.
160 """
112 - return None
161 + return PoolingParamsUpdate()
162
163 @abstractmethod
164 def forward(
@@ -127,9 +176,8 @@ def get_prompt_lens(
176 if isinstance(pooling_metadata, V1PoolingMetadata):
177 return pooling_metadata.prompt_lens
178
130 - assert isinstance(hidden_states, torch.Tensor)
179 return PoolingTensors.from_pooling_metadata(
132 - pooling_metadata, hidden_states.device).prompt_lens
180 + pooling_metadata, hidden_states[0].device).prompt_lens
181
182
183 def get_prompt_token_ids(
@@ -149,6 +197,21 @@ def get_prompt_token_ids(
197 ]
198
199
200 +def get_tasks(pooling_metadata: PoolingMetadata) -> list[PoolingTask]:
201 + if isinstance(pooling_metadata, V0PoolingMetadata):
202 + pooling_params = [p for _, p in pooling_metadata.seq_groups]
203 + else:
204 + pooling_params = pooling_metadata.pooling_params
205 +
206 + tasks: list[PoolingTask] = [
207 + task for pooling_param in pooling_params
208 + if (task := pooling_param.task) is not None
209 + ]
210 + assert len(pooling_params) == len(tasks)
211 +
212 + return tasks
213 +
214 +
215 def get_classification_activation_function(config: PretrainedConfig):
216 return PoolerClassify()
217
@@ -172,7 +235,8 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):
235 return PoolerScore()
236
237
175 -def build_output(all_data: torch.Tensor) -> PoolerOutput:
238 +def build_output(
239 + all_data: Union[torch.Tensor, list[torch.Tensor]], ) -> PoolerOutput:
240 all_outputs = [PoolingSequenceGroupOutput(data) for data in all_data]
241 return PoolerOutput(outputs=all_outputs)
242
@@ -193,12 +257,12 @@ class PoolingMethod(nn.Module, ABC):
257 raise NotImplementedError(f"Unsupported method: {pooling_type}")
258
259 @abstractmethod
196 - def get_pooling_updates(
197 - self,
198 - task: PoolingTask,
199 - ) -> Optional[PoolingParamsUpdate]:
260 + def get_supported_tasks(self) -> Set[PoolingTask]:
261 raise NotImplementedError
262
263 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
264 + return PoolingParamsUpdate()
265 +
266 @abstractmethod
267 def forward_one(
268 self,
@@ -237,16 +301,8 @@ class PoolingMethod(nn.Module, ABC):
301
302 class CLSPool(PoolingMethod):
303
240 - def get_pooling_updates(
241 - self,
242 - task: PoolingTask,
243 - ) -> Optional[PoolingParamsUpdate]:
244 - # The equalities are split up to keep mypy happy
245 - if (task == "encode" or task == "embed" or task == "classify"
246 - or task == "score"):
247 - return PoolingParamsUpdate()
248 -
249 - assert_never(task)
304 + def get_supported_tasks(self) -> Set[PoolingTask]:
305 + return {"encode", "embed", "classify", "score"}
306
307 def forward_one(
308 self,
@@ -270,16 +326,8 @@ class CLSPool(PoolingMethod):
326
327 class LastPool(PoolingMethod):
328
273 - def get_pooling_updates(
274 - self,
275 - task: PoolingTask,
276 - ) -> Optional[PoolingParamsUpdate]:
277 - # The equalities are split up to keep mypy happy
278 - if (task == "encode" or task == "embed" or task == "classify"
279 - or task == "score"):
280 - return PoolingParamsUpdate()
281 -
282 - assert_never(task)
329 + def get_supported_tasks(self) -> Set[PoolingTask]:
330 + return {"encode", "embed", "classify", "score"}
331
332 def forward_one(
333 self,
@@ -299,18 +347,8 @@ class LastPool(PoolingMethod):
347
348 class AllPool(PoolingMethod):
349
302 - def get_pooling_updates(
303 - self,
304 - task: PoolingTask,
305 - ) -> Optional[PoolingParamsUpdate]:
306 - if task == "encode":
307 - return PoolingParamsUpdate()
308 -
309 - # The equalities are split up to keep mypy happy
310 - if task == "embed" or task == "classify" or task == "score":
311 - return None
312 -
313 - assert_never(task)
350 + def get_supported_tasks(self) -> Set[PoolingTask]:
351 + return {"encode"}
352
353 def forward_one(
354 self,
@@ -327,28 +365,13 @@ class AllPool(PoolingMethod):
365 hidden_states: torch.Tensor,
366 prompt_lens: torch.Tensor,
367 ) -> Union[list[torch.Tensor], torch.Tensor]:
330 - offset = 0
331 - pooled_data = list[torch.Tensor]()
332 -
333 - for prompt_len in prompt_lens:
334 - pooled_data.append(hidden_states[offset:offset + prompt_len])
335 - offset += prompt_len
336 -
337 - return pooled_data
368 + return list(hidden_states.split_with_sizes(prompt_lens.tolist()))
369
370
371 class MeanPool(PoolingMethod):
372
342 - def get_pooling_updates(
343 - self,
344 - task: PoolingTask,
345 - ) -> Optional[PoolingParamsUpdate]:
346 - # The equalities are split up to keep mypy happy
347 - if (task == "encode" or task == "embed" or task == "classify"
348 - or task == "score"):
349 - return PoolingParamsUpdate()
350 -
351 - assert_never(task)
373 + def get_supported_tasks(self) -> Set[PoolingTask]:
374 + return {"encode", "embed", "classify", "score"}
375
376 def forward_one(
377 self,
@@ -529,24 +552,6 @@ class SimplePooler(Pooler):
552 3. Returns structured results as `PoolerOutput`.
553 """
554
532 - @classmethod
533 - def from_config_with_defaults( # type: ignore[override]
534 - cls,
535 - pooler_config: PoolerConfig,
536 - pooling_type: PoolingType,
537 - normalize: bool,
538 - softmax: bool,
539 - ) -> "SimplePooler":
540 - resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
541 - pooler_config=pooler_config,
542 - pooling_type=pooling_type,
543 - normalize=normalize,
544 - softmax=softmax,
545 - )
546 - assert resolved_config.pooling_type != PoolingType.STEP
547 -
548 - return cls.from_config(resolved_config)
549 -
555 @classmethod
556 def from_config(
557 cls,
@@ -563,10 +568,10 @@ class SimplePooler(Pooler):
568 self.pooling = pooling
569 self.head = head
570
566 - def get_pooling_updates(
567 - self,
568 - task: PoolingTask,
569 - ) -> Optional[PoolingParamsUpdate]:
571 + def get_supported_tasks(self) -> Set[PoolingTask]:
572 + return self.pooling.get_supported_tasks()
573 +
574 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
575 return self.pooling.get_pooling_updates(task)
576
577 def forward(
@@ -627,18 +632,11 @@ class StepPooler(Pooler):
632
633 return pooled_data
634
630 - def get_pooling_updates(
631 - self,
632 - task: PoolingTask,
633 - ) -> Optional[PoolingParamsUpdate]:
634 - if task == "encode":
635 - return PoolingParamsUpdate(requires_token_ids=True)
635 + def get_supported_tasks(self) -> Set[PoolingTask]:
636 + return {"encode"}
637
637 - # The equalities are split up to keep mypy happy
638 - if task == "embed" or task == "classify" or task == "score":
639 - return None
640 -
641 - assert_never(task)
638 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
639 + return PoolingParamsUpdate(requires_token_ids=True)
640
641 def forward(
642 self,
@@ -650,68 +648,43 @@ class StepPooler(Pooler):
648 return build_output(pooled_data)
649
650
653 -PoolingFn = Callable[
654 - [Union[torch.Tensor, list[torch.Tensor]], PoolingMetadata],
655 - Union[torch.Tensor, list[torch.Tensor]]]
656 -ClassifierFn = Callable[[torch.Tensor], torch.Tensor]
657 -
658 -
659 -class ClassifierPooler(nn.Module):
651 +class ClassifierPooler(Pooler):
652 """A pooling layer for classification tasks.
653
654 This layer does the following:
655 1. Applies a classification layer to the hidden states.
656 2. Optionally applies a pooler layer.
665 - 3. Applies an activation function to the output. In the case of
666 - classification models it is either sigmoid or softmax. In the
667 - case of scoring models, the same behavior is configuration
668 - dependent, as in the sentence-transformers library.
657 + 3. Applies an activation function to the output.
658 """
659
660 + @staticmethod
661 + def act_fn_for_seq_cls(config: ModelConfig):
662 + return get_classification_activation_function(config.hf_config)
663 +
664 + @staticmethod
665 + def act_fn_for_cross_encoder(config: ModelConfig):
666 + return get_cross_encoder_activation_function(config.hf_config)
667 +
668 def __init__(
669 self,
673 - config: ModelConfig,
670 pooling: PoolingFn,
671 classifier: ClassifierFn,
676 - act_fn: Optional[PoolerActivation] = None,
672 + act_fn: PoolerActivation,
673 ) -> None:
674 super().__init__()
675
676 self.pooling = pooling
677 self.classifier = classifier
678 + self.act_fn = act_fn
679
683 - self.classification_act_fn = get_classification_activation_function(
684 - config.hf_config) if act_fn is None else act_fn
685 - self.cross_encoder_act_fn = get_cross_encoder_activation_function(
686 - config.hf_config) if act_fn is None else act_fn
687 -
688 - def _get_act_fn(self, task: PoolingTask):
689 - if task == "encode" or task == "classify":
690 - return self.classification_act_fn
691 - if task == "score":
692 - return self.cross_encoder_act_fn
693 -
694 - raise ValueError(f"Unsupported task: {task!r}")
695 -
696 - def get_pooling_updates(
697 - self,
698 - task: PoolingTask,
699 - ) -> Optional[PoolingParamsUpdate]:
700 - # The equalities are split up to keep mypy happy
701 - if task == "encode" or task == "classify" or task == "score":
702 - return PoolingParamsUpdate()
703 -
704 - if task == "embed":
705 - return None
706 -
707 - assert_never(task)
680 + def get_supported_tasks(self) -> Set[PoolingTask]:
681 + return {"classify", "score"}
682
683 def forward(
684 self,
685 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
686 pooling_metadata: PoolingMetadata,
687 ) -> PoolerOutput:
714 - """Pools sentence pair scores from the hidden_states."""
688 pooled_data = self.pooling(hidden_states, pooling_metadata)
689
690 # apply classifier once on the full batch if possible
@@ -722,28 +695,59 @@ class ClassifierPooler(nn.Module):
695 else:
696 pooled_output = [self.classifier(data) for data in pooled_data]
697
725 - task_list: list[PoolingTask]
726 - if isinstance(pooling_metadata, V0PoolingMetadata):
727 - task_list = [
728 - task for _, pooling_param in pooling_metadata.seq_groups
729 - if (task := pooling_param.task) is not None
730 - ]
731 - else:
732 - task_list = [
733 - task for pooling_param in pooling_metadata.pooling_params
734 - if (task := pooling_param.task) is not None
735 - ]
698 + scores = self.act_fn(pooled_output)
699 +
700 + return build_output(scores)
701 +
702 +
703 +class DispatchPooler(Pooler):
704 + """Dispatches calls to a sub-pooler based on the pooling task."""
705 +
706 + def __init__(self, poolers_by_task: Mapping[PoolingTask, Pooler]) -> None:
707 + super().__init__()
708 +
709 + for task, pooler in poolers_by_task.items():
710 + if task not in pooler.get_supported_tasks():
711 + raise ValueError(
712 + f"{pooler=} does not support {task=}. "
713 + f"Supported tasks: {pooler.get_supported_tasks()}")
714 +
715 + self.poolers_by_task = poolers_by_task
716 +
717 + def get_supported_tasks(self) -> Set[PoolingTask]:
718 + return set(self.poolers_by_task)
719
737 - assert len(task_list) == len(pooled_output)
720 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
721 + return self.poolers_by_task[task].get_pooling_updates(task)
722
739 - # shape of scores: (batch_size, num_labels)
740 - if len(set(task_list)) <= 1:
741 - act_fn = self._get_act_fn(task_list[0])
742 - scores = act_fn(pooled_output)
723 + def forward(
724 + self,
725 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],
726 + pooling_metadata: PoolingMetadata,
727 + ) -> PoolerOutput:
728 + poolers_by_task = self.poolers_by_task
729 +
730 + if isinstance(hidden_states, list):
731 + hidden_states_lst = hidden_states
732 else:
744 - scores = torch.stack([
745 - self._get_act_fn(task)(vecs)
746 - for task, vecs in zip(task_list, pooled_output)
747 - ])
733 + prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)
734 + hidden_states_lst = list(hidden_states.split(prompt_lens.tolist()))
735
749 - return build_output(scores)
736 + outputs = list[PoolingSequenceGroupOutput]()
737 + offset = 0
738 + for task, group in groupby(get_tasks(pooling_metadata)):
739 + if not (pooler := poolers_by_task.get(task)):
740 + raise ValueError(
741 + f"Unsupported task: {task} "
742 + f"Supported tasks: {self.get_supported_tasks()}")
743 +
744 + num_items = len(list(group))
745 + group_output: PoolerOutput = pooler(
746 + hidden_states_lst[offset:offset + num_items],
747 + pooling_metadata[offset:offset + num_items],
748 + )
749 +
750 + outputs.extend(group_output.outputs)
751 + offset += num_items
752 +
753 + return PoolerOutput(outputs)
vllm/model_executor/models/adapters.py
+51 -57
@@ -13,7 +13,6 @@ from .interfaces_base import VllmModelForPooling, is_pooling_model
13
14 if TYPE_CHECKING:
15 from vllm.config import VllmConfig
16 - from vllm.model_executor.layers.pooler import PoolingType
16
17 _T = TypeVar("_T", bound=type[nn.Module])
18
@@ -34,16 +33,8 @@ def _get_pooling_model_name(orig_model_name: str, pooling_suffix: str) -> str:
33 return model_name + pooling_suffix
34
35
37 -def _create_pooling_model_cls(
38 - orig_cls: _T,
39 - *,
40 - default_pooling_type: "PoolingType",
41 - default_normalize: bool,
42 - default_softmax: bool,
43 -) -> _T:
36 +def _create_pooling_model_cls(orig_cls: _T) -> _T:
37 # Lazy import
45 - from vllm.model_executor.layers.pooler import Pooler
46 -
38 from .utils import AutoWeightsLoader, WeightsMapper
39
40 class ModelForPooling(orig_cls, VllmModelForPooling):
@@ -71,15 +62,7 @@ def _create_pooling_model_cls(
62 self._init_pooler(vllm_config, prefix=prefix)
63
64 def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
74 - pooler_config = vllm_config.model_config.pooler_config
75 - assert pooler_config is not None
76 -
77 - self.pooler = Pooler.from_config_with_defaults(
78 - pooler_config,
79 - pooling_type=default_pooling_type,
80 - normalize=default_normalize,
81 - softmax=default_softmax,
82 - )
65 + raise NotImplementedError
66
67 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
68 # TODO: Support uninitialized params tracking
@@ -132,14 +115,20 @@ def as_embedding_model(cls: _T) -> _T:
115 return cls
116
117 # Lazy import
135 - from vllm.model_executor.layers.pooler import PoolingType
136 -
137 - ModelForEmbedding = _create_pooling_model_cls(
138 - cls,
139 - default_pooling_type=PoolingType.LAST,
140 - default_normalize=True,
141 - default_softmax=False,
142 - )
118 + from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
119 +
120 + class ModelForEmbedding(_create_pooling_model_cls(cls)):
121 +
122 + def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
123 + pooler_config = vllm_config.model_config.pooler_config
124 + assert pooler_config is not None
125 +
126 + self.pooler = DispatchPooler(
127 + {
128 + "encode": Pooler.for_encode(pooler_config),
129 + "embed": Pooler.for_embed(pooler_config),
130 + }, )
131 +
132 ModelForEmbedding.__name__ = \
133 _get_pooling_model_name(cls.__name__, "ForEmbedding")
134
@@ -165,20 +154,14 @@ def as_seq_cls_model(cls: _T) -> _T:
154 # Lazy import
155 from vllm.model_executor.layers.linear import RowParallelLinear
156 from vllm.model_executor.layers.pooler import (ClassifierPooler,
168 - PoolingType, SimplePooler)
157 + DispatchPooler, Pooler,
158 + PoolingMethod, PoolingType)
159 from vllm.model_executor.models.interfaces import SupportsCrossEncoding
160 from vllm.sequence import IntermediateTensors
161
162 from .utils import maybe_prefix
163
174 - ModelForPooling = _create_pooling_model_cls(
175 - cls,
176 - default_pooling_type=PoolingType.LAST,
177 - default_normalize=False,
178 - default_softmax=True,
179 - )
180 -
181 - class ModelForSequenceClassification(ModelForPooling,
164 + class ModelForSequenceClassification(_create_pooling_model_cls(cls),
165 SupportsCrossEncoding):
166
167 def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
@@ -198,19 +181,28 @@ def as_seq_cls_model(cls: _T) -> _T:
181 pooler_config = vllm_config.model_config.pooler_config
182 assert pooler_config is not None
183
201 - pooler = SimplePooler.from_config_with_defaults(
202 - pooler_config,
203 - pooling_type=PoolingType.LAST,
204 - normalize=False,
205 - softmax=True,
206 - )
207 -
208 - self.pooler = ClassifierPooler(
209 - vllm_config.model_config,
210 - pooling=pooler.pooling,
211 - classifier=self._classifier,
212 - act_fn=pooler.head.activation,
213 - )
184 + pooling_type_str = pooler_config.pooling_type
185 + pooling_type = (PoolingType.LAST if pooling_type_str is None else
186 + PoolingType[pooling_type_str])
187 +
188 + self.pooler = DispatchPooler({
189 + "encode":
190 + Pooler.for_encode(pooler_config),
191 + "classify":
192 + ClassifierPooler(
193 + pooling=PoolingMethod.from_pooling_type(pooling_type),
194 + classifier=self._classifier,
195 + act_fn=ClassifierPooler.act_fn_for_seq_cls(
196 + vllm_config.model_config),
197 + ),
198 + "score":
199 + ClassifierPooler(
200 + pooling=PoolingMethod.from_pooling_type(pooling_type),
201 + classifier=self._classifier,
202 + act_fn=ClassifierPooler.act_fn_for_cross_encoder(
203 + vllm_config.model_config),
204 + ),
205 + })
206
207 def _classifier(self, x: torch.Tensor):
208 x, _ = self.score(x.float())
@@ -259,14 +251,16 @@ def as_reward_model(cls: _T) -> _T:
251 return cls
252
253 # Lazy import
262 - from vllm.model_executor.layers.pooler import PoolingType
263 -
264 - ModelForReward = _create_pooling_model_cls(
265 - cls,
266 - default_pooling_type=PoolingType.ALL,
267 - default_normalize=False,
268 - default_softmax=False,
269 - )
254 + from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
255 +
256 + class ModelForReward(_create_pooling_model_cls(cls)):
257 +
258 + def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
259 + pooler_config = vllm_config.model_config.pooler_config
260 + assert pooler_config is not None
261 +
262 + self.pooler = DispatchPooler(
263 + {"encode": Pooler.for_encode(pooler_config)}, )
264
265 ModelForReward.__name__ = \
266 _get_pooling_model_name(cls.__name__, "ForReward")
vllm/model_executor/models/bert.py
+99 -33
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from collections.abc import Iterable
4 +from collections.abc import Iterable, Set
5 from typing import Optional, Union
6
7 import torch
@@ -17,7 +17,8 @@ from vllm.model_executor.layers.activation import get_act_fn
17 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
18 QKVParallelLinear,
19 RowParallelLinear)
20 -from vllm.model_executor.layers.pooler import (ClassifierPooler, Pooler,
20 +from vllm.model_executor.layers.pooler import (ClassifierPooler,
21 + DispatchPooler, Pooler,
22 PoolingMethod,
23 PoolingParamsUpdate,
24 PoolingType)
@@ -92,20 +93,29 @@ class BertPooler(Pooler):
93 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
94 self.activation = nn.Tanh()
95
95 - def get_pooling_updates(
96 - self,
97 - task: PoolingTask,
98 - ) -> Optional[PoolingParamsUpdate]:
96 + def get_supported_tasks(self) -> Set[PoolingTask]:
97 + return self.pooling.get_supported_tasks()
98 +
99 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
100 return self.pooling.get_pooling_updates(task)
101
102 + def _head(self, pooled_output: torch.Tensor):
103 + pooled_output = self.dense(pooled_output)
104 + pooled_output = self.activation(pooled_output)
105 + return pooled_output
106 +
107 def forward(
108 self,
109 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
110 pooling_metadata: PoolingMetadata,
111 ) -> Union[torch.Tensor, list[torch.Tensor]]:
112 pooled_output = self.pooling(hidden_states, pooling_metadata)
107 - pooled_output = self.dense(pooled_output)
108 - pooled_output = self.activation(pooled_output)
113 +
114 + if isinstance(pooled_output, list):
115 + pooled_output = [self._head(output) for output in pooled_output]
116 + else:
117 + pooled_output = self._head(pooled_output)
118 +
119 return pooled_output
120
121
@@ -333,18 +343,19 @@ class BertModel(nn.Module, SupportsQuant):
343
344 packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}
345
336 - def __init__(self,
337 - *,
338 - vllm_config: VllmConfig,
339 - prefix: str = "",
340 - embedding_class: type = BertEmbedding,
341 - add_pooling_layer: bool = False):
346 + def __init__(
347 + self,
348 + *,
349 + vllm_config: VllmConfig,
350 + prefix: str = "",
351 + embedding_class: type[nn.Module] = BertEmbedding,
352 + ) -> None:
353 super().__init__()
354 +
355 config = vllm_config.model_config.hf_config
356 self.embeddings = embedding_class(config)
357 self.encoder = BertEncoder(vllm_config=vllm_config,
358 prefix=f"{prefix}.encoder")
347 - self.pooler = BertPooler(config) if add_pooling_layer else None
359
360 def forward(
361 self,
@@ -366,8 +377,7 @@ class BertModel(nn.Module, SupportsQuant):
377 token_type_ids=token_type_ids)
378 return self.encoder(hidden_states)
379
369 - def load_weights(self, weights: Iterable[tuple[str,
370 - torch.Tensor]]) -> set[str]:
380 + def _load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
381 stacked_params_mapping = [
382 # (param_name, shard_name, shard_id)
383 ("qkv_proj", "query", "q"),
@@ -395,10 +405,43 @@ class BertModel(nn.Module, SupportsQuant):
405 if name in params_dict:
406 other_weights.append((name, loaded_weight))
407
398 - loader = AutoWeightsLoader(
399 - self,
400 - skip_prefixes=(["pooler."] if self.pooler is None else []),
408 + return other_weights, loaded_stacked_params
409 +
410 + def load_weights(self, weights: Iterable[tuple[str,
411 + torch.Tensor]]) -> set[str]:
412 + other_weights, loaded_stacked_params = self._load_weights(weights)
413 +
414 + loader = AutoWeightsLoader(self, skip_prefixes=["pooler."])
415 + loaded_params = loader.load_weights(other_weights)
416 + loaded_params.update(loaded_stacked_params)
417 + return loaded_params
418 +
419 +
420 +class BertPoolingModel(BertModel):
421 +
422 + is_pooling_model = True
423 +
424 + def __init__(
425 + self,
426 + *,
427 + vllm_config: VllmConfig,
428 + prefix: str = "",
429 + embedding_class: type[nn.Module] = BertEmbedding,
430 + ) -> None:
431 + super().__init__(
432 + vllm_config=vllm_config,
433 + prefix=prefix,
434 + embedding_class=embedding_class,
435 )
436 +
437 + config = vllm_config.model_config.hf_config
438 + self.pooler = BertPooler(config)
439 +
440 + def load_weights(self, weights: Iterable[tuple[str,
441 + torch.Tensor]]) -> set[str]:
442 + other_weights, loaded_stacked_params = self._load_weights(weights)
443 +
444 + loader = AutoWeightsLoader(self)
445 loaded_params = loader.load_weights(other_weights)
446 loaded_params.update(loaded_stacked_params)
447 return loaded_params
@@ -421,6 +464,8 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
464 super().__init__()
465
466 pooler_config = vllm_config.model_config.pooler_config
467 + assert pooler_config is not None
468 +
469 self.model = self._build_model(vllm_config=vllm_config,
470 prefix=maybe_prefix(prefix, "model"))
471 self.pooler = self._build_pooler(pooler_config)
@@ -456,10 +501,15 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
501 embedding_class=BertEmbedding)
502
503 def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
459 - return Pooler.from_config_with_defaults(pooler_config,
460 - pooling_type=PoolingType.CLS,
461 - normalize=True,
462 - softmax=False)
504 + return DispatchPooler({
505 + "encode":
506 + Pooler.for_encode(pooler_config),
507 + "embed":
508 + Pooler.for_embed(
509 + pooler_config,
510 + default_pooling_type=PoolingType.CLS,
511 + ),
512 + })
513
514
515 class BertForSequenceClassification(nn.Module, SupportsV0Only,
@@ -481,16 +531,32 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,
531 config = vllm_config.model_config.hf_config
532
533 self.num_labels = config.num_labels
484 - self.bert = BertModel(vllm_config=vllm_config,
485 - prefix=maybe_prefix(prefix, "bert"),
486 - embedding_class=BertEmbedding,
487 - add_pooling_layer=True)
534 + self.bert = BertPoolingModel(vllm_config=vllm_config,
535 + prefix=maybe_prefix(prefix, "bert"),
536 + embedding_class=BertEmbedding)
537 self.classifier = nn.Linear(config.hidden_size, config.num_labels)
489 - self.pooler = ClassifierPooler(
490 - vllm_config.model_config,
491 - pooling=self.bert.pooler,
492 - classifier=self.classifier,
493 - )
538 +
539 + pooler_config = vllm_config.model_config.pooler_config
540 + assert pooler_config is not None
541 +
542 + self.pooler = DispatchPooler({
543 + "encode":
544 + Pooler.for_encode(pooler_config),
545 + "classify":
546 + ClassifierPooler(
547 + pooling=self.bert.pooler,
548 + classifier=self.classifier,
549 + act_fn=ClassifierPooler.act_fn_for_seq_cls(
550 + vllm_config.model_config),
551 + ),
552 + "score":
553 + ClassifierPooler(
554 + pooling=self.bert.pooler,
555 + classifier=self.classifier,
556 + act_fn=ClassifierPooler.act_fn_for_cross_encoder(
557 + vllm_config.model_config),
558 + ),
559 + })
560
561 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
562 loader = AutoWeightsLoader(self)
vllm/model_executor/models/gpt2.py
+10 -6
@@ -43,7 +43,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43 from vllm.model_executor.sampling_metadata import SamplingMetadata
44 from vllm.sequence import IntermediateTensors
45
46 -from ..layers.pooler import Pooler, PoolingType
46 +from ..layers.pooler import DispatchPooler, Pooler
47 from .interfaces import SupportsPP
48 from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
49 make_empty_intermediate_tensors_factory, make_layers,
@@ -339,12 +339,16 @@ class GPT2ForSequenceClassification(nn.Module):
339 self.transformer = GPT2Model(vllm_config=vllm_config,
340 prefix=maybe_prefix(prefix, "gpt2"))
341 self.score = nn.Linear(config.n_embd, config.num_labels, bias=False)
342 +
343 pooler_config = vllm_config.model_config.pooler_config
343 - self.pooler = Pooler.from_config_with_defaults(
344 - pooler_config,
345 - pooling_type=PoolingType.LAST,
346 - normalize=False,
347 - softmax=True)
344 + assert pooler_config is not None
345 +
346 + self.pooler = DispatchPooler({
347 + "encode":
348 + Pooler.for_encode(pooler_config),
349 + "classify":
350 + Pooler.for_classify(pooler_config, classifier=None),
351 + })
352
353 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
354 loader = AutoWeightsLoader(self)
vllm/model_executor/models/gritlm.py
+19 -20
@@ -1,17 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 +from collections.abc import Set
4 from typing import Optional, Union
5
6 import numpy as np
7 import torch
8 import torch.nn as nn
9 -from typing_extensions import assert_never
9
10 from vllm.config import ModelConfig, VllmConfig
11 from vllm.logger import init_logger
13 -from vllm.model_executor.layers.pooler import (Pooler, PoolerHead,
14 - PoolerNormalize,
12 +from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
13 + PoolerHead, PoolerNormalize,
14 PoolingParamsUpdate,
15 build_output, get_prompt_lens,
16 get_prompt_token_ids)
@@ -135,18 +134,11 @@ class GritLMMeanPool(nn.Module):
134
135 return instruction_len
136
138 - def get_pooling_updates(
139 - self,
140 - task: PoolingTask,
141 - ) -> Optional[PoolingParamsUpdate]:
142 - # The equalities are split up to keep mypy happy
143 - if task == "encode" or task == "embed":
144 - return PoolingParamsUpdate(requires_token_ids=True)
145 -
146 - if task == "classify" or task == "score":
147 - return None
137 + def get_supported_tasks(self) -> Set[PoolingTask]:
138 + return {"encode", "embed"}
139
149 - assert_never(task)
140 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
141 + return PoolingParamsUpdate(requires_token_ids=True)
142
143 def forward_one(
144 self,
@@ -207,10 +199,10 @@ class GritLMPooler(Pooler):
199 self.pooling = GritLMMeanPool(model_config)
200 self.head = PoolerHead(PoolerNormalize())
201
210 - def get_pooling_updates(
211 - self,
212 - task: PoolingTask,
213 - ) -> Optional[PoolingParamsUpdate]:
202 + def get_supported_tasks(self) -> Set[PoolingTask]:
203 + return self.pooling.get_supported_tasks()
204 +
205 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
206 return self.pooling.get_pooling_updates(task)
207
208 def forward(
@@ -262,4 +254,11 @@ class GritLM(LlamaForCausalLM, SupportsV0Only):
254
255 super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
256
265 - self.pooler = GritLMPooler(vllm_config.model_config)
257 + pooler_config = vllm_config.model_config.pooler_config
258 + if pooler_config is not None:
259 + self.pooler = DispatchPooler({
260 + "encode":
261 + Pooler.for_encode(pooler_config),
262 + "embed":
263 + GritLMPooler(vllm_config.model_config),
264 + })
vllm/model_executor/models/internlm2.py
+5 -7
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
22 QKVParallelLinear,
23 RowParallelLinear)
24 from vllm.model_executor.layers.logits_processor import LogitsProcessor
25 -from vllm.model_executor.layers.pooler import Pooler, PoolingType
25 +from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
26 from vllm.model_executor.layers.quantization import QuantizationConfig
27 from vllm.model_executor.layers.rotary_embedding import get_rope
28 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -429,12 +429,10 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):
429 )
430
431 pooler_config = vllm_config.model_config.pooler_config
432 - self.pooler = Pooler.from_config_with_defaults(
433 - pooler_config,
434 - pooling_type=PoolingType.ALL,
435 - normalize=False,
436 - softmax=False,
437 - )
432 + assert pooler_config is not None
433 +
434 + self.pooler = DispatchPooler(
435 + {"encode": Pooler.for_encode(pooler_config)}, )
436
437 def forward(
438 self,
vllm/model_executor/models/jamba.py
+14 -15
@@ -19,8 +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 (ClassifierPooler, PoolingType,
23 - SimplePooler)
22 +from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
23 + PoolingType)
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)
@@ -584,16 +584,15 @@ class JambaForSequenceClassification(JambaForCausalLM):
584 pooler_config = vllm_config.model_config.pooler_config
585 assert pooler_config is not None
586
587 - pooler = SimplePooler.from_config_with_defaults(
588 - pooler_config,
589 - pooling_type=PoolingType.LAST,
590 - normalize=False,
591 - softmax=False,
592 - )
593 -
594 - self.pooler = ClassifierPooler(
595 - vllm_config.model_config,
596 - pooling=pooler.pooling,
597 - classifier=self.score,
598 - act_fn=pooler.head.activation,
599 - )
587 + self.pooler = DispatchPooler({
588 + "encode":
589 + Pooler.for_encode(pooler_config),
590 + "classify":
591 + Pooler.for_classify(
592 + pooler_config,
593 + classifier=self.score,
594 + default_pooling_type=PoolingType.LAST,
595 + default_normalize=False,
596 + default_softmax=False,
597 + ),
598 + })
vllm/model_executor/models/jina_vl.py
+12 -6
@@ -12,7 +12,7 @@ from vllm.inputs import TokensPrompt
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
14 RowParallelLinear)
15 -from vllm.model_executor.layers.pooler import Pooler, PoolingType
15 +from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
16 from vllm.multimodal import MULTIMODAL_REGISTRY
17 from vllm.sequence import IntermediateTensors
18
@@ -96,11 +96,17 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,
96
97 self.score = JinaVLScorer(config)
98
99 - self.pooler = Pooler.from_config_with_defaults(
100 - pooler_config,
101 - pooling_type=PoolingType.LAST,
102 - normalize=False,
103 - softmax=True)
99 + pooler_config = vllm_config.model_config.pooler_config
100 + assert pooler_config is not None
101 +
102 + self.pooler = DispatchPooler({
103 + "encode":
104 + Pooler.for_encode(pooler_config),
105 + "classify":
106 + Pooler.for_classify(pooler_config, classifier=None),
107 + "score":
108 + Pooler.for_classify(pooler_config, classifier=None),
109 + })
110
111 @classmethod
112 def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
vllm/model_executor/models/modernbert.py
+38 -12
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Iterable
3 +from collections.abc import Iterable, Set
4 from typing import Optional, Union
5
6 import torch
@@ -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, Pooler,
16 +from vllm.model_executor.layers.pooler import (ClassifierPooler,
17 + DispatchPooler, Pooler,
18 PoolingMethod,
19 PoolingParamsUpdate,
20 PoolingType)
@@ -271,19 +272,27 @@ class ModernBertPooler(Pooler):
272 eps=config.norm_eps,
273 bias=config.norm_bias)
274
274 - def get_pooling_updates(
275 - self,
276 - task: PoolingTask,
277 - ) -> Optional[PoolingParamsUpdate]:
275 + def get_supported_tasks(self) -> Set[PoolingTask]:
276 + return self.pooling.get_supported_tasks()
277 +
278 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
279 return self.pooling.get_pooling_updates(task)
280
281 + def _head(self, pooled_output: torch.Tensor):
282 + return self.norm(self.act(self.dense(pooled_output)))
283 +
284 def forward(
285 self,
286 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
287 pooling_metadata: PoolingMetadata,
288 ) -> Union[torch.Tensor, list[torch.Tensor]]:
289 pooled_output = self.pooling(hidden_states, pooling_metadata)
286 - pooled_output = self.norm(self.act(self.dense(pooled_output)))
290 +
291 + if isinstance(pooled_output, list):
292 + pooled_output = [self._head(output) for output in pooled_output]
293 + else:
294 + pooled_output = self._head(pooled_output)
295 +
296 return pooled_output
297
298
@@ -299,11 +308,28 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,
308 self.model = ModernBertModel(vllm_config=vllm_config,
309 prefix=maybe_prefix(prefix, "modernbert"))
310 self.classifier = nn.Linear(config.hidden_size, config.num_labels)
302 - self.pooler = ClassifierPooler(
303 - vllm_config.model_config,
304 - pooling=ModernBertPooler(config),
305 - classifier=self.classifier,
306 - )
311 +
312 + pooler_config = vllm_config.model_config.pooler_config
313 + assert pooler_config is not None
314 +
315 + self.pooler = DispatchPooler({
316 + "encode":
317 + Pooler.for_encode(pooler_config),
318 + "classify":
319 + ClassifierPooler(
320 + pooling=ModernBertPooler(config),
321 + classifier=self.classifier,
322 + act_fn=ClassifierPooler.act_fn_for_seq_cls(
323 + vllm_config.model_config),
324 + ),
325 + "score":
326 + ClassifierPooler(
327 + pooling=ModernBertPooler(config),
328 + classifier=self.classifier,
329 + act_fn=ClassifierPooler.act_fn_for_cross_encoder(
330 + vllm_config.model_config),
331 + ),
332 + })
333
334 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
335
vllm/model_executor/models/qwen2_rm.py
+21 -14
@@ -15,7 +15,8 @@ from torch import nn
15 from vllm.config import VllmConfig
16 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17 RowParallelLinear)
18 -from vllm.model_executor.layers.pooler import Pooler, PoolingType, SimplePooler
18 +from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
19 + PoolingType)
20 from vllm.sequence import IntermediateTensors
21
22 from .interfaces import SupportsLoRA, SupportsPP
@@ -26,7 +27,7 @@ from .utils import AutoWeightsLoader, maybe_prefix
27 class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
28
29 is_pooling_model = True
29 - pooler: SimplePooler
30 + pooler: Pooler
31
32 packed_modules_mapping = {
33 "qkv_proj": [
@@ -94,12 +95,12 @@ class Qwen2ForRewardModel(Qwen2RewardBaseModel):
95 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
96 vllm_config.model_config.hf_config.num_labels = 1
97 super().__init__(vllm_config=vllm_config, prefix=prefix)
98 +
99 pooler_config = vllm_config.model_config.pooler_config
98 - self.pooler = Pooler.from_config_with_defaults(
99 - pooler_config,
100 - pooling_type=PoolingType.ALL,
101 - normalize=False,
102 - softmax=False)
100 + assert pooler_config is not None
101 +
102 + self.pooler = DispatchPooler(
103 + {"encode": Pooler.for_encode(pooler_config)}, )
104
105
106 class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
@@ -107,11 +108,17 @@ class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
108 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
109 vllm_config.model_config.hf_config.num_labels = 2
110 super().__init__(vllm_config=vllm_config, prefix=prefix)
111 +
112 pooler_config = vllm_config.model_config.pooler_config
111 - self.pooler = Pooler.from_config_with_defaults(
112 - pooler_config,
113 - pooling_type=PoolingType.STEP,
114 - normalize=False,
115 - softmax=True,
116 - step_tag_id=151651,
117 - )
113 + assert pooler_config is not None
114 +
115 + self.pooler = DispatchPooler({
116 + "encode":
117 + Pooler.for_encode(
118 + pooler_config,
119 + default_pooling_type=PoolingType.STEP,
120 + default_normalize=False,
121 + default_softmax=True,
122 + default_step_tag_id=151651,
123 + )
124 + })
vllm/model_executor/models/roberta.py
+27 -17
@@ -9,7 +9,8 @@ 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, CLSPool
12 +from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,
13 + DispatchPooler, Pooler)
14 from vllm.model_executor.layers.vocab_parallel_embedding import (
15 VocabParallelEmbedding)
16 from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel
@@ -63,16 +64,10 @@ class RobertaEmbedding(nn.Module):
64 # References:
65 # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L133
66 # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L1669
66 - pos_list = []
67 - token_list = []
68 - offset = 0
69 - for seq_len in seq_lens:
70 - pos_list.append(position_ids[offset:offset + seq_len])
71 - token_list.append(input_ids[offset:offset + seq_len])
72 - offset += seq_len
73 -
67 + seq_lens_list = seq_lens.tolist()
68 new_pos_list = []
75 - for positions, tokens in zip(pos_list, token_list):
69 + for positions, tokens in zip(position_ids.split(seq_lens_list),
70 + input_ids.split(seq_lens_list)):
71 # Verify assumption that incoming position are
72 # always a sequence from 0 to N.
73 expected_pos = torch.arange(positions.size()[0],
@@ -184,15 +179,30 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
179 self.num_labels = config.num_labels
180 self.roberta = BertModel(vllm_config=vllm_config,
181 prefix=maybe_prefix(prefix, "bert"),
187 - embedding_class=RobertaEmbedding,
188 - add_pooling_layer=False)
182 + embedding_class=RobertaEmbedding)
183 self.classifier = RobertaClassificationHead(config)
184
191 - self.pooler = ClassifierPooler(
192 - vllm_config.model_config,
193 - pooling=CLSPool(),
194 - classifier=self.classifier,
195 - )
185 + pooler_config = vllm_config.model_config.pooler_config
186 + assert pooler_config is not None
187 +
188 + self.pooler = DispatchPooler({
189 + "encode":
190 + Pooler.for_encode(pooler_config),
191 + "classify":
192 + ClassifierPooler(
193 + pooling=CLSPool(),
194 + classifier=self.classifier,
195 + act_fn=ClassifierPooler.act_fn_for_seq_cls(
196 + vllm_config.model_config),
197 + ),
198 + "score":
199 + ClassifierPooler(
200 + pooling=CLSPool(),
201 + classifier=self.classifier,
202 + act_fn=ClassifierPooler.act_fn_for_cross_encoder(
203 + vllm_config.model_config),
204 + ),
205 + })
206
207 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
208 loader = AutoWeightsLoader(self)
vllm/model_executor/pooling_metadata.py
+7
@@ -38,6 +38,13 @@ class PoolingMetadata:
38 f"seq_data={self.seq_data}, "
39 f"prompt_lens={self.prompt_lens})")
40
41 + def __getitem__(self, indices: slice):
42 + return PoolingMetadata(
43 + seq_groups=self.seq_groups[indices],
44 + seq_data=dict(list(self.seq_data.items())[indices]),
45 + prompt_lens=self.prompt_lens[indices],
46 + )
47 +
48
49 @dataclass
50 class PoolingTensors:
vllm/v1/pool/metadata.py
+8
@@ -15,3 +15,11 @@ class PoolingMetadata:
15 prompt_lens: torch.Tensor
16 prompt_token_ids: Optional[torch.Tensor]
17 pooling_params: list[PoolingParams]
18 +
19 + def __getitem__(self, indices: slice):
20 + return PoolingMetadata(
21 + prompt_lens=self.prompt_lens[indices],
22 + prompt_token_ids=None if self.prompt_token_ids is None else
23 + self.prompt_token_ids[indices],
24 + pooling_params=self.pooling_params[indices],
25 + )
vllm/v1/worker/gpu_model_runner.py
+4 -12
@@ -5,7 +5,7 @@ import copy
5 import gc
6 import time
7 from contextlib import contextmanager
8 -from typing import TYPE_CHECKING, Any, Optional, Union, cast, get_args
8 +from typing import TYPE_CHECKING, Any, Optional, Union, cast
9
10 import numpy as np
11 import torch
@@ -415,15 +415,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
415 generator = None
416
417 if pooling_params:
418 - assert pooling_params.task is not None, (
418 + assert (task := pooling_params.task) is not None, (
419 "You did not set `task` in the API")
420
421 model = cast(VllmModelForPooling, self.model)
422 - to_update = (model.pooler.get_pooling_updates(
423 - pooling_params.task))
424 - assert to_update is not None, (
425 - f"{pooling_params.task=} is not supported by the model")
426 -
422 + to_update = model.pooler.get_pooling_updates(task)
423 to_update.apply(pooling_params)
424
425 self.requests[req_id] = CachedRequestState(
@@ -1122,10 +1118,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1118 if not is_pooling_model(model):
1119 return []
1120
1125 - return [
1126 - task for task in get_args(PoolingTask)
1127 - if model.pooler.get_pooling_updates(task)
1128 - ]
1121 + return list(model.pooler.get_supported_tasks())
1122
1123 def apply_grammar_bitmask(
1124 self,
@@ -2247,7 +2240,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2240 dummy_pooling_params = PoolingParams(task=dummy_task)
2241
2242 to_update = model.pooler.get_pooling_updates(dummy_task)
2250 - assert to_update is not None
2243 to_update.apply(dummy_pooling_params)
2244
2245 dummy_metadata = PoolingMetadata(
vllm/v1/worker/tpu_model_runner.py
+2 -5
@@ -3,7 +3,7 @@
3 import bisect
4 import gc
5 import time
6 -from typing import TYPE_CHECKING, Any, Optional, cast, get_args
6 +from typing import TYPE_CHECKING, Any, Optional, cast
7 from unittest.mock import patch
8
9 import numpy as np
@@ -491,10 +491,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
491 if not is_pooling_model(model):
492 return []
493
494 - return [
495 - task for task in get_args(PoolingTask)
496 - if model.pooler.get_pooling_updates(task)
497 - ]
494 + return list(model.pooler.get_supported_tasks())
495
496 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
497 """
vllm/worker/model_runner_base.py
+2 -5
@@ -4,7 +4,7 @@
4 import dataclasses
5 from abc import ABC, abstractmethod
6 from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Type,
7 - TypeVar, get_args)
7 + TypeVar)
8
9 import torch
10 import torch.nn as nn
@@ -230,10 +230,7 @@ class ModelRunnerBase(ABC, Generic[T]):
230 if not is_pooling_model(model):
231 return []
232
233 - return [
234 - task for task in get_args(PoolingTask)
235 - if model.pooler.get_pooling_updates(task)
236 - ]
233 + return list(model.pooler.get_supported_tasks())
234
235 def execute_model(
236 self,
vllm/worker/pooling_model_runner.py
+3 -7
@@ -199,15 +199,11 @@ class PoolingModelRunner(
199
200 pooling_params = seq_group_metadata.pooling_params
201 assert pooling_params is not None
202 - assert pooling_params.task is not None, (
202 + assert (task := pooling_params.task) is not None, (
203 "You did not set `task` in the API")
204
205 - to_update = (cast(VllmModelForPooling,
206 - self.model).pooler.get_pooling_updates(
207 - pooling_params.task))
208 - assert to_update is not None, (
209 - f"{pooling_params.task=} is not supported by the model")
210 -
205 + model = cast(VllmModelForPooling, self.model)
206 + to_update = model.pooler.get_pooling_updates(task)
207 to_update.apply(pooling_params)
208
209 seq_groups.append((seq_ids, pooling_params))