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, 22:35 UTC
Enable headless models for pooling in the Transformers backend (#21767)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Aug 1, 2025 at 18:31 UTC
38c8bce8b652df87d111c04ddf849c38615000c7
5 files changed
+44
-9
tests/models/registry.py
+1
@@ -525,6 +525,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
525
}
526
527
_TRANSFORMERS_BACKEND_MODELS = {
528
+ "TransformersModel": _HfExamplesInfo("Qwen/Qwen3-Embedding-0.6B"),
529
"TransformersForCausalLM": _HfExamplesInfo("hmellor/Ilama-3.2-1B", trust_remote_code=True), # noqa: E501
530
"TransformersForMultimodalLM": _HfExamplesInfo("OpenGVLab/InternVL3-1B-hf"),
531
}
tests/models/test_transformers.py
+22
-6
@@ -34,8 +34,7 @@ def check_implementation(
34
35
with runner_test(model, **kwargs_test, **kwargs) as model_test:
36
model_config = model_test.llm.llm_engine.model_config
37
- assert model_config.architecture == (
38
- model_config._get_transformers_backend_cls())
37
+ assert model_config.using_transformers_backend()
38
39
outputs_test = model_test.generate_greedy_logprobs(*args)
40
@@ -135,8 +134,7 @@ def test_quantization(
134
enforce_eager=True,
135
**quantization_kwargs) as vllm_model: # type: ignore[arg-type]
136
model_config = vllm_model.llm.llm_engine.model_config
138
- assert model_config.architecture == (
139
- model_config._get_transformers_backend_cls())
137
+ assert model_config.using_transformers_backend()
138
139
transformers_outputs = vllm_model.generate_greedy_logprobs(
140
example_prompts, max_tokens=max_tokens, num_logprobs=num_logprobs)
@@ -149,6 +147,25 @@ def test_quantization(
147
)
148
149
150
+@pytest.mark.parametrize(
151
+ "model",
152
+ [
153
+ # Layers live in `layers`
154
+ "Qwen/Qwen3-Embedding-0.6B",
155
+ # Layers live in `model.layers`
156
+ "meta-llama/Llama-3.2-1B-Instruct"
157
+ ],
158
+)
159
+def test_embed_loading(vllm_runner, model):
160
+ with vllm_runner(model,
161
+ max_model_len=1024,
162
+ enforce_eager=True,
163
+ runner="pooling",
164
+ model_impl="transformers") as model_test:
165
+ model_config = model_test.llm.llm_engine.model_config
166
+ assert model_config.using_transformers_backend()
167
+
168
+
169
@pytest.mark.parametrize(
170
"model",
171
["jason9693/Qwen2.5-1.5B-apeach"],
@@ -169,8 +186,7 @@ def test_classify(
186
dtype=dtype,
187
model_impl="transformers") as vllm_model:
188
model_config = vllm_model.llm.llm_engine.model_config
172
- assert model_config.architecture == (
173
- model_config._get_transformers_backend_cls())
189
+ assert model_config.using_transformers_backend()
190
191
vllm_outputs = vllm_model.classify(example_prompts)
192
vllm/config.py
+7
-2
@@ -812,12 +812,17 @@ class ModelConfig:
812
def _get_transformers_backend_cls(self) -> str:
813
"""Determine which Transformers backend class will be used if
814
`model_impl` is set to `transformers` or `auto`."""
815
+ if getattr(self, "runner_type", self.runner) == "pooling":
816
+ return "TransformersModel"
817
if self.hf_config != self.hf_text_config:
818
# If 'hf_text_config' is the same as 'hf_config'. If not, it is
819
# probably a composite config, i.e. multimodal
820
return "TransformersForMultimodalLM"
819
- else:
820
- return "TransformersForCausalLM"
821
+ return "TransformersForCausalLM"
822
+
823
+ def using_transformers_backend(self) -> bool:
824
+ """Check if the model is using the Transformers backend class."""
825
+ return self.architecture == self._get_transformers_backend_cls()
826
827
@property
828
def registry(self):
vllm/model_executor/models/registry.py
+2
-1
@@ -270,8 +270,9 @@ _TRANSFORMERS_SUPPORTED_MODELS = {
270
}
271
272
_TRANSFORMERS_BACKEND_MODELS = {
273
- "TransformersForMultimodalLM": ("transformers", "TransformersForMultimodalLM"), # noqa: E501
273
+ "TransformersModel": ("transformers", "TransformersModel"),
274
"TransformersForCausalLM": ("transformers", "TransformersForCausalLM"),
275
+ "TransformersForMultimodalLM": ("transformers", "TransformersForMultimodalLM"), # noqa: E501
276
}
277
# yapf: enable
278
vllm/model_executor/models/transformers.py
+12
@@ -651,6 +651,18 @@ class TransformersBase(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
651
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
652
653
654
+@support_torch_compile
655
+class TransformersModel(TransformersBase):
656
+ hf_to_vllm_mapper = WeightsMapper(
657
+ orig_to_new_prefix={
658
+ # Add `model.` prefix for base model checkpoints
659
+ "": "model.",
660
+ # Remove `model.` from places it should not be
661
+ "model.model.": "model.",
662
+ "model.score": "score",
663
+ })
664
+
665
+
666
@support_torch_compile
667
class TransformersForCausalLM(TransformersBase):
668