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, 08:30 UTC
Prevent the cross-encoder logic from being applied to classification tasks (#18838)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Maximilien de Bayser committed
May 28, 2025 at 23:16 UTC
515b413ebf8feb8093556dbef21a02f9f6f877cc
5 files changed
+40
-24
vllm/model_executor/layers/pooler.py
+20
-14
@@ -6,10 +6,9 @@ from typing import Optional, Union
6
import torch
7
import torch.nn as nn
8
import torch.nn.functional as F
9
-from transformers import PretrainedConfig
9
from typing_extensions import assert_never
10
12
-from vllm.config import PoolerConfig
11
+from vllm.config import ModelConfig, PoolerConfig
12
from vllm.model_executor.pooling_metadata import (PoolingMetadata,
13
PoolingTensors)
14
from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
@@ -283,30 +282,37 @@ class Pooler(nn.Module):
282
)
283
284
286
-class CrossEncodingPooler(nn.Module):
287
- """A layer that pools specific information from hidden states.
285
+class ClassifierPooler(nn.Module):
286
+ """A pooling layer for classification tasks.
287
288
This layer does the following:
290
- 1. Extracts specific tokens or aggregates data based on pooling method.
291
- 2. Normalizes output if specified.
292
- 3. Returns structured results as `PoolerOutput`.
293
-
294
- Attributes:
295
- pooling_type: The type of pooling to use.
296
- normalize: Whether to normalize the pooled data.
289
+ 1. Applies a classification layer to the hidden states.
290
+ 2. Optionally applies a pooler layer.
291
+ 3. Applies an activation function to the output. In the case of
292
+ classification models it is either sigmoid or softmax. In the
293
+ case of scoring models, the same behavior is configuration
294
+ dependent, as in the sentence-transformers library.
295
"""
296
297
def __init__(
298
self,
301
- config: PretrainedConfig,
299
+ config: ModelConfig,
300
classifier: nn.Module,
301
pooler: Optional[nn.Module] = None,
302
):
303
super().__init__()
304
self.classifier = classifier
305
self.pooler = pooler
308
- self.default_activation_function = \
309
- get_cross_encoder_activation_function(config)
306
+
307
+ if config.task == "score":
308
+ self.default_activation_function = \
309
+ get_cross_encoder_activation_function(config.hf_config)
310
+ elif config.task == "classify":
311
+ self.default_activation_function = nn.Sigmoid() \
312
+ if config.hf_config.num_labels == 1 else nn.Softmax()
313
+ else:
314
+ raise NotImplementedError(f"task={config.task!r} is not supported"
315
+ " with the classification pooler")
316
317
def forward(
318
self,
vllm/model_executor/models/bert.py
+3
-3
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.activation import get_act_fn
16
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17
QKVParallelLinear,
18
RowParallelLinear)
19
-from vllm.model_executor.layers.pooler import (CrossEncodingPooler, Pooler,
19
+from vllm.model_executor.layers.pooler import (ClassifierPooler, Pooler,
20
PoolingType)
21
from vllm.model_executor.layers.quantization import QuantizationConfig
22
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -470,8 +470,8 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,
470
embedding_class=BertEmbedding,
471
add_pooling_layer=True)
472
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
473
- self._pooler = CrossEncodingPooler(config, self.classifier,
474
- self.bert.pooler)
473
+ self._pooler = ClassifierPooler(vllm_config.model_config,
474
+ self.classifier, self.bert.pooler)
475
476
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
477
vllm/model_executor/models/modernbert.py
+4
-3
@@ -12,7 +12,7 @@ from vllm.config import VllmConfig
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.linear import (QKVParallelLinear,
14
RowParallelLinear)
15
-from vllm.model_executor.layers.pooler import CrossEncodingPooler
15
+from vllm.model_executor.layers.pooler import ClassifierPooler
16
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
17
from vllm.model_executor.layers.vocab_parallel_embedding import (
18
VocabParallelEmbedding)
@@ -278,8 +278,9 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
278
self.model = ModernBertModel(vllm_config=vllm_config,
279
prefix=maybe_prefix(prefix, "modernbert"))
280
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
281
- self._pooler = CrossEncodingPooler(config, self.classifier,
282
- ModernBertPooler(config))
281
+ self._pooler = ClassifierPooler(vllm_config.model_config,
282
+ self.classifier,
283
+ ModernBertPooler(config))
284
285
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
286
vllm/model_executor/models/roberta.py
+4
-2
@@ -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 CrossEncodingPooler
12
+from vllm.model_executor.layers.pooler import ClassifierPooler
13
from vllm.model_executor.layers.vocab_parallel_embedding import (
14
VocabParallelEmbedding)
15
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -186,7 +186,9 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
186
embedding_class=RobertaEmbedding,
187
add_pooling_layer=False)
188
self.classifier = RobertaClassificationHead(config)
189
- self._pooler = CrossEncodingPooler(config, self.classifier)
189
+
190
+ self._pooler = ClassifierPooler(vllm_config.model_config,
191
+ self.classifier)
192
193
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
194
bert_weights, task_weights = roberta_task_weights_filter(weights)
vllm/transformers_utils/config.py
+9
-2
@@ -823,10 +823,17 @@ def try_get_generation_config(
823
824
825
def get_cross_encoder_activation_function(config: PretrainedConfig):
826
- if (hasattr(config, "sbert_ce_default_activation_function")
827
- and config.sbert_ce_default_activation_function is not None):
826
827
+ function_name: Optional[str] = None
828
+ if hasattr(config, "sentence_transformers") and "activation_fn" in \
829
+ config.sentence_transformers:
830
+ function_name = config.sentence_transformers["activation_fn"]
831
+
832
+ elif (hasattr(config, "sbert_ce_default_activation_function")
833
+ and config.sbert_ce_default_activation_function is not None):
834
function_name = config.sbert_ce_default_activation_function
835
+
836
+ if function_name is not None:
837
assert function_name.startswith("torch.nn.modules."), \
838
"Loading of activation functions is restricted to " \
839
"torch.nn.modules for security reasons"