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

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

[Model][0/N] Improve all pooling task | clean up (#25817)

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

wang.yuqi committed Oct 13, 2025 at 16:44 UTC 767c3ab8693aba5c5c629a1c9c6550d5cfab0cc4

19 files changed +197 -188

docs/models/supported_models.md

+1 -1

@@ -581,7 +581,7 @@ These models primarily support the [`LLM.encode`](./pooling_models.md#llmencode)

582

583 !!! note

584 - Named Entity Recognition (NER) usage, please refer to <gh-file:examples/offline_inference/pooling/ner.py>, <gh-file:examples/online_serving/pooling/ner.py>.

584 + Named Entity Recognition (NER) usage, please refer to <gh-file:examples/offline_inference/pooling/ner.py>, <gh-file:examples/online_serving/pooling/ner_client.py>.

585

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

587

examples/online_serving/pooling/README.md

+1 -1

@@ -15,7 +15,7 @@ python examples/online_serving/pooling/jinaai_rerank_client.py

15 ## Named Entity Recognition (NER) usage

17 ```bash

18 -python examples/online_serving/pooling/ner.py

18 +python examples/online_serving/pooling/ner_client.py

19 ```

21 ## Openai chat embedding for multimodal usage

examples/online_serving/pooling/ner_client.py renamed

tests/ci_envs.py

@@ -8,6 +8,8 @@ import os

8 from collections.abc import Callable

9 from typing import TYPE_CHECKING, Any

11 +from vllm.envs import maybe_convert_bool

12 +

13 if TYPE_CHECKING:

14 VLLM_CI_NO_SKIP: bool = False

15 VLLM_CI_DTYPE: str | None = None

@@ -25,6 +27,10 @@ environment_variables: dict[str, Callable[[], Any]] = {

27 "VLLM_CI_HEAD_DTYPE": lambda: os.getenv("VLLM_CI_HEAD_DTYPE", None),

28 # Allow changing the head dtype used by transformers in tests

29 "VLLM_CI_HF_DTYPE": lambda: os.getenv("VLLM_CI_HF_DTYPE", None),

30 + # Allow control over whether tests use enforce_eager

31 + "VLLM_CI_ENFORCE_EAGER": lambda: maybe_convert_bool(

32 + os.getenv("VLLM_CI_ENFORCE_EAGER", None)

33 + ),

34 }

tests/entrypoints/pooling/llm/test_classify.py

@@ -58,7 +58,9 @@ def test_pooling_params(llm: LLM):

58 )

61 +@pytest.mark.skip_global_cleanup

62 def test_encode_api(llm: LLM):

63 + # chunked prefill does not support all pooling

64 err_msg = "pooling_task must be one of.+"

65 with pytest.raises(ValueError, match=err_msg):

66 llm.encode(prompts, use_tqdm=False)

tests/entrypoints/pooling/llm/test_embedding.py

-1

@@ -35,7 +35,6 @@ def llm():

35 cleanup_dist_env_and_memory()

38 -@pytest.mark.skip_global_cleanup

38 def test_pooling_params(llm: LLM):

39 def get_outputs(normalize):

40 outputs = llm.embed(

tests/entrypoints/pooling/llm/test_encode.py

-1

@@ -74,7 +74,6 @@ def test_multiple_pooling_params(llm: LLM):

74 assert len(PROMPTS) == len(outputs)

77 -@pytest.mark.skip_global_cleanup

77 def test_right_side_truncation(llm: LLM):

78 # Embeddings models should truncate the end of the prompt

79 tokenizer = llm.get_tokenizer()

tests/entrypoints/pooling/llm/test_score.py

-1

@@ -33,7 +33,6 @@ def llm():

33 cleanup_dist_env_and_memory()

36 -@pytest.mark.skip_global_cleanup

36 def test_pooling_params(llm: LLM):

37 def get_outputs(activation):

38 text_1 = "What is the capital of France?"

tests/models/language/generation_ppl_test/ppl_utils.py

+6 -20

@@ -3,12 +3,15 @@

3 # Adapted from https://huggingface.co/docs/transformers/perplexity

4 from typing import cast

6 -import pytest

6 import torch

7 from datasets import load_dataset

9 import tests.ci_envs as ci_envs

11 -from tests.models.utils import GenerateModelInfo, TokensTextLogprobsPromptLogprobs

10 +from tests.models.utils import (

11 + GenerateModelInfo,

12 + TokensTextLogprobsPromptLogprobs,

13 + get_vllm_extra_kwargs,

14 +)

15 from vllm.logprobs import Logprob

17 # See #24485

@@ -25,27 +28,10 @@ def wikitext_ppl_test(

28 vllm_extra_kwargs=None,

29 atol=PPL_TOL,

30 ):

28 - # A model family has many models with the same architecture,

29 - # and we don't need to test each one.

30 - if not ci_envs.VLLM_CI_NO_SKIP and not model_info.enable_test:

31 - pytest.skip("Skipping test.")

31 + vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)

33 dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")

35 - # Allow vllm to test using the given dtype, such as float32

36 - vllm_extra_kwargs = vllm_extra_kwargs or {}

37 - vllm_extra_kwargs["dtype"] = ci_envs.VLLM_CI_DTYPE or model_info.dtype

38 -

39 - # Allow vllm to test using hf_overrides

40 - if model_info.hf_overrides is not None:

41 - vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides

42 -

43 - # Allow changing the head dtype used by vllm in tests

44 - if ci_envs.VLLM_CI_HEAD_DTYPE is not None:

45 - if "hf_overrides" not in vllm_extra_kwargs:

46 - vllm_extra_kwargs["hf_overrides"] = {}

47 - vllm_extra_kwargs["hf_overrides"]["head_dtype"] = ci_envs.VLLM_CI_HEAD_DTYPE

48 -

35 with vllm_runner(

36 model_info.name,

37 gpu_memory_utilization=0.7,

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

+47

@@ -0,0 +1,47 @@

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

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

3 +import pytest

4 +import torch

5 +from transformers import AutoModelForSequenceClassification

6 +

7 +

8 +@pytest.mark.parametrize(

9 + "model",

10 + ["nie3e/sentiment-polish-gpt2-small"],

11 +)

12 +@pytest.mark.parametrize("dtype", ["half"])

13 +def test_classify_models(

14 + hf_runner,

15 + vllm_runner,

16 + example_prompts,

17 + model: str,

18 + dtype: str,

19 +) -> None:

20 + with hf_runner(

21 + model, dtype=dtype, auto_cls=AutoModelForSequenceClassification

22 + ) as hf_model:

23 + hf_outputs = hf_model.classify(example_prompts)

24 +

25 + for head_dtype_str in ["float32", "model"]:

26 + with vllm_runner(

27 + model,

28 + max_model_len=512,

29 + dtype=dtype,

30 + hf_overrides={"head_dtype": head_dtype_str},

31 + ) as vllm_model:

32 + model_config = vllm_model.llm.llm_engine.model_config

33 + model_dtype = model_config.dtype

34 + head_dtype = model_config.head_dtype

35 +

36 + if head_dtype_str == "float32":

37 + assert head_dtype == torch.float32

38 + elif head_dtype_str == "model":

39 + assert head_dtype == model_dtype

40 +

41 + vllm_outputs = vllm_model.classify(example_prompts)

42 +

43 + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):

44 + hf_output = torch.tensor(hf_output).float()

45 + vllm_output = torch.tensor(vllm_output).float()

46 +

47 + assert torch.allclose(hf_output, vllm_output, atol=1e-2)

tests/models/language/pooling/test_splade_sparse_pooler.py

+9 -42

@@ -3,7 +3,6 @@

4 import types

6 -import numpy as np

6 import pytest

7 import torch

8 import torch.nn as nn

@@ -14,11 +13,12 @@ from vllm.model_executor.models.bert import (

13 )

15 # ---------------------------------------------------------------------

17 -# 1) Functional test: SPLADE formula correctness (no HF download needed)

16 +# Functional test: SPLADE formula correctness (no HF download needed)

17 # ---------------------------------------------------------------------

20 @pytest.mark.parametrize("B,T,H,V", [(2, 3, 5, 7)])

21 +@torch.inference_mode

22 def test_splade_pooler_matches_reference_formula(B, T, H, V):

23 """Ensure SPLADESparsePooler forward() matches the mathematical formula:

24 log1p(relu(logits)) -> max over sequence length (after masking)."""

@@ -26,9 +26,11 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):

27 # Prepare [B] sequences of shape [T, H]

28 hs_list = [torch.randn(T, H) for _ in range(B)]

29 + hs_tenser = torch.cat(hs_list)

31 # Simulate PoolingMetadata (only required fields)

32 prompt_lens = [T, T - 1]

33 + prompt_lens_tenser = torch.tensor(prompt_lens, dtype=torch.int32)

34 token_ids = torch.tensor(

35 [

36 [101, 5, 102], # Batch 0: [CLS], token, [SEP]

@@ -36,7 +38,9 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):

38 ],

39 dtype=torch.long,

40 )

39 - meta = types.SimpleNamespace(prompt_lens=prompt_lens, prompt_token_ids=token_ids)

41 + meta = types.SimpleNamespace(

42 + prompt_lens=prompt_lens_tenser, prompt_token_ids=token_ids

43 + )

45 # MLM head (prefer BertMLMHead, fallback to Linear if unavailable)

46 try:

@@ -46,10 +50,10 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):

51 # Forward pass through SPLADE pooler

52 pooler = SPLADESparsePooler(mlm_head=mlm_head, pooling="max", remove_cls_sep=True)

49 - pooled = pooler(hidden_states=hs_list, pooling_metadata=meta) # list of [V]

53 + pooled = pooler(hidden_states=hs_tenser, pooling_metadata=meta) # list of [V]

55 # Basic output checks

52 - assert isinstance(pooled, list) and len(pooled) == B

56 + assert isinstance(pooled, torch.Tensor) and len(pooled) == B

57 for vec in pooled:

58 assert vec.shape == (V,)

59 assert torch.isfinite(vec).all()

@@ -83,40 +87,3 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):

87 rtol=1e-4,

88 atol=1e-4,

89 )

86 -

87 -

88 -# ---------------------------------------------------------------------

89 -# 2) Integration smoke test: end-to-end embedding path wiring

90 -# ---------------------------------------------------------------------

91 -

92 -

93 -@pytest.mark.cpu_model

94 -def test_bert_splade_sparse_embed_smoke(vllm_runner, monkeypatch):

95 - """Ensure BertSpladeSparseEmbeddingModel loads and produces sparse embeddings."""

96 - from transformers import AutoTokenizer

97 -

98 - MODEL_ID = "hf-internal-testing/tiny-random-bert"

99 - hf_overrides = {"architectures": ["BertSpladeSparseEmbeddingModel"]}

100 -

101 - # Enforce CPU-only execution (optional)

102 - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "")

103 - monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")

104 -

105 - tok = AutoTokenizer.from_pretrained(MODEL_ID)

106 - vocab_size = tok.vocab_size

107 -

108 - # The embed path should route through SPLADESparsePooler

109 - with vllm_runner(

110 - MODEL_ID,

111 - runner="pooling",

112 - max_model_len=64,

113 - hf_overrides=hf_overrides,

114 - ) as vm:

115 - outs = vm.embed(["hello world", "splade sparse test"])

116 -

117 - # Basic sanity checks

118 - assert len(outs) == 2

119 - assert outs[0].shape[0] == vocab_size

120 - assert outs[1].shape[0] == vocab_size

121 - assert np.isfinite(outs[0]).all() and (outs[0] >= 0).all()

122 - assert np.isfinite(outs[1]).all() and (outs[1] >= 0).all()

tests/models/language/pooling_mteb_test/mteb_utils.py

+14 -41

@@ -6,12 +6,16 @@ from collections.abc import Sequence

7 import mteb

8 import numpy as np

9 -import pytest

9 import requests

10 import torch

12 import tests.ci_envs as ci_envs

14 -from tests.models.utils import EmbedModelInfo, RerankModelInfo, check_embeddings_close

13 +from tests.models.utils import (

14 + EmbedModelInfo,

15 + RerankModelInfo,

16 + check_embeddings_close,

17 + get_vllm_extra_kwargs,

18 +)

20 # Most embedding models on the STS12 task (See #17175):

21 # - Model implementation and minor changes in tensor dtype

@@ -165,28 +169,11 @@ def mteb_test_embed_models(

169 hf_model_callback=None,

170 atol=MTEB_EMBED_TOL,

171 ):

168 - # A model family has many models with the same architecture,

169 - # and we don't need to test each one.

170 - if not ci_envs.VLLM_CI_NO_SKIP and not model_info.enable_test:

171 - pytest.skip("Skipping test.")

172 + vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)

173

174 # Test embed_dims, isnan and whether to use normalize

175 example_prompts = ["The chef prepared a delicious meal." * 1000]

176

176 - # Allow vllm to test using the given dtype, such as float32

177 - vllm_extra_kwargs = vllm_extra_kwargs or {}

178 - vllm_extra_kwargs["dtype"] = ci_envs.VLLM_CI_DTYPE or model_info.dtype

179 -

180 - # Allow vllm to test using hf_overrides

181 - if model_info.hf_overrides is not None:

182 - vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides

183 -

184 - # Allow changing the head dtype used by vllm in tests

185 - if ci_envs.VLLM_CI_HEAD_DTYPE is not None:

186 - if "hf_overrides" not in vllm_extra_kwargs:

187 - vllm_extra_kwargs["hf_overrides"] = {}

188 - vllm_extra_kwargs["hf_overrides"]["head_dtype"] = ci_envs.VLLM_CI_HEAD_DTYPE

189 -

177 with vllm_runner(

178 model_info.name,

179 runner="pooling",

@@ -212,9 +199,12 @@ def mteb_test_embed_models(

199 vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype

200 head_dtype = model_config.head_dtype

201

215 - # Test embed_dims, isnan and whether to use normalize

202 + # Test embedding_size, isnan and whether to use normalize

203 vllm_outputs = vllm_model.embed(example_prompts, truncate_prompt_tokens=-1)

217 - assert not torch.any(torch.isnan(torch.tensor(vllm_outputs)))

204 + outputs_tensor = torch.tensor(vllm_outputs)

205 + assert not torch.any(torch.isnan(outputs_tensor))

206 + embedding_size = model_config.embedding_size

207 + assert torch.tensor(vllm_outputs).shape[-1] == embedding_size

208

209 # Accelerate mteb test by setting

210 # SentenceTransformers mteb score to a constant

@@ -231,7 +221,7 @@ def mteb_test_embed_models(

221 st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)

222 st_dtype = next(hf_model.model.parameters()).dtype

223

234 - # Test embed_dims and whether to use normalize

224 + # Check embeddings close to hf outputs

225 hf_outputs = hf_model.encode(example_prompts)

226 check_embeddings_close(

227 embeddings_0_lst=hf_outputs,

@@ -323,24 +313,7 @@ def mteb_test_rerank_models(

313 vllm_mteb_encoder=VllmMtebEncoder,

314 atol=MTEB_RERANK_TOL,

315 ):

326 - # A model family has many models with the same architecture,

327 - # and we don't need to test each one.

328 - if not ci_envs.VLLM_CI_NO_SKIP and not model_info.enable_test:

329 - pytest.skip("Skipping test.")

330 -

331 - # Allow vllm to test using the given dtype, such as float32

332 - vllm_extra_kwargs = vllm_extra_kwargs or {}

333 - vllm_extra_kwargs["dtype"] = ci_envs.VLLM_CI_DTYPE or model_info.dtype

334 -

335 - # Allow vllm to test using hf_overrides

336 - if model_info.hf_overrides is not None:

337 - vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides

338 -

339 - # Allow changing the head dtype used by vllm in tests

340 - if ci_envs.VLLM_CI_HEAD_DTYPE is not None:

341 - if "hf_overrides" not in vllm_extra_kwargs:

342 - vllm_extra_kwargs["hf_overrides"] = {}

343 - vllm_extra_kwargs["hf_overrides"]["head_dtype"] = ci_envs.VLLM_CI_HEAD_DTYPE

316 + vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)

317

318 with vllm_runner(

319 model_info.name,

tests/models/utils.py

+30

@@ -15,6 +15,7 @@ from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs

15 from vllm.multimodal.processing import InputProcessingContext

16 from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

18 +from .. import ci_envs

19 from .registry import HF_EXAMPLE_MODELS

21 TokensText = tuple[list[int], str]

@@ -414,6 +415,35 @@ class GenerateModelInfo(ModelInfo):

415 hf_ppl: float | None = None

416

417

418 +def get_vllm_extra_kwargs(model_info: ModelInfo, vllm_extra_kwargs):

419 + # A model family has many models with the same architecture,

420 + # and we don't need to test each one.

421 + if not ci_envs.VLLM_CI_NO_SKIP and not model_info.enable_test:

422 + import pytest

423 +

424 + pytest.skip("Skipping test.")

425 +

426 + # Allow vllm to test using the given dtype, such as float32

427 + vllm_extra_kwargs = vllm_extra_kwargs or {}

428 + vllm_extra_kwargs["dtype"] = ci_envs.VLLM_CI_DTYPE or model_info.dtype

429 +

430 + # Allow vllm to test using hf_overrides

431 + if model_info.hf_overrides is not None:

432 + vllm_extra_kwargs["hf_overrides"] = model_info.hf_overrides

433 +

434 + # Allow changing the head dtype used by vllm in tests

435 + if ci_envs.VLLM_CI_HEAD_DTYPE is not None:

436 + if "hf_overrides" not in vllm_extra_kwargs:

437 + vllm_extra_kwargs["hf_overrides"] = {}

438 + vllm_extra_kwargs["hf_overrides"]["head_dtype"] = ci_envs.VLLM_CI_HEAD_DTYPE

439 +

440 + # Allow control over whether tests use enforce_eager

441 + if ci_envs.VLLM_CI_ENFORCE_EAGER is not None:

442 + vllm_extra_kwargs["enforce_eager"] = ci_envs.VLLM_CI_ENFORCE_EAGER

443 +

444 + return vllm_extra_kwargs

445 +

446 +

447 def dummy_hf_overrides(

448 hf_config: PretrainedConfig,

449 *,

vllm/config/model.py

+15

@@ -30,6 +30,7 @@ from vllm.transformers_utils.config import (

30 get_sentence_transformer_tokenizer_config,

31 is_encoder_decoder,

32 is_interleaved,

33 + try_get_dense_modules,

34 try_get_generation_config,

35 try_get_safetensors_metadata,

36 try_get_tokenizer_config,

@@ -1681,6 +1682,20 @@ class ModelConfig:

1682 logger.debug_once("head dtype: %s", head_dtype)

1683 return head_dtype

1684

1685 + @property

1686 + def hidden_size(self):

1687 + if hasattr(self.hf_config, "hidden_size"):

1688 + return self.hf_config.hidden_size

1689 + text_config = self.hf_config.get_text_config()

1690 + return text_config.hidden_size

1691 +

1692 + @property

1693 + def embedding_size(self):

1694 + dense_modules = try_get_dense_modules(self.model, revision=self.revision)

1695 + if dense_modules is not None:

1696 + return dense_modules[-1]["out_features"]

1697 + return self.hidden_size

1698 +

1699 def get_and_verify_max_len(self, max_model_len: int):

1700 # Consider max_model_len in tokenizer_config only when

1701 # pooling models use absolute position_embedding.

vllm/model_executor/models/adapters.py

+17 -32

@@ -13,7 +13,10 @@ from vllm.config import VllmConfig

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.activation import get_act_fn

15 from vllm.model_executor.models.config import VerifyAndUpdateConfig

16 -from vllm.transformers_utils.config import get_hf_file_bytes, get_hf_file_to_dict

16 +from vllm.transformers_utils.config import (

17 + get_hf_file_bytes,

18 + try_get_dense_modules,

19 +)

21 from .interfaces_base import VllmModelForPooling, is_pooling_model

@@ -35,43 +38,25 @@ _GENERATE_SUFFIXES = [

38 def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None:

39 """Load Sentence-Transformers Dense projection layers."""

38 - try:

39 - modules = get_hf_file_to_dict(

40 - "modules.json", model_config.model, model_config.revision

41 - )

42 - if not modules:

43 - return None

44 -

45 - if isinstance(modules, dict):

46 - modules = modules.get("modules", [])

41 + dense_modules = try_get_dense_modules(

42 + model_config.model, revision=model_config.revision

43 + )

48 - dense_modules = [

49 - m for m in modules if m.get("type") == "sentence_transformers.models.Dense"

50 - ]

51 - if not dense_modules:

52 - return None

45 + if dense_modules is None:

46 + return

48 + try:

49 layers = []

55 - for module in dense_modules:

56 - folder = module.get("path", "")

57 -

58 - config_path = f"{folder}/config.json" if folder else "config.json"

59 - layer_config = get_hf_file_to_dict(

60 - config_path, model_config.model, model_config.revision

61 - )

62 - if not layer_config:

63 - continue

64 -

50 + for layer_config in dense_modules:

51 + folder = layer_config["folder"]

52 linear = nn.Linear(

66 - layer_config.get("in_features", 768),

67 - layer_config.get("out_features", 768),

53 + layer_config["in_features"],

54 + layer_config["out_features"],

55 bias=layer_config.get("bias", True),

56 dtype=model_config.head_dtype,

57 )

71 -

58 if not _load_dense_weights(linear, folder, model_config):

59 continue

74 -

60 layers.append(linear)

61 if act_name := layer_config.get("activation_function"):

62 layers.append(get_act_fn(act_name))

@@ -303,18 +288,18 @@ def as_seq_cls_model(cls: _T) -> _T:

288 from vllm.model_executor.models.interfaces import SupportsCrossEncoding

289 from vllm.sequence import IntermediateTensors

290

306 - from .utils import get_model_hidden_size, maybe_prefix

291 + from .utils import maybe_prefix

292

293 class ModelForSequenceClassification(

294 _create_pooling_model_cls(cls), SupportsCrossEncoding

295 ):

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

297 config = vllm_config.model_config.hf_config

298 + model_config = vllm_config.model_config

299 quant_config = vllm_config.quant_config

314 - hidden_size = get_model_hidden_size(config)

300

301 self.score = ReplicatedLinear(

317 - hidden_size,

302 + model_config.hidden_size,

303 config.num_labels,

304 bias=False,

305 params_dtype=torch.float32,

vllm/model_executor/models/gpt2.py

+5 -2

@@ -50,7 +50,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader

50 from vllm.sequence import IntermediateTensors

52 from ..layers.pooler import DispatchPooler, Pooler

53 -from .interfaces import SupportsPP

53 +from .interfaces import SupportsCrossEncoding, SupportsPP

54 from .utils import (

55 AutoWeightsLoader,

56 is_pp_missing_parameter,

@@ -321,7 +321,7 @@ class GPT2LMHeadModel(nn.Module, SupportsPP):

321 return loader.load_weights(weights)

322

323

324 -class GPT2ForSequenceClassification(nn.Module):

324 +class GPT2ForSequenceClassification(nn.Module, SupportsCrossEncoding):

325 """GPT2 Model for sequence classification.

326

327 This class expands GPT2Model with pooling and score functions - last token

@@ -358,6 +358,9 @@ class GPT2ForSequenceClassification(nn.Module):

358 }

359 )

360

361 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

362 + return self.transformer.get_input_embeddings(input_ids)

363 +

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

365 loader = AutoWeightsLoader(self)

366 return loader.load_weights(weights)

vllm/model_executor/models/gritlm.py

+10 -39

@@ -148,37 +148,6 @@ class GritLMMeanPool(nn.Module):

148 def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:

149 return PoolingParamsUpdate(requires_token_ids=True)

150

151 - def forward_one(

152 - self,

153 - hidden_states: torch.Tensor,

154 - prompt_len: torch.Tensor | None = None,

155 - instr_len: torch.Tensor | None = None,

156 - ) -> torch.Tensor:

157 - assert prompt_len is None or prompt_len == hidden_states.shape[0], (

158 - "partial prefill not supported with MEAN pooling"

159 - )

160 -

161 - return hidden_states[instr_len:].mean(dim=0, dtype=torch.float32)

162 -

163 - def forward_all(

164 - self,

165 - hidden_states: torch.Tensor,

166 - prompt_lens: torch.Tensor,

167 - instr_lens: torch.Tensor,

168 - ) -> list[torch.Tensor] | torch.Tensor:

169 - offset = 0

170 - pooled_data = list[torch.Tensor]()

171 -

172 - for prompt_len, instr_len in zip(prompt_lens, instr_lens):

173 - pooled_data.append(

174 - hidden_states[offset + instr_len : offset + prompt_len].mean(

175 - dim=0, dtype=torch.float32

176 - )

177 - )

178 - offset += prompt_len

179 -

180 - return pooled_data

181 -

151 def forward(

152 self,

153 hidden_states: torch.Tensor | list[torch.Tensor],

@@ -190,18 +159,20 @@ class GritLMMeanPool(nn.Module):

159 self._get_instruction_len(token_ids.cpu().numpy())

160 for token_ids in get_prompt_token_ids(pooling_metadata)

161 ],

193 - device=prompt_lens.device,

162 + device="cpu",

163 )

164

196 - if isinstance(hidden_states, list):

197 - return [

198 - self.forward_one(h, prompt_len, instr_len)

199 - for h, prompt_len, instr_len in zip(

200 - hidden_states, prompt_lens, instr_lens

165 + offset = 0

166 + pooled_data = list[torch.Tensor]()

167 + for prompt_len, instr_len in zip(prompt_lens, instr_lens):

168 + pooled_data.append(

169 + hidden_states[offset + instr_len : offset + prompt_len].mean(

170 + dim=0, dtype=torch.float32

171 )

202 - ]

172 + )

173 + offset += prompt_len

174

204 - return self.forward_all(hidden_states, prompt_lens, instr_lens)

175 + return pooled_data

176

177

178 class GritLMPooler(Pooler):

vllm/model_executor/models/utils.py

-7

@@ -777,13 +777,6 @@ def fast_topk(

777 return torch.topk(values, topk, dim=dim)

778

779

780 -def get_model_hidden_size(hf_config: PretrainedConfig) -> int:

781 - if hasattr(hf_config, "hidden_size"):

782 - return hf_config.hidden_size

783 - text_config = hf_config.get_text_config()

784 - return text_config.hidden_size

785 -

786 -

780 # Chunk x along the num_tokens axis for sequence parallelism

781 # NOTE: This is wrapped in a torch custom op to work around the following issue:

782 # The output tensor can have a sequence length 0 at small input sequence lengths

vllm/transformers_utils/config.py

+34

@@ -1049,6 +1049,40 @@ def try_get_tokenizer_config(

1049 return None

1050

1051

1052 +@cache

1053 +def try_get_dense_modules(

1054 + model: str | Path,

1055 + revision: str | None = None,

1056 +) -> list[dict[str, Any]] | None:

1057 + try:

1058 + modules = get_hf_file_to_dict("modules.json", model, revision)

1059 + if not modules:

1060 + return None

1061 +

1062 + if isinstance(modules, dict):

1063 + modules = modules.get("modules", [])

1064 +

1065 + dense_modules = [

1066 + m for m in modules if m.get("type") == "sentence_transformers.models.Dense"

1067 + ]

1068 + if not dense_modules:

1069 + return None

1070 +

1071 + layer_configs = []

1072 + for module in dense_modules:

1073 + folder = module.get("path", "")

1074 +

1075 + config_path = f"{folder}/config.json" if folder else "config.json"

1076 + layer_config = get_hf_file_to_dict(config_path, model, revision)

1077 + if not layer_config:

1078 + continue

1079 + layer_config["folder"] = folder

1080 + layer_configs.append(layer_config)

1081 + return layer_configs

1082 + except Exception:

1083 + return None

1084 +

1085 +

1086 def get_safetensors_params_metadata(

1087 model: str,

1088 *,