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 11, 21:15 UTC
[Model][7/N] Improve all pooling task | Deprecation as_reward_model. Extract hidden states prefer using new multi-vector retrieval API (#26686)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Dec 8, 2025 at 16:10 UTC
9e77ffca3f41e0e73879098f1686a4c82b8619d9
8 files changed
+88
-58
docs/models/pooling_models.md
+8
-2
@@ -33,8 +33,8 @@ shown in the table below.
33
| Architecture | `--convert` | Supported pooling tasks |
34
|-------------------------------------------------|-------------|---------------------------------------|
35
| `*ForTextEncoding`, `*EmbeddingModel`, `*Model` | `embed` | `token_embed`, `embed` |
36
+| `*ForRewardModeling`, `*RewardModel` | `embed` | `token_embed`, `embed` |
37
| `*For*Classification`, `*ClassificationModel` | `classify` | `token_classify`, `classify`, `score` |
37
-| `*ForRewardModeling`, `*RewardModel` | `reward` | `token_classify` |
38
39
!!! tip
40
You can explicitly set `--convert <type>` to specify how to convert the model.
@@ -70,7 +70,6 @@ the pooler assigned to each task has the following attributes by default:
70
71
| Task | Pooling Type | Normalization | Softmax |
72
|------------|--------------|---------------|---------|
73
-| `reward` | `ALL` | ❌ | ❌ |
73
| `embed` | `LAST` | ✅︎ | ❌ |
74
| `classify` | `LAST` | ❌ | ✅︎ |
75
@@ -318,3 +317,10 @@ We have split the `encode` task into two more specific token-wise tasks: `token_
317
### Remove softmax from PoolingParams
318
319
We are going to remove `softmax` and `activation` from `PoolingParams`. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function.
320
+
321
+### as_reward_model
322
+
323
+Pooling models now default support all pooling, you can use it without any settings.
324
+
325
+- Extracting hidden states prefers using `token_embed` task.
326
+- Reward models prefers using `token_classify` task.
docs/models/supported_models.md
+1
-8
@@ -581,16 +581,9 @@ These models primarily support the [`LLM.reward`](./pooling_models.md#llmreward)
581
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
582
|--------------|--------|-------------------|----------------------|---------------------------|
583
| `InternLM2ForRewardModel` | InternLM2-based | `internlm/internlm2-1_8b-reward`, `internlm/internlm2-7b-reward`, etc. | ✅︎ | ✅︎ |
584
-| `LlamaForCausalLM`<sup>C</sup> | Llama-based | `peiyi9979/math-shepherd-mistral-7b-prm`, etc. | ✅︎ | ✅︎ |
584
+| `LlamaForCausalLM` | Llama-based | `peiyi9979/math-shepherd-mistral-7b-prm`, etc. | ✅︎ | ✅︎ |
585
| `Qwen2ForRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-RM-72B`, etc. | ✅︎ | ✅︎ |
586
| `Qwen2ForProcessRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-PRM-7B`, etc. | ✅︎ | ✅︎ |
587
-| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
588
-
589
-<sup>C</sup> Automatically converted into a reward model via `--convert reward`. ([details](./pooling_models.md#model-conversion))
590
-\* Feature support is the same as that of the original model.
591
-
592
-If your model is not in the above list, we will try to automatically convert the model using
593
-[as_reward_model][vllm.model_executor.models.adapters.as_reward_model]. By default, we return the hidden states of each token directly.
587
588
!!! important
589
For process-supervised reward models such as `peiyi9979/math-shepherd-mistral-7b-prm`, the pooling config should be set explicitly,
examples/pooling/token_embed/jina_embeddings_v4.py
new
+71
@@ -0,0 +1,71 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm import LLM
7
+from vllm.inputs.data import TextPrompt
8
+from vllm.multimodal.utils import fetch_image
9
+
10
+# Initialize model
11
+model = LLM(
12
+ model="jinaai/jina-embeddings-v4-vllm-text-matching",
13
+ runner="pooling",
14
+ max_model_len=1024,
15
+ gpu_memory_utilization=0.8,
16
+)
17
+
18
+# Create text prompts
19
+text1 = "Ein wunderschöner Sonnenuntergang am Strand"
20
+text1_prompt = TextPrompt(prompt=f"Query: {text1}")
21
+
22
+text2 = "浜辺に沈む美しい夕日"
23
+text2_prompt = TextPrompt(prompt=f"Query: {text2}")
24
+
25
+# Create image prompt
26
+image = fetch_image(
27
+ "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
28
+)
29
+image_prompt = TextPrompt(
30
+ prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
31
+ multi_modal_data={"image": image},
32
+)
33
+
34
+# Encode all prompts
35
+prompts = [text1_prompt, text2_prompt, image_prompt]
36
+outputs = model.encode(prompts, pooling_task="token_embed")
37
+
38
+
39
+def get_embeddings(outputs):
40
+ VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
41
+
42
+ embeddings = []
43
+ for output in outputs:
44
+ if VISION_START_TOKEN_ID in output.prompt_token_ids:
45
+ # Gather only vision tokens
46
+ img_start_pos = torch.where(
47
+ torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
48
+ )[0][0]
49
+ img_end_pos = torch.where(
50
+ torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
51
+ )[0][0]
52
+ embeddings_tensor = output.outputs.data.detach().clone()[
53
+ img_start_pos : img_end_pos + 1
54
+ ]
55
+ else:
56
+ # Use all tokens for text-only prompts
57
+ embeddings_tensor = output.outputs.data.detach().clone()
58
+
59
+ # Pool and normalize embeddings
60
+ pooled_output = (
61
+ embeddings_tensor.sum(dim=0, dtype=torch.float32)
62
+ / embeddings_tensor.shape[0]
63
+ )
64
+ embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
65
+ return embeddings
66
+
67
+
68
+embeddings = get_embeddings(outputs)
69
+
70
+for embedding in embeddings:
71
+ print(embedding.shape)
tests/models/test_registry.py
-2
@@ -13,7 +13,6 @@ from vllm.model_executor.models import (
13
)
14
from vllm.model_executor.models.adapters import (
15
as_embedding_model,
16
- as_reward_model,
16
as_seq_cls_model,
17
)
18
from vllm.model_executor.models.registry import (
@@ -46,7 +45,6 @@ def test_registry_imports(model_arch):
45
# All vLLM models should be convertible to a pooling model
46
assert is_pooling_model(as_seq_cls_model(model_cls))
47
assert is_pooling_model(as_embedding_model(model_cls))
49
- assert is_pooling_model(as_reward_model(model_cls))
48
49
if model_arch in _MULTIMODAL_MODELS:
50
assert supports_multimodal(model_cls)
tests/test_config.py
+1
-1
@@ -97,7 +97,7 @@ def test_update_config():
97
("intfloat/multilingual-e5-small", "pooling", "none", "embed"),
98
("jason9693/Qwen2.5-1.5B-apeach", "pooling", "classify", "classify"),
99
("cross-encoder/ms-marco-MiniLM-L-6-v2", "pooling", "none", "classify"),
100
- ("Qwen/Qwen2.5-Math-RM-72B", "pooling", "none", "reward"),
100
+ ("Qwen/Qwen2.5-Math-RM-72B", "pooling", "none", "embed"),
101
("openai/whisper-small", "generate", "none", "transcription"),
102
],
103
)
vllm/config/model.py
+7
-3
@@ -516,7 +516,11 @@ class ModelConfig:
516
if task == "classify":
517
return "classify"
518
if task == "reward":
519
- return "reward"
519
+ logger.warning(
520
+ "Pooling models now default support all pooling; "
521
+ "you can use it without any settings."
522
+ )
523
+ return "embed"
524
if task == "score":
525
new_task = self._get_default_pooling_task(architectures)
526
return "classify" if new_task == "classify" else "embed"
@@ -1899,8 +1903,8 @@ _SUFFIX_TO_DEFAULTS: list[tuple[str, tuple[RunnerType, ConvertType]]] = [
1903
("ForImageClassification", ("pooling", "classify")),
1904
("ForVideoClassification", ("pooling", "classify")),
1905
("ClassificationModel", ("pooling", "classify")),
1902
- ("ForRewardModeling", ("pooling", "reward")),
1903
- ("RewardModel", ("pooling", "reward")),
1906
+ ("ForRewardModeling", ("pooling", "embed")),
1907
+ ("RewardModel", ("pooling", "embed")),
1908
# Let other `*Model`s take priority
1909
("Model", ("pooling", "embed")),
1910
]
vllm/model_executor/model_loader/utils.py
-4
@@ -167,7 +167,6 @@ _MODEL_ARCH_BY_HASH = dict[int, tuple[type[nn.Module], str]]()
167
def _get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module], str]:
168
from vllm.model_executor.models.adapters import (
169
as_embedding_model,
170
- as_reward_model,
170
as_seq_cls_model,
171
try_create_mm_pooling_model_cls,
172
)
@@ -207,9 +206,6 @@ def _get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module],
206
elif convert_type == "classify":
207
logger.debug_once("Converting to sequence classification model.")
208
model_cls = as_seq_cls_model(model_cls)
210
- elif convert_type == "reward":
211
- logger.debug_once("Converting to reward model.")
212
- model_cls = as_reward_model(model_cls)
209
else:
210
assert_never(convert_type)
211
vllm/model_executor/models/adapters.py
-38
@@ -346,44 +346,6 @@ def as_seq_cls_model(cls: _T) -> _T:
346
return ModelForSequenceClassification # type: ignore
347
348
349
-def as_reward_model(cls: _T) -> _T:
350
- """
351
- Subclass an existing vLLM model to support reward modeling.
352
-
353
- By default, we return the hidden states of each token directly.
354
-
355
- Note:
356
- We assume that no extra layers are added to the original model;
357
- please implement your own model if this is not the case.
358
- """
359
- # Avoid modifying existing reward models
360
- if is_pooling_model(cls):
361
- return cls
362
-
363
- # Lazy import
364
- from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
365
-
366
- from .interfaces_base import default_pooling_type
367
-
368
- @default_pooling_type("ALL")
369
- class ModelForReward(_create_pooling_model_cls(cls)):
370
- def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
371
- pooler_config = vllm_config.model_config.pooler_config
372
- assert pooler_config is not None
373
-
374
- self.pooler = DispatchPooler(
375
- {
376
- "token_classify": Pooler.for_token_classify(
377
- pooler_config=pooler_config
378
- )
379
- }
380
- )
381
-
382
- ModelForReward.__name__ = _get_pooling_model_name(cls.__name__, "ForReward")
383
-
384
- return ModelForReward # type: ignore
385
-
386
-
349
class SequenceClassificationConfig(VerifyAndUpdateConfig):
350
@staticmethod
351
def verify_and_update_config(vllm_config: "VllmConfig") -> None: