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

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

[Model][3/N] Automatic conversion of CrossEncoding model (#20168)

Signed-off-by: wang.yuqi <noooop@126.com>

wang.yuqi committed Jul 4, 2025 at 20:47 UTC 2e26f9156ac1eb4cc402221ffb7aa98fe10d9907

8 files changed +234 -133

docs/models/supported_models.md

+15 -6

@@ -477,12 +477,20 @@ If your model is not in the above list, we will try to automatically convert the

477

478 Specified using `--task score`.

479

481 -|---------------------------------------|-------------------|--------------------------------------------------------------------------------------|-----------------------|

481 +|---------------------------------------|-------------------|--------------------------------------------------------------------------------------|---------------------|

487 +

488 +!!! note

489 + Load the official original `mxbai-rerank-v2` by using the following command.

490 +

491 + ```bash

492 + vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}'

493 + ```

494

495 !!! note

496 Load the official original `Qwen3 Reranker` by using the following command. More information can be found at: <gh-file:examples/offline_inference/qwen3_reranker.py>.

@@ -490,6 +498,7 @@ Specified using `--task score`.

498 ```bash

499 vllm serve Qwen/Qwen3-Reranker-0.6B --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}'

500 ```

501 +

502 [](){ #supported-mm-models }

503

504 ## List of Multimodal Language Models

tests/models/language/pooling/test_embedding.py

+9 -1

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import os

4 +from typing import Optional

6 import pytest

@@ -74,6 +75,13 @@ def test_models(

75 vllm_extra_kwargs["override_pooler_config"] = \

76 PoolerConfig(pooling_type="MEAN", normalize=False)

78 + max_model_len: Optional[int] = 512

79 + if model in [

80 + "sentence-transformers/all-MiniLM-L12-v2",

81 + "sentence-transformers/stsb-roberta-base-v2"

82 + ]:

83 + max_model_len = None

84 +

85 # The example_prompts has ending "\n", for example:

86 # "Write a short story about a robot that dreams for the first time.\n"

87 # sentence_transformers will strip the input texts, see:

@@ -87,7 +95,7 @@ def test_models(

96 with vllm_runner(model,

97 task="embed",

90 - max_model_len=512,

98 + max_model_len=max_model_len,

99 **vllm_extra_kwargs) as vllm_model:

100 vllm_outputs = vllm_model.embed(example_prompts)

101

tests/models/language/pooling/test_gte.py

+12 -4

@@ -56,10 +56,16 @@ MODELS = [

56 enable_test=False),

57 ]

59 +V1FlashAttentionImpNotSupported = [

60 + "Alibaba-NLP/gte-Qwen2-1.5B-instruct", "Alibaba-NLP/gte-modernbert-base"

61 +]

62 +

64 @pytest.mark.parametrize("model_info", MODELS)

61 -def test_embed_models_mteb(hf_runner, vllm_runner,

62 - model_info: EmbedModelInfo) -> None:

65 +def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo,

66 + monkeypatch) -> None:

67 + if model_info.name in V1FlashAttentionImpNotSupported:

68 + monkeypatch.setenv("VLLM_USE_V1", "0")

70 vllm_extra_kwargs: dict[str, Any] = {}

71 if model_info.architecture == "GteNewModel":

@@ -71,8 +77,10 @@ def test_embed_models_mteb(hf_runner, vllm_runner,

78 @pytest.mark.parametrize("model_info", MODELS)

79 def test_embed_models_correctness(hf_runner, vllm_runner,

74 - model_info: EmbedModelInfo,

75 - example_prompts) -> None:

80 + model_info: EmbedModelInfo, example_prompts,

81 + monkeypatch) -> None:

82 + if model_info.name in V1FlashAttentionImpNotSupported:

83 + monkeypatch.setenv("VLLM_USE_V1", "0")

85 vllm_extra_kwargs: dict[str, Any] = {}

86 if model_info.architecture == "GteNewModel":

tests/models/language/pooling/test_mxbai_rerank.py new

+84

@@ -0,0 +1,84 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Any

4 +

5 +import pytest

6 +import torch

7 +

8 +from tests.conftest import HfRunner

9 +

10 +from .mteb_utils import RerankModelInfo, mteb_test_rerank_models

11 +

12 +RERANK_MODELS = [

13 + RerankModelInfo("mixedbread-ai/mxbai-rerank-base-v2",

14 + architecture="Qwen2ForSequenceClassification",

15 + dtype="float32",

16 + enable_test=True),

17 + RerankModelInfo("mixedbread-ai/mxbai-rerank-large-v2",

18 + architecture="Qwen2ForSequenceClassification",

19 + dtype="float32",

20 + enable_test=False)

21 +]

22 +

23 +

24 +class MxbaiRerankerHfRunner(HfRunner):

25 +

26 + def __init__(self,

27 + model_name: str,

28 + dtype: str = "auto",

29 + *args: Any,

30 + **kwargs: Any) -> None:

31 + from transformers import AutoModelForCausalLM, AutoTokenizer

32 + super().__init__(model_name, dtype, auto_cls=AutoModelForCausalLM)

33 +

34 + self.tokenizer = AutoTokenizer.from_pretrained(model_name,

35 + padding_side='left')

36 + self.yes_loc = self.tokenizer.convert_tokens_to_ids("1")

37 + self.no_loc = self.tokenizer.convert_tokens_to_ids("0")

38 +

39 + def predict(self, prompts: list[list[str]], *args,

40 + **kwargs) -> torch.Tensor:

41 +

42 + def process_inputs(pairs):

43 + inputs = self.tokenizer(pairs,

44 + padding=False,

45 + truncation='longest_first',

46 + return_attention_mask=False)

47 + for i, ele in enumerate(inputs['input_ids']):

48 + inputs['input_ids'][i] = ele

49 + inputs = self.tokenizer.pad(inputs,

50 + padding=True,

51 + return_tensors="pt")

52 + for key in inputs:

53 + inputs[key] = inputs[key].to(self.model.device)

54 + return inputs

55 +

56 + @torch.no_grad()

57 + def compute_logits(inputs):

58 + logits = self.model(**inputs).logits[:, -1, :]

59 + yes_logits = logits[:, self.yes_loc]

60 + no_logits = logits[:, self.no_loc]

61 + logits = yes_logits - no_logits

62 + scores = logits.float().sigmoid()

63 + return scores

64 +

65 + scores = []

66 + for prompt in prompts:

67 + inputs = process_inputs([prompt])

68 + score = compute_logits(inputs)

69 + scores.append(score[0].item())

70 + return torch.Tensor(scores)

71 +

72 +

73 +@pytest.mark.parametrize("model_info", RERANK_MODELS)

74 +def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

75 + vllm_extra_kwargs: dict[str, Any] = {}

76 + if model_info.architecture == "Qwen2ForSequenceClassification":

77 + vllm_extra_kwargs["hf_overrides"] = {

78 + "architectures": ["Qwen2ForSequenceClassification"],

79 + "classifier_from_token": ["0", "1"],

80 + "method": "from_2_way_softmax",

81 + }

82 +

83 + mteb_test_rerank_models(MxbaiRerankerHfRunner, vllm_runner, model_info,

84 + vllm_extra_kwargs)

vllm/config.py

+10 -3

@@ -466,6 +466,9 @@ class ModelConfig:

466 "affect the random state of the Python process that "

467 "launched vLLM.", self.seed)

468

469 + # Keep set served_model_name before maybe_model_redirect(self.model)

470 + self.served_model_name = get_served_model_name(self.model,

471 + self.served_model_name)

472 self.model = maybe_model_redirect(self.model)

473 # The tokenizer is consistent with the model by default.

474 if self.tokenizer is None:

@@ -609,8 +612,6 @@ class ModelConfig:

612

613 self.original_max_model_len = self.max_model_len

614 self.max_model_len = self.get_and_verify_max_len(self.max_model_len)

612 - self.served_model_name = get_served_model_name(self.model,

613 - self.served_model_name)

615 self.multimodal_config = self._init_multimodal_config()

616 if not self.skip_tokenizer_init:

617 self._verify_tokenizer_mode()

@@ -1420,7 +1421,7 @@ class ModelConfig:

1421

1422 @property

1423 def is_cross_encoder(self) -> bool:

1423 - return self.registry.is_cross_encoder_model(self.architectures)

1424 + return self.task == "classify"

1425

1426 @property

1427 def use_mla(self) -> bool:

@@ -4762,6 +4763,12 @@ class VllmConfig:

4763 if cls is not None:

4764 cls.verify_and_update_config(self)

4765

4766 + if self.model_config.task == "classify":

4767 + # Maybe convert ForCausalLM into ForSequenceClassification model.

4768 + from vllm.model_executor.models.adapters import (

4769 + SequenceClassificationConfig)

4770 + SequenceClassificationConfig.verify_and_update_config(self)

4771 +

4772 def __str__(self):

4773 return (

4774 f"model={self.model_config.model!r},"

vllm/model_executor/models/adapters.py

+99 -3

@@ -2,14 +2,17 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from collections.abc import Iterable

5 -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union

5 +from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union, cast

7 import torch

8 import torch.nn as nn

10 +from vllm.model_executor.models.config import VerifyAndUpdateConfig

11 +

12 from .interfaces_base import VllmModelForPooling, is_pooling_model

14 if TYPE_CHECKING:

15 + from vllm.config import VllmConfig

16 from vllm.model_executor.layers.pooler import PoolingType

18 _T = TypeVar("_T", bound=type[nn.Module])

@@ -39,7 +42,6 @@ def _create_pooling_model_cls(

42 default_softmax: bool,

43 ) -> _T:

44 # Lazy import

42 - from vllm.config import VllmConfig

45 from vllm.model_executor.layers.pooler import Pooler, PoolerOutput

46 from vllm.model_executor.pooling_metadata import PoolingMetadata

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

164 return cls

165

166 # Lazy import

165 - from vllm.config import VllmConfig

167 from vllm.model_executor.layers.linear import RowParallelLinear

168 from vllm.model_executor.layers.pooler import PoolerOutput, PoolingType

169 from vllm.model_executor.models.interfaces import SupportsCrossEncoding

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

194 config = vllm_config.model_config.hf_config

195 quant_config = vllm_config.quant_config

196

197 + self.vllm_config = vllm_config

198 self.task = vllm_config.model_config.task

199 self.pooling_type = (

200 vllm_config.model_config.pooler_config.pooling_type)

@@ -242,6 +244,17 @@ def as_seq_cls_model(cls: _T) -> _T:

244 ]

245 return PoolerOutput(outputs=pooled_outputs)

246

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

248 + tokens = getattr(self.config, "classifier_from_token", None)

249 + method = getattr(self.config, "method", None)

250 +

251 + if tokens is None and method is None:

252 + return super().load_weights(weights)

253 + else:

254 + # Online convert ForCausalLM into

255 + # ForSequenceClassification model.

256 + return seq_cls_model_loader(self, weights)

257 +

258

259 ModelForSequenceClassification.__name__ = \

260 _get_pooling_model_name(cls.__name__, "ForSequenceClassification")

@@ -277,3 +290,86 @@ def as_reward_model(cls: _T) -> _T:

290 _get_pooling_model_name(cls.__name__, "ForReward")

291

292 return ModelForReward # type: ignore

293 +

294 +

295 +class SequenceClassificationConfig(VerifyAndUpdateConfig):

296 +

297 + @staticmethod

298 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:

299 + config = vllm_config.model_config.hf_config

300 + method = getattr(config, "method", None)

301 + tokens = getattr(config, "classifier_from_token", None)

302 +

303 + if method is None:

304 + return

305 +

306 + assert tokens is not None

307 + assert method in SEQ_CLS_LOAD_METHODS, f"method {method} not supported"

308 +

309 + if method == "from_2_way_softmax":

310 + assert len(tokens) == 2

311 + config.num_labels = 1

312 + else:

313 + config.num_labels = len(tokens)

314 +

315 +

316 +def load_weights_using_from_2_way_softmax(

317 + model, weights: Iterable[tuple[str, torch.Tensor]]):

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

319 + from vllm.model_executor.layers.vocab_parallel_embedding import (

320 + ParallelLMHead)

321 + from vllm.model_executor.models.utils import AutoWeightsLoader

322 +

323 + model_config = model.vllm_config.model_config

324 + tokens = getattr(model.config, "classifier_from_token", [])

325 + tokens = cast(list[int], tokens)

326 + assert len(tokens) == 2

327 +

328 + device = model.score.weight.device

329 +

330 + if model.config.tie_word_embeddings:

331 + model.lm_head = model.model.embed_tokens

332 + else:

333 + model.lm_head = ParallelLMHead(model.config.vocab_size,

334 + model.config.hidden_size,

335 + quant_config=model.quant_config)

336 +

337 + loader = AutoWeightsLoader(model)

338 + loaded_weights = loader.load_weights(weights)

339 +

340 + from vllm.transformers_utils.tokenizer import get_tokenizer

341 + tokenizer = get_tokenizer(model_config.tokenizer,

342 + revision=model_config.tokenizer_revision,

343 + tokenizer_mode=model_config.tokenizer_mode,

344 + trust_remote_code=model_config.trust_remote_code)

345 +

346 + false_id = tokenizer.convert_tokens_to_ids(tokens[0])

347 + true_id = tokenizer.convert_tokens_to_ids(tokens[1])

348 + weight = model.lm_head.weight.data[true_id].to(device).to(

349 + torch.float32) - model.lm_head.weight.data[false_id].to(device).to(

350 + torch.float32)

351 + model.score.weight.data.copy_(weight)

352 +

353 + del model.lm_head

354 + loaded_weights.add("score.weight")

355 + loaded_weights.discard("lm_head.weight")

356 + return loaded_weights

357 +

358 +

359 +SEQ_CLS_LOAD_METHODS = {

360 + "from_2_way_softmax": load_weights_using_from_2_way_softmax,

361 +}

362 +

363 +

364 +def seq_cls_model_loader(model, weights: Iterable[tuple[str, torch.Tensor]]):

365 + # Online convert ForCausalLM into ForSequenceClassification model.

366 + # - from_2_way_softmax:

367 + # - Qwen3ForCausalLM

368 + # - Qwen3-Reranker

369 + # - Qwen2ForCausalLM

370 + # - mxbai-rerank-v2

371 +

372 + config = model.vllm_config.model_config.hf_config

373 + method = getattr(config, "method", None)

374 + assert method in SEQ_CLS_LOAD_METHODS, f"method {method} not supported"

375 + return SEQ_CLS_LOAD_METHODS[method](model, weights)

vllm/model_executor/models/config.py

+1 -1

@@ -167,7 +167,7 @@ class Qwen3ForSequenceClassificationConfig(VerifyAndUpdateConfig):

167 assert tokens is not None and len(tokens) == 2, \

168 ("Try loading the original Qwen3 Reranker?, see: "

169 "https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen3_reranker.py")

170 - config.num_labels = 1

170 + vllm_config.model_config.hf_config.method = "from_2_way_softmax"

171

172

173 class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):

vllm/model_executor/models/qwen3.py

+4 -115

@@ -38,15 +38,14 @@ from vllm.model_executor.layers.layernorm import RMSNorm

38 from vllm.model_executor.layers.linear import (QKVParallelLinear,

39 RowParallelLinear)

40 from vllm.model_executor.layers.logits_processor import LogitsProcessor

41 -from vllm.model_executor.layers.pooler import Pooler, PoolingType

41 from vllm.model_executor.layers.quantization import QuantizationConfig

42 from vllm.model_executor.layers.rotary_embedding import get_rope

43 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

45 -from vllm.model_executor.pooling_metadata import PoolingMetadata

44 from vllm.model_executor.sampling_metadata import SamplingMetadata

47 -from vllm.sequence import IntermediateTensors, PoolerOutput

45 +from vllm.sequence import IntermediateTensors

49 -from .interfaces import SupportsCrossEncoding, SupportsLoRA, SupportsPP

47 +from .adapters import as_seq_cls_model

48 +from .interfaces import SupportsLoRA, SupportsPP

49 from .qwen2 import Qwen2MLP as Qwen3MLP

50 from .qwen2 import Qwen2Model

51 from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix

@@ -323,114 +322,4 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

322 return loader.load_weights(weights)

323

324

326 -class Qwen3ForSequenceClassification(nn.Module, SupportsLoRA,

327 - SupportsCrossEncoding):

328 -

329 - def __init__(

330 - self,

331 - vllm_config: "VllmConfig",

332 - prefix: str = "",

333 - ) -> None:

334 - super().__init__()

335 -

336 - config = vllm_config.model_config.hf_config

337 - quant_config = vllm_config.quant_config

338 - pooler_config = vllm_config.model_config.pooler_config

339 -

340 - self.vllm_config = vllm_config

341 - self.config = config

342 - self.quant_config = quant_config

343 - self.prefix = prefix

344 - self.model = Qwen3Model(vllm_config=vllm_config,

345 - prefix=maybe_prefix(prefix, "model"))

346 - self.score = RowParallelLinear(config.hidden_size,

347 - config.num_labels,

348 - quant_config=quant_config,

349 - input_is_parallel=False,

350 - bias=False,

351 - prefix=maybe_prefix(prefix, "score"))

352 -

353 - self._pooler = Pooler.from_config_with_defaults(

354 - pooler_config,

355 - pooling_type=PoolingType.LAST,

356 - normalize=False,

357 - softmax=True)

358 -

359 - def forward(

360 - self,

361 - input_ids: torch.Tensor,

362 - positions: torch.Tensor,

363 - intermediate_tensors: Optional[IntermediateTensors] = None,

364 - inputs_embeds: Optional[torch.Tensor] = None,

365 - ) -> torch.Tensor:

366 - return self.model(input_ids=input_ids,

367 - positions=positions,

368 - inputs_embeds=inputs_embeds,

369 - intermediate_tensors=intermediate_tensors)

370 -

371 - def pooler(

372 - self,

373 - hidden_states: torch.Tensor,

374 - pooling_metadata: PoolingMetadata,

375 - ) -> Optional[PoolerOutput]:

376 - hidden_states = self._pooler.extract_states(hidden_states,

377 - pooling_metadata)

378 -

379 - if isinstance(hidden_states, list):

380 - logits = [self.score(state)[0] for state in hidden_states]

381 - else:

382 - logits, _ = self.score(hidden_states)

383 -

384 - pooled_data = self._pooler.head(logits, pooling_metadata)

385 - pooled_outputs = [

386 - self._pooler.build_output(data.squeeze(-1)) for data in pooled_data

387 - ]

388 - return PoolerOutput(outputs=pooled_outputs)

389 -

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

391 - is_original_qwen3_reranker = getattr(self.config,

392 - "is_original_qwen3_reranker",

393 - False)

394 -

395 - if not is_original_qwen3_reranker:

396 - loader = AutoWeightsLoader(self)

397 - return loader.load_weights(weights)

398 -

399 - return self.load_weights_from_original_qwen3_reranker(weights)

400 -

401 - def load_weights_from_original_qwen3_reranker(

402 - self, weights: Iterable[tuple[str, torch.Tensor]]):

403 -

404 - model_config = self.vllm_config.model_config

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

406 - device = self.score.weight.device

407 -

408 - if self.config.tie_word_embeddings:

409 - self.lm_head = self.model.embed_tokens

410 - else:

411 - self.lm_head = ParallelLMHead(self.config.vocab_size,

412 - self.config.hidden_size,

413 - quant_config=self.quant_config,

414 - prefix=maybe_prefix(

415 - self.prefix, "lm_head"))

416 -

417 - loader = AutoWeightsLoader(self)

418 - loaded_weights = loader.load_weights(weights)

419 -

420 - from vllm.transformers_utils.tokenizer import get_tokenizer

421 - tokenizer = get_tokenizer(

422 - model_config.tokenizer,

423 - revision=model_config.tokenizer_revision,

424 - tokenizer_mode=model_config.tokenizer_mode,

425 - trust_remote_code=model_config.trust_remote_code)

426 -

427 - a = tokenizer.convert_tokens_to_ids(tokens[0])

428 - b = tokenizer.convert_tokens_to_ids(tokens[1])

429 - weight = self.lm_head.weight.data[b].to(

430 - device) - self.lm_head.weight.data[a].to(device)

431 - self.score.weight.data.copy_(weight)

432 -

433 - del self.lm_head

434 - loaded_weights.add("score.weight")

435 - loaded_weights.discard("lm_head.weight")

436 - return loaded_weights

325 +Qwen3ForSequenceClassification = as_seq_cls_model(Qwen3ForCausalLM)