main
py 393 lines 13.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """vLLM support for microsoft/Phi-4-reasoning-vision-15B.
4
5 Architecture: Siglip2 vision tower + MLP projector + Phi3 language model.
6 """
7
8 import math
9 from collections.abc import Iterable, Mapping, Sequence
10 from typing import Annotated, Any, Literal
11
12 import torch
13 import torch.nn as nn
14 from transformers import BatchFeature, PreTrainedConfig, Siglip2VisionConfig
15
16 from vllm.config import VllmConfig
17 from vllm.config.multimodal import MultiModalDummyOptions
18 from vllm.inputs import MultiModalDataDict
19 from vllm.logger import init_logger
20 from vllm.multimodal import MULTIMODAL_REGISTRY
21 from vllm.multimodal.inputs import (
22 MultiModalFieldConfig,
23 MultiModalKwargsItems,
24 )
25 from vllm.multimodal.parse import (
26 ImageSize,
27 MultiModalDataItems,
28 )
29 from vllm.multimodal.processing import (
30 BaseDummyInputsBuilder,
31 PromptReplacement,
32 PromptUpdate,
33 )
34 from vllm.multimodal.processing.processor import (
35 BaseMultiModalProcessor,
36 BaseProcessingInfo,
37 cached_encode,
38 )
39 from vllm.sequence import IntermediateTensors
40 from vllm.utils.tensor_schema import TensorSchema, TensorShape
41
42 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
43 from .lfm2_siglip2 import Siglip2Model
44 from .llava import LlavaMultiModalProjector
45 from .utils import (
46 AutoWeightsLoader,
47 WeightsMapper,
48 init_vllm_registered_model,
49 maybe_prefix,
50 )
51
52 logger = init_logger(__name__)
53
54 IMAGE_TOKEN_INDEX = -200
55 DEFAULT_IMAGE_TOKEN = "<image>"
56
57 # The HF processor replaces "<image>" with IMAGE_TOKEN_INDEX (-200) in input_ids.
58 # Negative token IDs cause OverflowError during decoding, so we remap to a real
59 # in-vocabulary token. The Phi-4-reasoning-vision tokenizer ships with reserved
60 # dummy tokens (<|dummy_0|> … <|dummy_83|>); we reuse the first one as the
61 # image placeholder. This mirrors how Phi-3-vision uses its dedicated <|image|>
62 # token (ID 32044).
63 _IMAGE_TOKEN_ID = 100256 # <|dummy_0|> in the Phi-4 tokenizer
64
65
66 # ---------------------------------------------------------------------------
67 # Processing
68 # ---------------------------------------------------------------------------
69
70
71 class Phi4SiglipProcessingInfo(BaseProcessingInfo):
72 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
73 return {"image": None}
74
75 def _get_vision_config(self) -> dict:
76 return self.get_hf_config().vision_config # type: ignore[attr-defined]
77
78 def _get_patch_size(self) -> int:
79 vc = self._get_vision_config()
80 if isinstance(vc, dict):
81 return vc.get("patch_size", 16)
82 return getattr(vc, "patch_size", 16)
83
84 def _get_max_num_patches(self) -> int:
85 return getattr(self.get_hf_config(), "max_num_patches", 3600)
86
87 def _get_min_num_patches(self) -> int:
88 return getattr(self.get_hf_config(), "min_num_patches", 256)
89
90 def get_num_image_tokens(
91 self,
92 *,
93 image_width: int,
94 image_height: int,
95 ) -> int:
96 patch_size = self._get_patch_size()
97 min_patches = self._get_min_num_patches()
98 max_patches = self._get_max_num_patches()
99
100 num_patches_h = image_height // patch_size
101 num_patches_w = image_width // patch_size
102 num_patches = max(num_patches_h * num_patches_w, 1)
103 num_patches = max(min(num_patches, max_patches), min_patches)
104 return num_patches
105
106 def get_image_size_with_most_features(self) -> ImageSize:
107 patch_size = self._get_patch_size()
108 max_patches = self._get_max_num_patches()
109 side = int(math.sqrt(max_patches)) * patch_size
110 return ImageSize(width=side, height=side)
111
112 def get_mm_max_tokens_per_item(
113 self, seq_len: int, mm_counts: Mapping[str, int]
114 ) -> Mapping[str, int]:
115 return {"image": self._get_max_num_patches()}
116
117
118 class Phi4SiglipDummyInputsBuilder(
119 BaseDummyInputsBuilder[Phi4SiglipProcessingInfo],
120 ):
121 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
122 num_images = mm_counts.get("image", 0)
123 return DEFAULT_IMAGE_TOKEN * num_images
124
125 def get_dummy_mm_data(
126 self,
127 seq_len: int,
128 mm_counts: Mapping[str, int],
129 mm_options: MultiModalDummyOptions,
130 ) -> MultiModalDataDict:
131 size = self.info.get_image_size_with_most_features()
132 return {
133 "image": self._get_dummy_images(
134 width=size.width,
135 height=size.height,
136 num_images=mm_counts.get("image", 0),
137 overrides=mm_options.get("image"),
138 ),
139 }
140
141
142 class Phi4SiglipMultiModalProcessor(
143 BaseMultiModalProcessor[Phi4SiglipProcessingInfo],
144 ):
145 def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
146 return self.dummy_inputs.get_dummy_text(mm_counts)
147
148 def _get_mm_fields_config(
149 self,
150 hf_inputs: BatchFeature,
151 hf_processor_mm_kwargs: Mapping[str, object],
152 ) -> Mapping[str, MultiModalFieldConfig]:
153 return dict(
154 pixel_values=MultiModalFieldConfig.batched("image"),
155 pixel_attention_mask=MultiModalFieldConfig.batched("image"),
156 spatial_shapes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
157 )
158
159 def _get_prompt_updates(
160 self,
161 mm_items: MultiModalDataItems,
162 hf_processor_mm_kwargs: Mapping[str, Any],
163 out_mm_kwargs: MultiModalKwargsItems,
164 ) -> Sequence[PromptUpdate]:
165 tokenizer = self.info.get_tokenizer()
166
167 def get_replacement(item_idx: int):
168 # Read the actual patch grid from the NaFlex processor's
169 # spatial_shapes output (same pattern as LFM2-VL). This avoids
170 # predicting from raw image dimensions, which can diverge from
171 # the NaFlex resize/tile logic.
172 out_item = out_mm_kwargs["image"][item_idx]
173 spatial_shapes = out_item["spatial_shapes"].data
174 assert isinstance(spatial_shapes, torch.Tensor)
175 num_tokens = int(spatial_shapes.prod().item())
176 return [_IMAGE_TOKEN_ID] * num_tokens
177
178 return [
179 PromptReplacement(
180 modality="image",
181 target=cached_encode(
182 tokenizer, DEFAULT_IMAGE_TOKEN, add_special_tokens=False
183 ),
184 replacement=get_replacement,
185 ),
186 ]
187
188
189 # ---------------------------------------------------------------------------
190 # Input schemas
191 # ---------------------------------------------------------------------------
192
193
194 class Phi4SiglipImagePixelInputs(TensorSchema):
195 """Dimensions:
196 - bn: Batch size * number of images
197 - d: Max number of patches (padded across images in the batch)
198 - fd: Features per patch (patch_size * patch_size * channels)
199 """
200
201 type: Literal["pixel_values"] = "pixel_values"
202 pixel_values: Annotated[torch.Tensor, TensorShape("bn", "d", "fd")]
203 pixel_attention_mask: Annotated[torch.Tensor, TensorShape("bn", "d")]
204 spatial_shapes: Annotated[torch.Tensor, TensorShape("bn", 2)]
205
206
207 # ---------------------------------------------------------------------------
208 # Model
209 # ---------------------------------------------------------------------------
210
211
212 @MULTIMODAL_REGISTRY.register_processor(
213 Phi4SiglipMultiModalProcessor,
214 info=Phi4SiglipProcessingInfo,
215 dummy_inputs=Phi4SiglipDummyInputsBuilder,
216 )
217 class Phi4ForCausalLMV(nn.Module, SupportsMultiModal, SupportsPP):
218 hf_to_vllm_mapper = WeightsMapper(
219 orig_to_new_prefix={
220 "model.vision_tower.vision_tower.vision_model.head.": None,
221 "model.vision_tower.vision_tower.": "vision_tower.",
222 "model.mm_projector.0.": "multi_modal_projector.linear_1.",
223 "model.mm_projector.2.": "multi_modal_projector.linear_2.",
224 "lm_head.": "language_model.lm_head.",
225 "model.": "language_model.model.",
226 },
227 )
228
229 @classmethod
230 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
231 if modality.startswith("image"):
232 return DEFAULT_IMAGE_TOKEN
233 raise ValueError("Only image modality is supported")
234
235 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
236 super().__init__()
237
238 config: PreTrainedConfig = vllm_config.model_config.hf_config
239 quant_config = vllm_config.quant_config
240 self.config = config
241
242 vision_config_dict: dict = getattr(config, "vision_config", {})
243 if isinstance(vision_config_dict, dict):
244 if "patch_size" not in vision_config_dict:
245 vision_config_dict["patch_size"] = 16
246 siglip2_config = Siglip2VisionConfig(**vision_config_dict)
247 else:
248 siglip2_config = vision_config_dict
249
250 vision_hidden_size: int = config.mm_hidden_size # type: ignore[attr-defined]
251 text_hidden_size: int = config.hidden_size # type: ignore[attr-defined]
252
253 with self._mark_tower_model(vllm_config, "image"):
254 layer_idx = -2
255 num_hidden_layers = siglip2_config.num_hidden_layers + layer_idx + 1
256
257 self.vision_tower = Siglip2Model(
258 siglip2_config,
259 quant_config=quant_config,
260 num_hidden_layers_override=num_hidden_layers,
261 require_post_norm=False,
262 prefix=maybe_prefix(prefix, "vision_tower"),
263 )
264 self.multi_modal_projector = LlavaMultiModalProjector(
265 vision_hidden_size=vision_hidden_size,
266 text_hidden_size=text_hidden_size,
267 projector_hidden_act="gelu",
268 multimodal_projector_bias=True,
269 quant_config=quant_config,
270 prefix=maybe_prefix(prefix, "multi_modal_projector"),
271 )
272
273 with self._mark_language_model(vllm_config):
274 self.language_model = init_vllm_registered_model(
275 vllm_config=vllm_config,
276 hf_config=config,
277 prefix=maybe_prefix(prefix, "language_model"),
278 architectures=["Phi3ForCausalLM"],
279 )
280
281 self.make_empty_intermediate_tensors = (
282 self.language_model.make_empty_intermediate_tensors
283 )
284
285 self.configure_mm_token_handling(
286 vocab_size=config.vocab_size, # type: ignore[attr-defined]
287 mm_token_ids=[_IMAGE_TOKEN_ID],
288 )
289
290 def _packed_from_padded(
291 self,
292 pixel_values: torch.Tensor,
293 pixel_attention_mask: torch.Tensor,
294 spatial_shapes: torch.Tensor,
295 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
296 """Convert padded NaFlex tensors to packed format for Siglip2Model."""
297 valid_counts = pixel_attention_mask.sum(dim=1).to(torch.int32)
298 pixel_values_packed = pixel_values[pixel_attention_mask.bool()]
299 cu_seqlens = torch.zeros(
300 len(valid_counts) + 1,
301 dtype=torch.int32,
302 device=pixel_values.device,
303 )
304 cu_seqlens[1:] = valid_counts.cumsum(0)
305 max_seqlen = valid_counts.max()
306 return (
307 pixel_values_packed,
308 spatial_shapes,
309 cu_seqlens,
310 max_seqlen,
311 )
312
313 def _parse_and_validate_image_input(
314 self, **kwargs: object
315 ) -> Phi4SiglipImagePixelInputs | None:
316 pixel_values = kwargs.pop("pixel_values", None)
317 pixel_attention_mask = kwargs.pop("pixel_attention_mask", None)
318 spatial_shapes = kwargs.pop("spatial_shapes", None)
319 if pixel_values is None:
320 return None
321
322 return Phi4SiglipImagePixelInputs(
323 type="pixel_values",
324 pixel_values=pixel_values,
325 pixel_attention_mask=pixel_attention_mask,
326 spatial_shapes=spatial_shapes,
327 )
328
329 def _process_image_input(
330 self, image_input: Phi4SiglipImagePixelInputs
331 ) -> MultiModalEmbeddings:
332 pixel_values = image_input["pixel_values"]
333 pixel_attention_mask = image_input["pixel_attention_mask"]
334 spatial_shapes = image_input["spatial_shapes"]
335
336 (
337 pixel_values_packed,
338 spatial_shapes_packed,
339 cu_seqlens,
340 max_seqlen,
341 ) = self._packed_from_padded(pixel_values, pixel_attention_mask, spatial_shapes)
342
343 vision_features = self.vision_tower(
344 pixel_values_packed=pixel_values_packed,
345 spatial_shapes=spatial_shapes_packed,
346 cu_seqlens=cu_seqlens,
347 max_seqlen=max_seqlen,
348 select_layers=[-2],
349 )
350
351 if vision_features.dim() == 3:
352 vision_features = vision_features.squeeze(0)
353
354 image_features = self.multi_modal_projector(vision_features)
355
356 valid_counts = pixel_attention_mask.sum(dim=1).tolist()
357 return torch.split(image_features, [int(c) for c in valid_counts])
358
359 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
360 image_input = self._parse_and_validate_image_input(**kwargs)
361 if image_input is None:
362 return []
363
364 return self._process_image_input(image_input)
365
366 def forward(
367 self,
368 input_ids: torch.Tensor | None,
369 positions: torch.Tensor,
370 intermediate_tensors: IntermediateTensors | None = None,
371 inputs_embeds: torch.Tensor | None = None,
372 **kwargs: object,
373 ) -> torch.Tensor | IntermediateTensors:
374 if intermediate_tensors is not None:
375 inputs_embeds = None
376
377 hidden_states = self.language_model.model(
378 input_ids,
379 positions,
380 intermediate_tensors,
381 inputs_embeds=inputs_embeds,
382 )
383 return hidden_states
384
385 def compute_logits(
386 self,
387 hidden_states: torch.Tensor,
388 ) -> torch.Tensor | None:
389 return self.language_model.compute_logits(hidden_states)
390
391 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
392 loader = AutoWeightsLoader(self)
393 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)