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 9, 21:48 UTC
[Bugfix] Fix Granite Vision / Don't use Siglip Pooling Head Nested Models by Default (#32299)
Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>
Alex Brooks committed
Jan 20, 2026 at 20:11 UTC
27b81e010d09cb3f23eae5e4972f53ce4c0baf46
5 files changed
+83
-16
tests/models/multimodal/generation/test_common.py
+8
@@ -397,6 +397,14 @@ VLM_TEST_SETTINGS = {
397
vllm_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}},
398
patch_hf_runner=model_utils.gemma3_patch_hf_runner,
399
),
400
+ "granite_vision": VLMTestInfo(
401
+ models=["ibm-granite/granite-vision-3.3-2b"],
402
+ test_type=(VLMTestType.IMAGE),
403
+ prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}\n<|assistant|>\n",
404
+ max_model_len=8192,
405
+ auto_cls=AutoModelForImageTextToText,
406
+ vllm_output_post_proc=model_utils.llava_image_vllm_to_hf_output,
407
+ ),
408
"glm4v": VLMTestInfo(
409
models=["zai-org/glm-4v-9b"],
410
test_type=VLMTestType.IMAGE,
tests/models/multimodal/generation/vlm_utils/model_utils.py
+4
-2
@@ -124,8 +124,10 @@ def _llava_vllm_to_hf_output(
124
if token_id != mm_token_id or output_ids[idx - 1] != mm_token_id
125
]
126
127
- assert output_str[0] == " "
128
- hf_output_str = output_str[1:]
127
+ # output_str[0] is not " " in some cases, e.g., Granite Vision,
128
+ # but for most llava based models, this is the case
129
+ hf_output_str = output_str[1:] if output_str[0] == " " else output_str
130
+
131
if hf_output_ids[-1] == eos_token_id:
132
hf_output_str = hf_output_str + tokenizer.decode(eos_token_id)
133
tests/models/registry.py
+1
@@ -692,6 +692,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
692
trust_remote_code=True,
693
min_transformers_version="5.0",
694
),
695
+ "GraniteVision": _HfExamplesInfo("ibm-granite/granite-vision-3.3-2b"),
696
"GraniteSpeechForConditionalGeneration": _HfExamplesInfo(
697
"ibm-granite/granite-speech-3.3-2b"
698
),
vllm/model_executor/models/siglip.py
+55
-13
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable, Iterable, Mapping
5
-from functools import cached_property
5
+from functools import cached_property, partial
6
from typing import Annotated, Literal
7
8
import torch
@@ -705,6 +705,7 @@ class SiglipVisionTransformer(nn.Module):
705
num_hidden_layers_override: int | None = None,
706
require_post_norm: bool | None = None,
707
prefix: str = "",
708
+ use_head: bool | None = False,
709
) -> None:
710
super().__init__()
711
@@ -738,16 +739,30 @@ class SiglipVisionTransformer(nn.Module):
739
else:
740
self.post_layernorm = None
741
741
- self.use_head = (
742
- True if not hasattr(config, "vision_use_head") else config.vision_use_head
743
- )
744
- if self.use_head:
745
- self.head = SiglipMultiheadAttentionPoolingHead(
742
+ # Fall back to the config if a bool is not provided explicitly;
743
+ # note that many config types, including SiglipVisionConfig,
744
+ # do not have vision_use_head as a defined attribute.
745
+ if isinstance(use_head, bool):
746
+ self.use_head = use_head
747
+ else:
748
+ self.use_head = (
749
+ True
750
+ if not hasattr(config, "vision_use_head")
751
+ else config.vision_use_head
752
+ )
753
+
754
+ # Only create and load the head weights if we actually need them
755
+ self.head = (
756
+ SiglipMultiheadAttentionPoolingHead(
757
config=config,
758
quant_config=quant_config,
759
multimodal_config=multimodal_config,
760
prefix=f"{prefix}.head",
761
)
762
+ if self.use_head
763
+ else None
764
+ )
765
+ self.last_hs_proc = partial(self.maybe_layer_norm_and_apply_head)
766
767
@property
768
def dtype(self):
@@ -776,23 +791,37 @@ class SiglipVisionTransformer(nn.Module):
791
return_all_hidden_states=select_layers is not None,
792
)
793
779
- if self.post_layernorm is not None:
780
- encoder_outputs = self.post_layernorm(encoder_outputs)
781
-
782
- if self.use_head:
783
- encoder_outputs = self.head(encoder_outputs)
784
-
785
- # stacks feature layers if needed
794
+ # In the case that we have multiple feature layers,
795
+ # we stack and concatenate them into a tensor.
796
+ # NOTE: post layer norm and the attention pooling head
797
+ # are handled by last_hs_proc, which runs before applying
798
+ # the vision feature selection strategy.
799
encoder_outputs = resolve_visual_encoder_outputs(
800
encoder_outputs,
801
None,
802
select_layers=select_layers,
803
max_possible_layers=self.config.num_hidden_layers,
804
+ last_hs_proc=self.last_hs_proc,
805
feature_select_strategy=feature_select_strategy,
806
)
807
808
return encoder_outputs
809
810
+ def maybe_layer_norm_and_apply_head(
811
+ self, encoder_outputs: torch.Tensor
812
+ ) -> torch.Tensor:
813
+ """Apply the post layer norm and head if they are enabled,
814
+ given the last hidden states tensor.
815
+
816
+ args:
817
+ encoder_outputs: The last hidden states from the visual encoder.
818
+ """
819
+ if self.post_layernorm is not None:
820
+ encoder_outputs = self.post_layernorm(encoder_outputs)
821
+ if self.head is not None:
822
+ encoder_outputs = self.head(encoder_outputs)
823
+ return encoder_outputs
824
+
825
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
826
stacked_params_mapping = [
827
# (param_name, shard_name, shard_id)
@@ -809,6 +838,11 @@ class SiglipVisionTransformer(nn.Module):
838
if name.startswith("post_layernorm") and self.post_layernorm is None:
839
continue
840
841
+ # if the model configuration is not going to use
842
+ # the pooling head for inference, don't load its weights
843
+ if self.head is None and name.startswith("head"):
844
+ continue
845
+
846
# omit layers when num_hidden_layers_override is set
847
if name.startswith("encoder.layers"):
848
layer_idx = int(name.split(".")[2])
@@ -841,6 +875,7 @@ class SiglipVisionModel(nn.Module):
875
num_hidden_layers_override: int | None = None,
876
require_post_norm: bool | None = None,
877
prefix: str = "",
878
+ use_head: bool | None = False,
879
) -> None:
880
super().__init__()
881
@@ -852,6 +887,7 @@ class SiglipVisionModel(nn.Module):
887
num_hidden_layers_override=num_hidden_layers_override,
888
require_post_norm=require_post_norm,
889
prefix=f"{prefix}.vision_model",
890
+ use_head=use_head,
891
)
892
893
def get_input_embeddings(self) -> nn.Module:
@@ -898,6 +934,11 @@ class SiglipVisionModel(nn.Module):
934
):
935
continue
936
937
+ # if the model configuration is not going to use
938
+ # the pooling head for inference, don't load its weights
939
+ if self.vision_model.head is None and name.startswith("vision_model.head"):
940
+ continue
941
+
942
# omit layers when num_hidden_layers_override is set
943
if name.startswith("vision_model.encoder.layers"):
944
layer_idx = int(name.split(".")[3])
@@ -1048,6 +1089,7 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
1089
quant_config=quant_config,
1090
multimodal_config=multimodal_config,
1091
prefix=maybe_prefix(prefix, "vision_model"),
1092
+ use_head=None, # Allows potential pooling head
1093
)
1094
1095
pooler_config = vllm_config.model_config.pooler_config
vllm/model_executor/models/vision.py
+15
-1
@@ -154,6 +154,7 @@ def resolve_visual_encoder_outputs(
154
*,
155
select_layers: list[int] | None = None,
156
max_possible_layers: int | None = None,
157
+ last_hs_proc: Callable[[torch.Tensor], torch.Tensor] | None = None,
158
feature_select_strategy: VisionFeatureSelectStrategy | None = None,
159
) -> torch.Tensor:
160
"""Given the outputs a visual encoder module that may correspond to the
@@ -166,6 +167,11 @@ def resolve_visual_encoder_outputs(
167
select_layers: Optional layer indices to grab from the encoder
168
outputs; if provided, encoder outputs must be a list.
169
max_possible_layers: Total layers in the fully loaded visual encoder.
170
+ last_hs_proc: Optional callable to be applied to the last layer if it
171
+ is used, e.g., pooling head for Siglip. This is done prior to
172
+ feature selection and layer normalization. If select_layers are
173
+ provided, the output of last_hs_proc must be able to be
174
+ concatenated with the other select_layers along the last dimension.
175
feature_select_strategy: Defines how to select the hidden states
176
from each layer.
177
"""
@@ -176,6 +182,11 @@ def resolve_visual_encoder_outputs(
182
"`select_layers` is not provided"
183
)
184
185
+ # Preprocess the encoder outputs as needed, e.g., map head
186
+ # and layer norm for siglip, which runs before feature selection
187
+ if last_hs_proc is not None:
188
+ encoder_outputs = last_hs_proc(encoder_outputs)
189
+
190
if feature_select_strategy is not None:
191
select_features = _get_vision_feature_selector(feature_select_strategy)
192
encoder_outputs = select_features(encoder_outputs)
@@ -205,12 +216,15 @@ def resolve_visual_encoder_outputs(
216
for layer_idx in select_layers
217
]
218
219
+ uses_last_layer = select_layers[-1] in (max_possible_layers - 1, -1)
220
+ if last_hs_proc is not None and uses_last_layer:
221
+ hs_pool[-1] = last_hs_proc(hs_pool[-1])
222
+
223
if feature_select_strategy is not None:
224
select_features = _get_vision_feature_selector(feature_select_strategy)
225
hs_pool = [select_features(hs) for hs in hs_pool]
226
227
# Apply post-norm on the final hidden state if we are using it
213
- uses_last_layer = select_layers[-1] in (max_possible_layers - 1, -1)
228
if post_layer_norm is not None and uses_last_layer:
229
hs_pool[-1] = post_layer_norm(hs_pool[-1])
230