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

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

Fix pooling adapters for Transformers backend (#27338)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Oct 24, 2025 at 04:23 UTC 1f9460c4c19c18905260134bac730a51385ca675

6 files changed +97 -74

vllm/config/model.py

+11 -16

@@ -793,25 +793,20 @@ class ModelConfig:

793 cls += "MoE" if self.get_num_experts() > 1 else ""

794 # Check if the architecture we're wrapping has defaults

795 runner = None

796 - convert = None

796 + task = None

797 if defaults := try_match_architecture_defaults(self.architectures[0]):

798 - _, (runner, convert) = defaults

799 - # Overwrite with user-specified values

798 + _, (runner, task) = defaults

799 + # User specified value take precedence

800 if self.runner != "auto":

801 runner = self.runner

802 - if self.convert not in {"auto", "none"}:

803 - convert = self.convert

804 - # Fall back to default values if still not set

805 - if runner is None:

806 - runner = "generate"

807 - if convert in {None, "none"}:

808 - convert = "embed"

809 - # Resolve Transformers backend task

810 - if runner == "pooling":

811 - if convert == "embed":

812 - return cls + "EmbeddingModel"

813 - if convert == "classify":

814 - return cls + "ForSequenceClassification"

802 + # Only consider Transformers backend pooling classes if we're wrapping an

803 + # architecture that defaults to pooling. Otherwise, we return the LM class

804 + # and use adapters.

805 + if runner == "pooling" and task in {"embed", "classify"}:

806 + if task == "embed":

807 + cls += "EmbeddingModel"

808 + elif task == "classify":

809 + cls += "ForSequenceClassification"

810 else:

811 cls += "ForCausalLM"

812 return cls

vllm/model_executor/models/adapters.py

+49 -45

@@ -283,7 +283,6 @@ def as_seq_cls_model(cls: _T) -> _T:

283 Pooler,

284 )

285 from vllm.model_executor.models.interfaces import SupportsCrossEncoding

286 - from vllm.sequence import IntermediateTensors

286

287 from .utils import maybe_prefix

288

@@ -291,13 +290,13 @@ def as_seq_cls_model(cls: _T) -> _T:

290 _create_pooling_model_cls(cls), SupportsCrossEncoding

291 ):

292 def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):

294 - config = vllm_config.model_config.hf_config

293 + text_config = vllm_config.model_config.hf_config.get_text_config()

294 model_config = vllm_config.model_config

295 quant_config = vllm_config.quant_config

296

297 self.score = ReplicatedLinear(

298 model_config.hidden_size,

300 - config.num_labels,

299 + text_config.num_labels,

300 bias=False,

301 params_dtype=vllm_config.model_config.head_dtype,

302 quant_config=quant_config,

@@ -322,20 +321,10 @@ def as_seq_cls_model(cls: _T) -> _T:

321 }

322 )

323

325 - def forward(

326 - self,

327 - input_ids: torch.Tensor,

328 - positions: torch.Tensor,

329 - intermediate_tensors: IntermediateTensors | None = None,

330 - inputs_embeds: torch.Tensor | None = None,

331 - ) -> torch.Tensor:

332 - return super().forward(

333 - input_ids, positions, intermediate_tensors, inputs_embeds

334 - )

335 -

324 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

337 - tokens = getattr(self.config, "classifier_from_token", None)

338 - method = getattr(self.config, "method", None)

325 + text_config = self.config.get_text_config()

326 + tokens = getattr(text_config, "classifier_from_token", None)

327 + method = getattr(text_config, "method", None)

328

329 if tokens is None and method is None:

330 return super().load_weights(weights)

@@ -392,9 +381,9 @@ def as_reward_model(cls: _T) -> _T:

381 class SequenceClassificationConfig(VerifyAndUpdateConfig):

382 @staticmethod

383 def verify_and_update_config(vllm_config: "VllmConfig") -> None:

395 - config = vllm_config.model_config.hf_config

396 - method = getattr(config, "method", None)

397 - tokens = getattr(config, "classifier_from_token", None)

384 + text_config = vllm_config.model_config.hf_config.get_text_config()

385 + method = getattr(text_config, "method", None)

386 + tokens = getattr(text_config, "classifier_from_token", None)

387

388 if method is None:

389 return

@@ -404,13 +393,13 @@ class SequenceClassificationConfig(VerifyAndUpdateConfig):

393

394 if method == "from_2_way_softmax":

395 assert len(tokens) == 2

407 - config.num_labels = 1

396 + text_config.num_labels = 1

397 else:

409 - config.num_labels = len(tokens)

398 + text_config.num_labels = len(tokens)

399

400 # `llm as reranker` defaults to not using pad_token

412 - use_pad_token = getattr(config, "use_pad_token", False)

413 - config.use_pad_token = use_pad_token

401 + use_pad_token = getattr(text_config, "use_pad_token", False)

402 + text_config.use_pad_token = use_pad_token

403

404

405 def load_weights_using_from_2_way_softmax(

@@ -419,24 +408,31 @@ def load_weights_using_from_2_way_softmax(

408 # refer to https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3

409 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

410 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

422 - from vllm.model_executor.models.utils import AutoWeightsLoader

411

412 model_config = model.vllm_config.model_config

413 + quant_config = model.vllm_config.quant_config

414 + text_config = model.config.get_text_config()

415

426 - tokens = getattr(model.config, "classifier_from_token", [])

416 + tokens = getattr(text_config, "classifier_from_token", [])

417 tokens = cast(list[int], tokens)

418 assert len(tokens) == 2

419

430 - if model.config.tie_word_embeddings:

431 - model.lm_head = model.model.embed_tokens

432 - else:

433 - quant_config = model.vllm_config.quant_config

434 - model.lm_head = ParallelLMHead(

435 - model.config.vocab_size, model.config.hidden_size, quant_config=quant_config

420 + model.lm_head = ParallelLMHead(

421 + text_config.vocab_size, text_config.hidden_size, quant_config=quant_config

422 + )

423 + if text_config.tie_word_embeddings:

424 + # embed_tokens is the assumed name for input embeddings. If the model does not

425 + # have this attribute, we fallback to get_input_embeddings(), which is used by

426 + # the Transformers backend.

427 + embed_tokens = (

428 + model.model.embed_tokens

429 + if hasattr(model.model, "embed_tokens")

430 + else model.model.get_input_embeddings()

431 )

432 + model.lm_head = model.lm_head.tie_weights(embed_tokens)

433

438 - loader = AutoWeightsLoader(model)

439 - loaded_weights = loader.load_weights(weights)

434 + # Skip ModelForSequenceClassification in MRO to avoid infinite recursion

435 + loaded_weights = type(model).__mro__[1].load_weights(model, weights)

436

437 from vllm.transformers_utils.tokenizer import get_tokenizer

438

@@ -466,23 +462,31 @@ def load_weights_using_from_2_way_softmax(

462 def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Tensor]]):

463 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

464 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

469 - from vllm.model_executor.models.utils import AutoWeightsLoader

465

466 model_config = model.vllm_config.model_config

472 - tokens = getattr(model.config, "classifier_from_token", [])

467 + quant_config = model.vllm_config.quant_config

468 + text_config = model.config.get_text_config()

469 +

470 + tokens = getattr(text_config, "classifier_from_token", [])

471 tokens = cast(list[int], tokens)

472 assert len(tokens) > 0

473

476 - if model.config.tie_word_embeddings:

477 - model.lm_head = model.model.embed_tokens

478 - else:

479 - quant_config = model.vllm_config.quant_config

480 - model.lm_head = ParallelLMHead(

481 - model.config.vocab_size, model.config.hidden_size, quant_config=quant_config

474 + model.lm_head = ParallelLMHead(

475 + text_config.vocab_size, text_config.hidden_size, quant_config=quant_config

476 + )

477 + if text_config.tie_word_embeddings:

478 + # embed_tokens is the assumed name for input embeddings. If the model does not

479 + # have this attribute, we fallback to get_input_embeddings(), which is used by

480 + # the Transformers backend.

481 + embed_tokens = (

482 + model.model.embed_tokens

483 + if hasattr(model.model, "embed_tokens")

484 + else model.model.get_input_embeddings()

485 )

486 + model.lm_head = model.lm_head.tie_weights(embed_tokens)

487

484 - loader = AutoWeightsLoader(model)

485 - loaded_weights = loader.load_weights(weights)

488 + # Skip ModelForSequenceClassification in MRO to avoid infinite recursion

489 + loaded_weights = type(model).__mro__[1].load_weights(model, weights)

490

491 from vllm.transformers_utils.tokenizer import get_tokenizer

492

@@ -523,7 +527,7 @@ def seq_cls_model_loader(model, weights: Iterable[tuple[str, torch.Tensor]]):

527 # - GemmaForCausalLM

528 # - bge-reranker-v2-gemma

529

526 - config = model.vllm_config.model_config.hf_config

527 - method = getattr(config, "method", None)

530 + text_config = model.vllm_config.model_config.hf_config.get_text_config()

531 + method = getattr(text_config, "method", None)

532 assert method in SEQ_CLS_LOAD_METHODS, f"method {method} not supported"

533 return SEQ_CLS_LOAD_METHODS[method](model, weights)

vllm/model_executor/models/transformers/base.py

+22

@@ -49,6 +49,7 @@ from vllm.model_executor.models.transformers.utils import (

49 from vllm.model_executor.models.utils import (

50 AutoWeightsLoader,

51 PPMissingLayer,

52 + WeightsMapper,

53 make_empty_intermediate_tensors_factory,

54 maybe_prefix,

55 )

@@ -92,6 +93,27 @@ ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward

93 class Base(nn.Module, VllmModel, SupportsQuant, SupportsLoRA, SupportsPP):

94 embedding_padding_modules = ["lm_head"]

95 embedding_modules = ["embed_tokens"] # TODO transformers will have a util to get it

96 + hf_to_vllm_mapper = WeightsMapper(

97 + orig_to_new_prefix={

98 + # Add `model.` prefix for base model checkpoints,

99 + # handling the case where it is already present

100 + "": "model.",

101 + "model.model.": "model.",

102 + # Heads will be adjacent to `model` (pooling included because of adapters)

103 + "model.lm_head.": "lm_head.",

104 + "model.score.": "classifier.",

105 + "model.classifier.": "classifier.",

106 + }

107 + )

108 +

109 + def __init_subclass__(cls, *args, **kwargs):

110 + """Merge hf_to_vllm_mapper in MRO from most specific to least specific."""

111 + super().__init_subclass__(*args, **kwargs)

112 + hf_to_vllm_mapper = WeightsMapper()

113 + for base in cls.__mro__:

114 + if base_hf_to_vllm_mapper := getattr(base, "hf_to_vllm_mapper", None):

115 + hf_to_vllm_mapper |= base_hf_to_vllm_mapper

116 + cls.hf_to_vllm_mapper = hf_to_vllm_mapper

117

118 def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):

119 super().__init__()

vllm/model_executor/models/transformers/legacy.py

-7

@@ -34,13 +34,6 @@ class LegacyMixin:

34 # Handle BERT-like models

35 "roberta": "model",

36 "bert": "model",

37 - # Add `model.` prefix for base model checkpoints

38 - "": "model.",

39 - # Remove `model.` prefix if it was already there

40 - "model.model.": "model.",

41 - # Classifier/scoring heads will be adjacent to `model`

42 - "model.score": "classifier",

43 - "model.classifier": "classifier",

37 },

38 orig_to_new_suffix={

39 # Replace legacy suffixes used for norms

vllm/model_executor/models/transformers/pooling.py

+7 -6

@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING

21 import torch

22 from transformers import AutoModelForSequenceClassification

24 +from vllm.config.utils import getattr_iter

25 from vllm.model_executor.layers.pooler import (

26 ClassifierPooler,

27 CLSPool,

@@ -82,14 +83,14 @@ class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):

83 if hasattr(module, "pooler") and module.pooler is None:

84 self.model.pooler = None

85 break

85 - if self.model.pooler is not None:

86 - raise ValueError(

87 - "Sequence classification models with pooling layers are not "

88 - "supported yet in the Transformers backend."

89 - )

87 # Unlike `lm_head`, `classifier` is not always `nn.Linear`.

92 - self.classifier = seq_cls_model.classifier

88 + self.classifier = getattr_iter(seq_cls_model, ["classifier", "score"], None)

89 + if self.classifier is None:

90 + raise ValueError(

91 + "Could not find `classifier` or `score` layer in the "

92 + "`AutoModelForSequenceClassification` instance."

93 + )

94 self.init_parameters(self.classifier, dtype=self.model_config.head_dtype)

96 class ClassifierWithReshape(self.classifier.__class__):

vllm/model_executor/models/utils.py

@@ -46,6 +46,14 @@ class WeightsMapper:

46 orig_to_new_prefix: WeightsMapping = field(default_factory=dict)

47 orig_to_new_suffix: WeightsMapping = field(default_factory=dict)

49 + def __or__(self, other: "WeightsMapper") -> "WeightsMapper":

50 + """Combine two `WeightsMapper`s by merging their mappings."""

51 + return WeightsMapper(

52 + orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr},

53 + orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix},

54 + orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},

55 + )

56 +

57 def _map_name(self, key: str) -> str | None:

58 for substr, new_key in self.orig_to_new_substr.items():

59 if substr in key: