@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e090b7b45b
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

Enable conversion of multimodal models to pooling tasks (#24451)

Signed-off-by: Max de Bayser <mbayser@br.ibm.com>

Maximilien de Bayser committed Sep 12, 2025 at 00:30 UTC e090b7b45b015e9a4158c6676da802e04f1d88f1
5 files changed +266 -59
tests/models/language/pooling/test_mm_classifier_conversion.py new
+114
@@ -0,0 +1,114 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.platforms import current_platform
5 +
6 +
7 +def test_idefics_multimodal(
8 + vllm_runner,
9 + monkeypatch,
10 +) -> None:
11 + if current_platform.is_rocm():
12 + # ROCm Triton FA does not currently support sliding window attention
13 + # switch to use ROCm CK FA backend
14 + monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")
15 +
16 + prompts = [
17 + "Hello, my name is",
18 + "The president of the United States is",
19 + "The capital of France is",
20 + "The future of AI is",
21 + ]
22 +
23 + with vllm_runner(model_name="HuggingFaceM4/Idefics3-8B-Llama3",
24 + runner="pooling",
25 + task="classify",
26 + convert="classify",
27 + load_format="dummy",
28 + max_model_len=512,
29 + enforce_eager=True,
30 + tensor_parallel_size=1,
31 + disable_log_stats=True,
32 + dtype="bfloat16") as vllm_model:
33 + llm = vllm_model.get_llm()
34 + outputs = llm.classify(prompts)
35 + for output in outputs:
36 + assert len(output.outputs.probs) == 2
37 +
38 +
39 +def update_config(config):
40 + config.text_config.update({
41 + "architectures": ["Gemma3ForSequenceClassification"],
42 + "classifier_from_token": ["A", "B", "C", "D", "E"],
43 + "method":
44 + "no_post_processing",
45 + "id2label": {
46 + "A": "Chair",
47 + "B": "Couch",
48 + "C": "Table",
49 + "D": "Bed",
50 + "E": "Cupboard"
51 + },
52 + })
53 + return config
54 +
55 +
56 +def test_gemma_multimodal(
57 + vllm_runner,
58 + monkeypatch,
59 +) -> None:
60 + if current_platform.is_rocm():
61 + # ROCm Triton FA does not currently support sliding window attention
62 + # switch to use ROCm CK FA backend
63 + monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")
64 +
65 + messages = [{
66 + "role":
67 + "system",
68 + "content":
69 + """
70 + You are a helpful assistant. You will be given a product description
71 + which may also include an image. Classify the following product into
72 + one of the categories:
73 +
74 + A = chair
75 + B = couch
76 + C = table
77 + D = bed
78 + E = cupboard
79 +
80 + You'll answer with exactly one letter (A, B, C, D, or E)."""
81 + }, {
82 + "role":
83 + "user",
84 + "content": [{
85 + "type": "image_url",
86 + "image_url": {
87 + "url":
88 + "https://upload.wikimedia.org/wikipedia/commons/c/c6/Set_of_fourteen_side_chairs_MET_DP110780.jpg"
89 + }
90 + }, {
91 + "type": "text",
92 + "text": "A fine 19th century piece of furniture."
93 + }]
94 + }]
95 +
96 + with vllm_runner(model_name="google/gemma-3-4b-it",
97 + runner="pooling",
98 + task="classify",
99 + convert="classify",
100 + load_format="auto",
101 + hf_overrides=update_config,
102 + override_pooler_config={"pooling_type": "LAST"},
103 + max_model_len=512,
104 + enforce_eager=True,
105 + tensor_parallel_size=1,
106 + disable_log_stats=True,
107 + dtype="bfloat16") as vllm_model:
108 +
109 + llm = vllm_model.get_llm()
110 + prompts = llm.preprocess_chat(messages)
111 +
112 + result = llm.classify(prompts)
113 + assert result[0].outputs.probs[0] > 0.95
114 + assert all(c < 0.05 for c in result[0].outputs.probs[1:])
\ No newline at end of file
vllm/entrypoints/llm.py
+90 -52
@@ -703,13 +703,10 @@ class LLM:
703
704 return outputs
705
706 - def chat(
706 + def preprocess_chat(
707 self,
708 messages: Union[list[ChatCompletionMessageParam],
709 list[list[ChatCompletionMessageParam]]],
710 - sampling_params: Optional[Union[SamplingParams,
711 - list[SamplingParams]]] = None,
712 - use_tqdm: Union[bool, Callable[..., tqdm]] = True,
710 lora_request: Optional[LoRARequest] = None,
711 chat_template: Optional[str] = None,
712 chat_template_content_format: ChatTemplateContentFormatOption = "auto",
@@ -718,56 +715,16 @@ class LLM:
715 tools: Optional[list[dict[str, Any]]] = None,
716 chat_template_kwargs: Optional[dict[str, Any]] = None,
717 mm_processor_kwargs: Optional[dict[str, Any]] = None,
721 - ) -> list[RequestOutput]:
718 + ) -> list[TokensPrompt]:
719 """
723 - Generate responses for a chat conversation.
724 -
725 - The chat conversation is converted into a text prompt using the
726 - tokenizer and calls the [generate][vllm.LLM.generate] method to generate
727 - the responses.
728 -
729 - Multi-modal inputs can be passed in the same way you would pass them
730 - to the OpenAI API.
731 -
732 - Args:
733 - messages: A list of conversations or a single conversation.
734 -
735 - - Each conversation is represented as a list of messages.
736 - - Each message is a dictionary with 'role' and 'content' keys.
737 -
738 - sampling_params: The sampling parameters for text generation.
739 - If None, we use the default sampling parameters. When it
740 - is a single value, it is applied to every prompt. When it
741 - is a list, the list must have the same length as the
742 - prompts and it is paired one by one with the prompt.
743 - use_tqdm: If `True`, shows a tqdm progress bar.
744 - If a callable (e.g., `functools.partial(tqdm, leave=False)`),
745 - it is used to create the progress bar.
746 - If `False`, no progress bar is created.
747 - lora_request: LoRA request to use for generation, if any.
748 - chat_template: The template to use for structuring the chat.
749 - If not provided, the model's default chat template will be used.
750 - chat_template_content_format: The format to render message content.
751 -
752 - - "string" will render the content as a string.
753 - Example: `"Who are you?"`
754 - - "openai" will render the content as a list of dictionaries,
755 - similar to OpenAI schema.
756 - Example: `[{"type": "text", "text": "Who are you?"}]`
757 -
758 - add_generation_prompt: If True, adds a generation template
759 - to each message.
760 - continue_final_message: If True, continues the final message in
761 - the conversation instead of starting a new one. Cannot be
762 - `True` if `add_generation_prompt` is also `True`.
763 - chat_template_kwargs: Additional kwargs to pass to the chat
764 - template.
765 - mm_processor_kwargs: Multimodal processor kwarg overrides for this
766 - chat request. Only used for offline requests.
720 + Generate prompt for a chat conversation. The pre-processed
721 + prompt can then be used as input for the other LLM methods.
722
723 + Refer to `chat` for a complete description of the arguments.
724 Returns:
769 - A list of `RequestOutput` objects containing the generated
770 - responses in the same order as the input messages.
725 + A list of `TokensPrompts` objects containing the tokenized
726 + prompt after chat template interpolation, and the
727 + pre-processed multi-modal inputs.
728 """
729 list_of_messages: list[list[ChatCompletionMessageParam]]
730
@@ -800,7 +757,7 @@ class LLM:
757 )
758 _chat_template_kwargs.update(chat_template_kwargs or {})
759
803 - prompts: list[Union[TokensPrompt, TextPrompt]] = []
760 + prompts: list[TokensPrompt] = []
761
762 for msgs in list_of_messages:
763 # NOTE: _parse_chat_message_content_parts() currently doesn't
@@ -844,6 +801,87 @@ class LLM:
801
802 prompts.append(prompt)
803
804 + return prompts
805 +
806 + def chat(
807 + self,
808 + messages: Union[list[ChatCompletionMessageParam],
809 + list[list[ChatCompletionMessageParam]]],
810 + sampling_params: Optional[Union[SamplingParams,
811 + list[SamplingParams]]] = None,
812 + use_tqdm: Union[bool, Callable[..., tqdm]] = True,
813 + lora_request: Optional[LoRARequest] = None,
814 + chat_template: Optional[str] = None,
815 + chat_template_content_format: ChatTemplateContentFormatOption = "auto",
816 + add_generation_prompt: bool = True,
817 + continue_final_message: bool = False,
818 + tools: Optional[list[dict[str, Any]]] = None,
819 + chat_template_kwargs: Optional[dict[str, Any]] = None,
820 + mm_processor_kwargs: Optional[dict[str, Any]] = None,
821 + ) -> list[RequestOutput]:
822 + """
823 + Generate responses for a chat conversation.
824 +
825 + The chat conversation is converted into a text prompt using the
826 + tokenizer and calls the [generate][vllm.LLM.generate] method to generate
827 + the responses.
828 +
829 + Multi-modal inputs can be passed in the same way you would pass them
830 + to the OpenAI API.
831 +
832 + Args:
833 + messages: A list of conversations or a single conversation.
834 +
835 + - Each conversation is represented as a list of messages.
836 + - Each message is a dictionary with 'role' and 'content' keys.
837 +
838 + sampling_params: The sampling parameters for text generation.
839 + If None, we use the default sampling parameters. When it
840 + is a single value, it is applied to every prompt. When it
841 + is a list, the list must have the same length as the
842 + prompts and it is paired one by one with the prompt.
843 + use_tqdm: If `True`, shows a tqdm progress bar.
844 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),
845 + it is used to create the progress bar.
846 + If `False`, no progress bar is created.
847 + lora_request: LoRA request to use for generation, if any.
848 + chat_template: The template to use for structuring the chat.
849 + If not provided, the model's default chat template will be used.
850 + chat_template_content_format: The format to render message content.
851 +
852 + - "string" will render the content as a string.
853 + Example: `"Who are you?"`
854 + - "openai" will render the content as a list of dictionaries,
855 + similar to OpenAI schema.
856 + Example: `[{"type": "text", "text": "Who are you?"}]`
857 +
858 + add_generation_prompt: If True, adds a generation template
859 + to each message.
860 + continue_final_message: If True, continues the final message in
861 + the conversation instead of starting a new one. Cannot be
862 + `True` if `add_generation_prompt` is also `True`.
863 + chat_template_kwargs: Additional kwargs to pass to the chat
864 + template.
865 + mm_processor_kwargs: Multimodal processor kwarg overrides for this
866 + chat request. Only used for offline requests.
867 +
868 + Returns:
869 + A list of `RequestOutput` objects containing the generated
870 + responses in the same order as the input messages.
871 + """
872 +
873 + prompts = self.preprocess_chat(
874 + messages=messages,
875 + lora_request=lora_request,
876 + chat_template=chat_template,
877 + chat_template_content_format=chat_template_content_format,
878 + add_generation_prompt=add_generation_prompt,
879 + continue_final_message=continue_final_message,
880 + tools=tools,
881 + chat_template_kwargs=chat_template_kwargs,
882 + mm_processor_kwargs=mm_processor_kwargs,
883 + )
884 +
885 return self.generate(
886 prompts,
887 sampling_params=sampling_params,
vllm/model_executor/model_loader/utils.py
+14 -4
@@ -19,10 +19,11 @@ from vllm.logger import init_logger
19 from vllm.model_executor.layers.linear import QKVCrossParallelLinear
20 from vllm.model_executor.layers.quantization.base_config import (
21 QuantizationConfig, QuantizeMethodBase)
22 -from vllm.model_executor.models.adapters import (as_embedding_model,
23 - as_reward_model,
24 - as_seq_cls_model)
25 -from vllm.model_executor.models.interfaces import SupportsQuant
22 +from vllm.model_executor.models.adapters import (
23 + as_embedding_model, as_reward_model, as_seq_cls_model,
24 + try_create_mm_pooling_model_cls)
25 +from vllm.model_executor.models.interfaces import (SupportsQuant,
26 + supports_multimodal)
27 from vllm.utils import is_pin_memory_available
28
29 logger = init_logger(__name__)
@@ -183,6 +184,15 @@ def get_model_architecture(
184 "performance may not be optimal.", arch)
185
186 convert_type = model_config.convert_type
187 + if convert_type != "none" and supports_multimodal(model_cls):
188 + logger.debug_once("Detected conversion of Multi Modal model.")
189 + converted = try_create_mm_pooling_model_cls(model_cls)
190 + if converted is not None:
191 + logger.debug_once("Creating wrapper class to forward pooler.")
192 + return converted, arch
193 + else:
194 + logger.debug_once("Attempting direct conversion.")
195 +
196 if convert_type == "none":
197 pass
198 elif convert_type == "embed":
vllm/model_executor/models/adapters.py
+43 -2
@@ -1,12 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import ast
5 +import inspect
6 from collections.abc import Iterable
7 from typing import TYPE_CHECKING, Any, Optional, TypeVar, cast
8
9 import torch
10 import torch.nn as nn
11
12 +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
@@ -129,6 +132,41 @@ def _get_pooling_model_name(orig_model_name: str, pooling_suffix: str) -> str:
132 return model_name + pooling_suffix
133
134
135 +def try_create_mm_pooling_model_cls(orig_cls: _T) -> _T:
136 +
137 + class CallVisitor(ast.NodeVisitor):
138 +
139 + def __init__(self):
140 + self.calls = []
141 +
142 + def visit_Call(self, node):
143 + if isinstance(node.func, ast.Name):
144 + self.calls.append(node.func.id)
145 + self.generic_visit(node)
146 +
147 + visitor = CallVisitor()
148 + visitor.visit(ast.parse(inspect.getsource(orig_cls)))
149 + if "init_vllm_registered_model" not in visitor.calls:
150 + return None
151 +
152 + class ModelForPooling(orig_cls, VllmModelForPooling):
153 +
154 + is_pooling_model = True
155 +
156 + def __init__(
157 + self,
158 + *,
159 + vllm_config: "VllmConfig",
160 + prefix: str = "",
161 + **kwargs: Any,
162 + ) -> None:
163 + super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
164 +
165 + self.pooler = self.get_language_model().pooler
166 +
167 + return ModelForPooling # type: ignore
168 +
169 +
170 def _create_pooling_model_cls(orig_cls: _T) -> _T:
171 # Lazy import
172 from .utils import AutoWeightsLoader, WeightsMapper
@@ -399,6 +437,7 @@ def load_weights_using_from_2_way_softmax(
437 from vllm.model_executor.models.utils import AutoWeightsLoader
438
439 model_config = model.vllm_config.model_config
440 +
441 tokens = getattr(model.config, "classifier_from_token", [])
442 tokens = cast(list[int], tokens)
443 assert len(tokens) == 2
@@ -406,9 +445,10 @@ def load_weights_using_from_2_way_softmax(
445 if model.config.tie_word_embeddings:
446 model.lm_head = model.model.embed_tokens
447 else:
448 + quant_config = model.vllm_config.quant_config
449 model.lm_head = ParallelLMHead(model.config.vocab_size,
450 model.config.hidden_size,
411 - quant_config=model.quant_config)
451 + quant_config=quant_config)
452
453 loader = AutoWeightsLoader(model)
454 loaded_weights = loader.load_weights(weights)
@@ -452,9 +492,10 @@ def load_weights_no_post_processing(model,
492 if model.config.tie_word_embeddings:
493 model.lm_head = model.model.embed_tokens
494 else:
495 + quant_config = model.vllm_config.quant_config
496 model.lm_head = ParallelLMHead(model.config.vocab_size,
497 model.config.hidden_size,
457 - quant_config=model.quant_config)
498 + quant_config=quant_config)
499
500 loader = AutoWeightsLoader(model)
501 loaded_weights = loader.load_weights(weights)
vllm/model_executor/models/gemma3_mm.py
+5 -1
@@ -512,7 +512,11 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
512 architectures=["Gemma3ForCausalLM"],
513 )
514 logit_scale = getattr(config, "logit_scale", 1.0)
515 - self.language_model.logits_processor.scale *= logit_scale
515 +
516 + if hasattr(self.language_model, "logits_processor"):
517 + # The logits processor can be unset if we're using
518 + # automatic conversion to pooling model.
519 + self.language_model.logits_processor.scale *= logit_scale
520
521 self.make_empty_intermediate_tensors = (
522 self.language_model.make_empty_intermediate_tensors)