@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c18b3b8e8b

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, 03:30 UTC

[Bugfix] Add `use_cross_encoder` flag to use correct activation in `ClassifierPooler` (#20527)

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

Cyrus Leung committed Jul 7, 2025 at 05:01 UTC c18b3b8e8bdcebaa150311d2c4911a6428480162

8 files changed +56 -41

vllm/entrypoints/llm.py

+1 -1

@@ -1204,7 +1204,7 @@ class LLM:

1204

1205 input_pairs = [(t1, t2) for t1, t2 in zip(text_1, text_2)]

1206

1207 - pooling_params = PoolingParams()

1207 + pooling_params = PoolingParams(use_cross_encoder=True)

1208

1209 tokenization_kwargs: dict[str, Any] = {}

1210 _validate_truncation_size(self.llm_engine.model_config.max_model_len,

vllm/entrypoints/openai/protocol.py

+6 -4

@@ -1156,8 +1156,9 @@ class ScoreRequest(OpenAIBaseModel):

1156

1157 # --8<-- [end:score-extra-params]

1158

1159 - def to_pooling_params(self):

1160 - return PoolingParams(additional_data=self.additional_data)

1159 + def to_pooling_params(self, *, use_cross_encoder: bool = False):

1160 + return PoolingParams(use_cross_encoder=use_cross_encoder,

1161 + additional_data=self.additional_data)

1162

1163

1164 class RerankRequest(OpenAIBaseModel):

@@ -1182,8 +1183,9 @@ class RerankRequest(OpenAIBaseModel):

1183

1184 # --8<-- [end:rerank-extra-params]

1185

1185 - def to_pooling_params(self):

1186 - return PoolingParams(additional_data=self.additional_data)

1186 + def to_pooling_params(self, *, use_cross_encoder: bool = False):

1187 + return PoolingParams(use_cross_encoder=use_cross_encoder,

1188 + additional_data=self.additional_data)

1189

1190

1191 class RerankDocument(BaseModel):

vllm/entrypoints/openai/serving_score.py

+4 -6

@@ -25,9 +25,7 @@ from vllm.logger import init_logger

25 from vllm.lora.request import LoRARequest

26 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput

27 from vllm.prompt_adapter.request import PromptAdapterRequest

28 -from vllm.transformers_utils.tokenizer import (AnyTokenizer, MistralTokenizer,

29 - PreTrainedTokenizer,

30 - PreTrainedTokenizerFast)

28 +from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer

29 from vllm.utils import make_async, merge_async_iterators

31 logger = init_logger(__name__)

@@ -50,7 +48,7 @@ class ServingScores(OpenAIServing):

49 async def _embedding_score(

50 self,

53 - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

51 + tokenizer: AnyTokenizer,

52 texts_1: list[str],

53 texts_2: list[str],

54 request: Union[RerankRequest, ScoreRequest],

@@ -141,7 +139,7 @@ class ServingScores(OpenAIServing):

139

140 async def _cross_encoding_score(

141 self,

144 - tokenizer: Union[AnyTokenizer],

142 + tokenizer: AnyTokenizer,

143 texts_1: list[str],

144 texts_2: list[str],

145 request: Union[RerankRequest, ScoreRequest],

@@ -190,7 +188,7 @@ class ServingScores(OpenAIServing):

188 # Schedule the request and get the result generator.

189 generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []

190

193 - pooling_params = request.to_pooling_params()

191 + pooling_params = request.to_pooling_params(use_cross_encoder=True)

192

193 for i, engine_prompt in enumerate(engine_prompts):

194 request_id_item = f"{request_id}-{i}"

vllm/model_executor/layers/pooler.py

+31 -11

@@ -15,6 +15,7 @@ from vllm.model_executor.pooling_metadata import ( # noqa: E501

15 from vllm.model_executor.pooling_metadata import PoolingTensors

16 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput

17 from vllm.transformers_utils.config import (

18 + get_classification_activation_function,

19 get_cross_encoder_activation_function)

20 from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata

@@ -388,15 +389,14 @@ class ClassifierPooler(nn.Module):

389 self.classifier = classifier

390 self.pooler = pooler

391

391 - if config.task == "score":

392 - self.default_activation_function = \

393 - get_cross_encoder_activation_function(config.hf_config)

394 - elif config.task == "classify":

395 - self.default_activation_function = nn.Sigmoid() \

396 - if config.hf_config.num_labels == 1 else nn.Softmax()

397 - else:

398 - raise NotImplementedError(f"task={config.task!r} is not supported"

399 - " with the classification pooler")

392 + self.classification_act_fn = get_classification_activation_function(

393 + config.hf_config)

394 + self.cross_encoder_act_fn = get_cross_encoder_activation_function(

395 + config.hf_config)

396 +

397 + def _get_act_fn(self, use_cross_encoder: bool):

398 + return (self.cross_encoder_act_fn

399 + if use_cross_encoder else self.classification_act_fn)

400

401 def get_prompt_lens(

402 self,

@@ -446,8 +446,28 @@ class ClassifierPooler(nn.Module):

446 # apply classifier once on the full batch if possible

447 pooled_output = self.classifier(pooled_output)

448

449 - # shape: (batch_size, num_labels)

450 - scores = self.default_activation_function(pooled_output)

449 + if isinstance(pooling_metadata, V0PoolingMetadata):

450 + use_cross_encoder_list = [

451 + pooling_param.use_cross_encoder

452 + for _, pooling_param in pooling_metadata.seq_groups

453 + ]

454 + else:

455 + use_cross_encoder_list = [

456 + pooling_param.use_cross_encoder

457 + for pooling_param in pooling_metadata.pooling_params

458 + ]

459 +

460 + # shape of scores: (batch_size, num_labels)

461 + if all(use_cross_encoder == use_cross_encoder_list[0]

462 + for use_cross_encoder in use_cross_encoder_list):

463 + act_fn = self._get_act_fn(use_cross_encoder_list[0])

464 + scores = act_fn(pooled_output)

465 + else:

466 + scores = torch.stack([

467 + self._get_act_fn(use_cross_encoder)(vecs)

468 + for use_cross_encoder, vecs in zip(use_cross_encoder_list,

469 + pooled_output)

470 + ])

471

472 pooled_outputs = [PoolingSequenceGroupOutput(data) for data in scores]

473 return PoolerOutput(outputs=pooled_outputs)

vllm/model_executor/models/bert.py

-5

@@ -25,8 +25,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

25 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

26 from vllm.model_executor.pooling_metadata import PoolingMetadata

27 from vllm.sequence import IntermediateTensors, PoolerOutput

28 -from vllm.transformers_utils.config import (

29 - get_cross_encoder_activation_function)

29 from .interfaces import SupportsCrossEncoding, SupportsQuant, SupportsV0Only

30 from .utils import WeightsMapper, maybe_prefix

@@ -462,9 +460,6 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,

460 super().__init__()

461 config = vllm_config.model_config.hf_config

462

465 - self.default_activation_function = \

466 - get_cross_encoder_activation_function(config)

467 -

463 self.num_labels = config.num_labels

464 self.bert = BertModel(vllm_config=vllm_config,

465 prefix=maybe_prefix(prefix, "bert"),

vllm/model_executor/models/roberta.py

-5

@@ -18,8 +18,6 @@ from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel

18 from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix

19 from vllm.model_executor.pooling_metadata import PoolingMetadata

20 from vllm.sequence import IntermediateTensors, PoolerOutput

21 -from vllm.transformers_utils.config import (

22 - get_cross_encoder_activation_function)

22 from .bert_with_rope import BertWithRope, JinaRobertaModel

23 from .interfaces import SupportsCrossEncoding, SupportsV0Only

@@ -178,9 +176,6 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,

176 super().__init__()

177 config = vllm_config.model_config.hf_config

178

181 - self.default_activation_function = \

182 - get_cross_encoder_activation_function(config)

183 -

179 self.num_labels = config.num_labels

180 self.roberta = BertModel(vllm_config=vllm_config,

181 prefix=maybe_prefix(prefix, "bert"),

vllm/pooling_params.py

@@ -24,12 +24,14 @@ class PoolingParams(

24 """

26 dimensions: Optional[int] = None

27 + use_cross_encoder: bool = False

28 additional_data: Optional[Any] = None

29 output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY

31 def clone(self) -> "PoolingParams":

32 """Returns a deep copy of the PoolingParams instance."""

33 return PoolingParams(dimensions=self.dimensions,

34 + use_cross_encoder=self.use_cross_encoder,

35 additional_data=self.additional_data)

37 def verify(self, model_config: "ModelConfig") -> None:

@@ -54,6 +56,7 @@ class PoolingParams(

56 def __repr__(self) -> str:

57 return (f"PoolingParams("

58 f"dimensions={self.dimensions}, "

59 + f"use_cross_encoder={self.use_cross_encoder}, "

60 f"additional_metadata={self.additional_data})")

62 def __post_init__(self) -> None:

vllm/transformers_utils/config.py

+11 -9

@@ -866,24 +866,26 @@ def try_get_generation_config(

866 return None

867

868

869 -def get_cross_encoder_activation_function(config: PretrainedConfig):

869 +def get_classification_activation_function(config: PretrainedConfig):

870 + return nn.Sigmoid() if config.num_labels == 1 else nn.Softmax()

871 +

872

873 +def get_cross_encoder_activation_function(config: PretrainedConfig):

874 function_name: Optional[str] = None

872 - if hasattr(config, "sentence_transformers") and "activation_fn" in \

873 - config.sentence_transformers:

875 + if (hasattr(config, "sentence_transformers")

876 + and "activation_fn" in config.sentence_transformers):

877 function_name = config.sentence_transformers["activation_fn"]

875 -

878 elif (hasattr(config, "sbert_ce_default_activation_function")

879 and config.sbert_ce_default_activation_function is not None):

880 function_name = config.sbert_ce_default_activation_function

881

882 if function_name is not None:

881 - assert function_name.startswith("torch.nn.modules."), \

882 - "Loading of activation functions is restricted to " \

883 - "torch.nn.modules for security reasons"

883 + assert function_name.startswith("torch.nn.modules."), (

884 + "Loading of activation functions is restricted to "

885 + "torch.nn.modules for security reasons")

886 return resolve_obj_by_qualname(function_name)()

885 - else:

886 - return nn.Sigmoid() if config.num_labels == 1 else nn.Identity()

887 +

888 + return nn.Sigmoid() if config.num_labels == 1 else nn.Identity()

889

890

891 def try_get_safetensors_metadata(