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 14, 05:30 UTC
[Model] Classification models support logit_bias / sigmoid_normalize (#24031)
Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
wang.yuqi committed
Sep 3, 2025 at 00:48 UTC
e0653f6c0b9f331af0877e7c7abc99a85efc3982
4 files changed
+38
-30
vllm/config/__init__.py
+24
-21
@@ -2651,24 +2651,46 @@ class PoolerConfig:
2651
## for embeddings models
2652
normalize: Optional[bool] = None
2653
"""
2654
- Whether to normalize the embeddings outputs.
2654
+ Whether to normalize the embeddings outputs. Defaults to True.
2655
"""
2656
dimensions: Optional[int] = None
2657
"""
2658
Reduce the dimensions of embeddings if model
2659
- support matryoshka representation.
2659
+ support matryoshka representation. Defaults to None.
2660
+ """
2661
+ enable_chunked_processing: Optional[bool] = None
2662
+ """
2663
+ Whether to enable chunked processing for long inputs that exceed the model's
2664
+ maximum position embeddings. When enabled, long inputs will be split into
2665
+ chunks, processed separately, and then aggregated using weighted averaging.
2666
+ This allows embedding models to handle arbitrarily long text without CUDA
2667
+ errors. Defaults to False.
2668
+ """
2669
+ max_embed_len: Optional[int] = None
2670
+ """
2671
+ Maximum input length allowed for embedding generation. When set, allows
2672
+ inputs longer than max_embed_len to be accepted for embedding models.
2673
+ When an input exceeds max_embed_len, it will be handled according to
2674
+ the original max_model_len validation logic.
2675
+ Defaults to None (i.e. set to max_model_len).
2676
"""
2677
2678
## for classification models
2679
activation: Optional[bool] = None
2680
"""
2681
Whether to apply activation function to the classification outputs.
2682
+ Defaults to True.
2683
+ """
2684
+ logit_bias: Optional[float] = None
2685
+ """
2686
+ If provided, apply classification logit biases. Defaults to None.
2687
"""
2688
2689
## for reward models
2690
softmax: Optional[bool] = None
2691
"""
2692
Whether to apply softmax to the reward outputs.
2693
+ Defaults to True.
2694
"""
2695
step_tag_id: Optional[int] = None
2696
"""
@@ -2683,25 +2705,6 @@ class PoolerConfig:
2705
``math-shepherd-mistral-7b-prm`` model.
2706
"""
2707
2686
- enable_chunked_processing: Optional[bool] = None
2687
- """
2688
- Whether to enable chunked processing for long inputs that exceed the model's
2689
- maximum position embeddings. When enabled, long inputs will be split into
2690
- chunks, processed separately, and then aggregated using weighted averaging.
2691
- This allows embedding models to handle arbitrarily long text without CUDA
2692
- errors. Defaults to False.
2693
- """
2694
-
2695
- max_embed_len: Optional[int] = None
2696
- """
2697
- Maximum input length allowed for embedding generation. When set, allows
2698
- inputs longer than max_embed_len to be accepted for embedding models.
2699
- This parameter enables accepting long inputs without requiring
2700
- VLLM_ALLOW_LONG_MAX_MODEL_LEN environment variable. When an input exceeds
2701
- max_embed_len, it will be handled according to the original max_model_len
2702
- validation logic. Defaults to None (i.e. set to max_model_len).
2703
- """
2704
-
2708
def compute_hash(self) -> str:
2709
"""
2710
WARNING: Whenever a new field is added to this config,
vllm/model_executor/layers/pooler.py
+8
@@ -633,9 +633,14 @@ class ClassifierPooler(Pooler):
633
) -> None:
634
super().__init__()
635
636
+ from vllm.config import get_current_vllm_config
637
+ vllm_config = get_current_vllm_config()
638
+
639
self.pooling = pooling
640
self.classifier = classifier
641
self.act_fn = act_fn or PoolerClassify()
642
+ self.logit_bias: Optional[
643
+ float] = vllm_config.model_config.pooler_config.logit_bias
644
645
def get_supported_tasks(self) -> Set[PoolingTask]:
646
return {"classify", "score"}
@@ -654,6 +659,9 @@ class ClassifierPooler(Pooler):
659
pooled_data = self.classifier(pooled_data)
660
# pooled_data shape: [batchsize, num_labels]
661
662
+ if self.logit_bias is not None:
663
+ pooled_data -= self.logit_bias
664
+
665
pooling_params = get_pooling_params(pooling_metadata)
666
flags = [p.activation for p in pooling_params]
667
vllm/model_executor/models/config.py
+3
-1
@@ -210,8 +210,10 @@ class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):
210
@staticmethod
211
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
212
config = vllm_config.model_config.hf_config
213
-
213
config.num_labels = 1
214
+ pooler_config = vllm_config.model_config.pooler_config
215
+ if pooler_config.logit_bias is None:
216
+ pooler_config.logit_bias = 2.65
217
218
219
class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
vllm/model_executor/models/jina_vl.py
+3
-8
@@ -92,17 +92,14 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,
92
pooler_config = vllm_config.model_config.pooler_config
93
assert pooler_config is not None
94
95
- # logit bias for sigmoid normalization
96
- self.LOGIT_BIAS = 2.65
97
-
95
self.score = JinaVLScorer(config)
96
self.pooler = DispatchPooler({
97
"encode":
98
Pooler.for_encode(pooler_config),
99
"classify":
103
- Pooler.for_classify(pooler_config, classifier=None),
100
+ Pooler.for_classify(pooler_config, classifier=self.score),
101
"score":
105
- Pooler.for_classify(pooler_config, classifier=None),
102
+ Pooler.for_classify(pooler_config, classifier=self.score),
103
})
104
105
@classmethod
@@ -137,9 +134,7 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,
134
inputs_embeds=inputs_embeds,
135
**kwargs,
136
)
140
-
141
- logits = self.score(hidden_states) - self.LOGIT_BIAS
142
- return logits
137
+ return hidden_states
138
139
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
140
loader = AutoWeightsLoader(self)