main
py 1,157 lines 38.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable, Iterable, Mapping
5 from functools import cached_property, partial
6 from typing import Annotated, Literal
7
8 import torch
9 from torch import nn
10 from transformers import (
11 BatchFeature,
12 SiglipConfig,
13 SiglipProcessor,
14 SiglipTextConfig,
15 SiglipVisionConfig,
16 )
17
18 from vllm.config import VllmConfig
19 from vllm.config.multimodal import MultiModalDummyOptions
20 from vllm.distributed import divide, get_tensor_model_parallel_world_size
21 from vllm.inputs import MultiModalDataDict, MultiModalInput
22 from vllm.model_executor.layers.activation import get_act_fn
23 from vllm.model_executor.layers.attention import (
24 EncoderOnlyAttention,
25 MMEncoderAttention,
26 )
27 from vllm.model_executor.layers.conv import Conv2dLayer
28 from vllm.model_executor.layers.layernorm import LayerNorm
29 from vllm.model_executor.layers.linear import (
30 ColumnParallelLinear,
31 QKVParallelLinear,
32 RowParallelLinear,
33 )
34 from vllm.model_executor.layers.pooler import DispatchPooler
35 from vllm.model_executor.layers.quantization import QuantizationConfig
36 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
37 from vllm.multimodal import MULTIMODAL_REGISTRY
38 from vllm.multimodal.inputs import (
39 MultiModalFieldConfig,
40 MultiModalKwargsItems,
41 )
42 from vllm.multimodal.parse import (
43 ImageProcessorItems,
44 ImageSize,
45 MultiModalDataItems,
46 )
47 from vllm.multimodal.processing import (
48 BaseDummyInputsBuilder,
49 BaseMultiModalProcessor,
50 BaseProcessingInfo,
51 ProcessorInputs,
52 PromptIndexTargets,
53 PromptReplacement,
54 PromptUpdate,
55 TimingContext,
56 )
57 from vllm.sequence import IntermediateTensors
58 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
59 from vllm.utils.tensor_schema import TensorSchema, TensorShape
60 from vllm.utils.torch_utils import async_tensor_h2d
61
62 from .clip import dual_encoder_has_text_tokens, merge_dual_encoder_text_and_vision
63 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant
64 from .interfaces_base import default_pooling_type
65 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
66 from .vision import (
67 VisionEncoderInfo,
68 VisionFeatureSelectStrategy,
69 VisionFeatureSelectStrategyStr,
70 get_num_selected_vision_tokens,
71 is_vit_use_data_parallel,
72 resolve_visual_encoder_outputs,
73 )
74
75
76 class SiglipImagePixelInputs(TensorSchema):
77 """Dimensions:
78 - bn: Batch size * number of images
79 - c: Number of channels (3)
80 - h: Height of each image
81 - w: Width of each image
82 """
83
84 type: Literal["pixel_values"]
85 data: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
86
87
88 _POOLING_TYPE_TO_STRATEGY: dict[str, VisionFeatureSelectStrategyStr] = {
89 "MEAN": "full",
90 "ALL": "full",
91 "CLS": "class",
92 }
93
94
95 def _get_vision_feature_select_strategy(
96 pooling_type: str,
97 ) -> VisionFeatureSelectStrategyStr:
98 try:
99 return _POOLING_TYPE_TO_STRATEGY[pooling_type]
100 except KeyError:
101 raise ValueError(
102 f"No feature selection strategy is defined for "
103 f"pooling_type: {pooling_type!r}"
104 ) from None
105
106
107 class SiglipProcessingInfo(BaseProcessingInfo):
108 def get_hf_config(self):
109 return self.ctx.get_hf_config(SiglipConfig)
110
111 def get_vision_encoder_info(self):
112 return SiglipEncoderInfo(self.get_hf_config())
113
114 def get_hf_processor(self, **kwargs: object):
115 return self.ctx.get_hf_processor(SiglipProcessor, **kwargs)
116
117 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
118 return {"image": 1}
119
120 def get_num_image_tokens(
121 self,
122 *,
123 image_width: int,
124 image_height: int,
125 ) -> int:
126 vision_encoder_info = self.get_vision_encoder_info()
127
128 pooler_config = self.ctx.model_config.pooler_config
129 assert pooler_config is not None
130 pooling_type = pooler_config.seq_pooling_type
131 assert pooling_type is not None
132
133 return get_num_selected_vision_tokens(
134 vision_encoder_info.get_num_image_tokens(
135 image_width=image_width,
136 image_height=image_height,
137 ),
138 _get_vision_feature_select_strategy(pooling_type),
139 )
140
141 def get_image_size_with_most_features(self) -> ImageSize:
142 vision_encoder_info = self.get_vision_encoder_info()
143 width = height = vision_encoder_info.get_image_size()
144 return ImageSize(width=width, height=height)
145
146 def get_max_image_tokens(self) -> int:
147 target_width, target_height = self.get_image_size_with_most_features()
148
149 return self.get_num_image_tokens(
150 image_width=target_width, image_height=target_height
151 )
152
153
154 class SiglipDummyInputsBuilder(BaseDummyInputsBuilder[SiglipProcessingInfo]):
155 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
156 return ""
157
158 def get_dummy_mm_data(
159 self,
160 seq_len: int,
161 mm_counts: Mapping[str, int],
162 mm_options: MultiModalDummyOptions,
163 ) -> MultiModalDataDict:
164 target_width, target_height = self.info.get_image_size_with_most_features()
165
166 return {
167 "image": self._get_dummy_images(
168 width=target_width,
169 height=target_height,
170 num_images=mm_counts.get("image", 0),
171 overrides=mm_options.get("image"),
172 )
173 }
174
175
176 class SiglipMultiModalProcessor(BaseMultiModalProcessor[SiglipProcessingInfo]):
177 @cached_property
178 def image_token_id(self) -> int:
179 tokenizer = self.info.get_tokenizer()
180 dummy_token_id = next(
181 token_id
182 for token_id in range(tokenizer.vocab_size)
183 if token_id not in tokenizer.all_special_ids
184 )
185
186 return dummy_token_id
187
188 def apply(
189 self,
190 inputs: ProcessorInputs,
191 timing_ctx: TimingContext,
192 ) -> MultiModalInput:
193 if inputs.mm_data_items:
194 special_tokens = self.info.get_tokenizer().all_special_ids
195 if all(tok in special_tokens for tok in inputs.prompt):
196 inputs.prompt = []
197 else:
198 raise ValueError(
199 "SigLIP accepts text-only or image-only inputs, not both! "
200 "You must pass an image with an empty token prompt."
201 )
202
203 return super().apply(inputs, timing_ctx)
204
205 def _get_mm_fields_config(
206 self,
207 hf_inputs: BatchFeature,
208 hf_processor_mm_kwargs: Mapping[str, object],
209 ) -> Mapping[str, MultiModalFieldConfig]:
210 return dict(pixel_values=MultiModalFieldConfig.batched("image"))
211
212 def _get_prompt_updates(
213 self,
214 mm_items: MultiModalDataItems,
215 hf_processor_mm_kwargs: Mapping[str, object],
216 out_mm_kwargs: MultiModalKwargsItems,
217 ) -> list[PromptUpdate]:
218 image_token_id = self.image_token_id
219
220 def get_replacement(item_idx: int):
221 images = mm_items.get_items("image", ImageProcessorItems)
222 image_size = images.get_image_size(item_idx)
223
224 num_image_tokens = self.info.get_num_image_tokens(
225 image_width=image_size.width, image_height=image_size.height
226 )
227 return [image_token_id] * num_image_tokens
228
229 return [
230 PromptReplacement(
231 modality="image",
232 target=PromptIndexTargets.start(),
233 replacement=get_replacement,
234 ),
235 ]
236
237
238 class SiglipEncoderInfo(VisionEncoderInfo[SiglipVisionConfig]):
239 def get_num_image_tokens(
240 self,
241 *,
242 image_width: int,
243 image_height: int,
244 ) -> int:
245 return self.get_patch_grid_length() ** 2
246
247 def get_image_size(self) -> int:
248 return self.vision_config.image_size
249
250 def get_patch_size(self) -> int:
251 return self.vision_config.patch_size
252
253 def get_patch_grid_length(self) -> int:
254 image_size, patch_size = self.get_image_size(), self.get_patch_size()
255 return image_size // patch_size
256
257
258 # Adapted from https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/siglip/modeling_siglip.py#L216
259 class SiglipVisionEmbeddings(nn.Module):
260 def __init__(self, config: SiglipVisionConfig):
261 super().__init__()
262 self.config = config
263 self.embed_dim = config.hidden_size
264 self.image_size = config.image_size
265 self.patch_size = config.patch_size
266
267 self.patch_embedding = Conv2dLayer(
268 in_channels=config.num_channels,
269 out_channels=self.embed_dim,
270 kernel_size=self.patch_size,
271 stride=self.patch_size,
272 padding="valid",
273 )
274
275 self.num_patches = (self.image_size // self.patch_size) ** 2
276 self.num_positions = self.num_patches
277 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
278 self.register_buffer(
279 "position_ids",
280 torch.arange(self.num_positions, dtype=torch.int64).expand((1, -1)),
281 persistent=False,
282 )
283
284 def interpolate_pos_encoding(
285 self, embeddings: torch.Tensor, height: int, width: int
286 ) -> torch.Tensor:
287 num_patches = embeddings.shape[1]
288 num_positions = self.position_embedding.weight.shape[1]
289 if num_patches == num_positions and height == width:
290 return self.position_embedding(self.position_ids)
291
292 patch_pos_embed = self.position_embedding.weight.unsqueeze(0)
293
294 dim = embeddings.shape[-1]
295
296 new_height = height // self.patch_size
297 new_width = width // self.patch_size
298
299 sqrt_num_positions = int(num_positions**0.5)
300 patch_pos_embed = patch_pos_embed.reshape(
301 1, sqrt_num_positions, sqrt_num_positions, dim
302 )
303 patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
304
305 patch_pos_embed = nn.functional.interpolate(
306 patch_pos_embed,
307 size=(new_height, new_width),
308 mode="bicubic",
309 align_corners=False,
310 )
311
312 patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
313 return patch_pos_embed
314
315 def forward(
316 self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False
317 ) -> torch.Tensor:
318 _, _, height, width = pixel_values.shape
319 target_dtype = self.patch_embedding.weight.dtype
320 patch_embeds = self.patch_embedding(
321 pixel_values.to(dtype=target_dtype)
322 ) # shape = [*, width, grid, grid]
323 embeddings = patch_embeds.flatten(2).transpose(1, 2)
324
325 if interpolate_pos_encoding:
326 embeddings += self.interpolate_pos_encoding(embeddings, height, width)
327 else:
328 embeddings += self.position_embedding(self.position_ids)
329 return embeddings
330
331
332 class SiglipAttention(nn.Module):
333 def __init__(
334 self,
335 config: SiglipVisionConfig | SiglipTextConfig,
336 quant_config: QuantizationConfig | None = None,
337 *,
338 prefix: str = "",
339 attn_cls: type[EncoderOnlyAttention] | type[MMEncoderAttention],
340 ) -> None:
341 super().__init__()
342
343 self.config = config
344 self.embed_dim = config.hidden_size
345 self.num_heads = config.num_attention_heads
346 self.head_dim = self.embed_dim // self.num_heads
347 if self.head_dim * self.num_heads != self.embed_dim:
348 raise ValueError(
349 f"embed_dim must be divisible by num_heads "
350 f"(got `embed_dim`: {self.embed_dim} and "
351 f"`num_heads`: {self.num_heads})."
352 )
353
354 self.scale = self.head_dim**-0.5
355
356 use_data_parallel = is_vit_use_data_parallel()
357 self.qkv_proj = QKVParallelLinear(
358 hidden_size=self.embed_dim,
359 head_size=self.head_dim,
360 total_num_heads=self.num_heads,
361 quant_config=quant_config,
362 prefix=f"{prefix}.qkv_proj",
363 disable_tp=use_data_parallel,
364 )
365
366 self.out_proj = RowParallelLinear(
367 input_size=self.embed_dim,
368 output_size=self.embed_dim,
369 quant_config=quant_config,
370 prefix=f"{prefix}.out_proj",
371 disable_tp=use_data_parallel,
372 )
373
374 self.tp_size = (
375 1 if use_data_parallel else get_tensor_model_parallel_world_size()
376 )
377 self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
378
379 if attn_cls == MMEncoderAttention:
380 self.attn = attn_cls(
381 self.num_heads_per_partition,
382 self.head_dim,
383 self.scale,
384 prefix=f"{prefix}.attn",
385 )
386 else:
387 self.attn = attn_cls(
388 self.num_heads_per_partition,
389 self.head_dim,
390 self.scale,
391 prefix=f"{prefix}.attn",
392 )
393
394 def forward(
395 self,
396 hidden_states: torch.Tensor,
397 ) -> tuple[torch.Tensor, None]:
398 """Input shape: Batch x Time x Channel."""
399 qkv_states, _ = self.qkv_proj(hidden_states)
400 query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
401 out = self.attn(query_states, key_states, value_states)
402 attn_output, _ = self.out_proj(out)
403
404 return attn_output, None
405
406
407 class SiglipMLP(nn.Module):
408 def __init__(
409 self,
410 config: SiglipVisionConfig | SiglipTextConfig,
411 quant_config: QuantizationConfig | None = None,
412 prefix: str = "",
413 ) -> None:
414 super().__init__()
415
416 self.config = config
417 use_data_parallel = is_vit_use_data_parallel()
418 self.activation_fn = get_act_fn(config.hidden_act)
419
420 # Special handling for BNB and torchao quantization
421 if quant_config and quant_config.get_name() in ["bitsandbytes", "torchao"]:
422 quantizable = True
423 else:
424 # For other quantization, we require the hidden size to be a
425 # multiple of 64
426 quantizable = (
427 config.hidden_size % 64 == 0 and config.intermediate_size % 64 == 0
428 )
429
430 self.fc1 = ColumnParallelLinear(
431 config.hidden_size,
432 config.intermediate_size,
433 quant_config=quant_config if quantizable else None,
434 prefix=f"{prefix}.fc1",
435 disable_tp=use_data_parallel,
436 )
437 self.fc2 = RowParallelLinear(
438 config.intermediate_size,
439 config.hidden_size,
440 quant_config=quant_config if quantizable else None,
441 prefix=f"{prefix}.fc2",
442 disable_tp=use_data_parallel,
443 )
444
445 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
446 hidden_states, _ = self.fc1(hidden_states)
447 hidden_states = self.activation_fn(hidden_states)
448 hidden_states, _ = self.fc2(hidden_states)
449 return hidden_states
450
451
452 class SiglipEncoderLayer(nn.Module):
453 def __init__(
454 self,
455 config: SiglipVisionConfig | SiglipTextConfig,
456 quant_config: QuantizationConfig | None = None,
457 *,
458 prefix: str = "",
459 attn_cls: type[EncoderOnlyAttention] | type[MMEncoderAttention],
460 ) -> None:
461 super().__init__()
462
463 self.embed_dim = config.hidden_size
464
465 self.self_attn = SiglipAttention(
466 config,
467 quant_config=quant_config,
468 prefix=f"{prefix}.self_attn",
469 attn_cls=attn_cls,
470 )
471 self.layer_norm1 = LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
472 self.mlp = SiglipMLP(
473 config,
474 quant_config=quant_config,
475 prefix=f"{prefix}.mlp",
476 )
477 self.layer_norm2 = LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
478
479 def forward(
480 self,
481 hidden_states: torch.Tensor,
482 ) -> tuple[torch.Tensor, None]:
483 residual = hidden_states
484
485 hidden_states = self.layer_norm1(hidden_states)
486 hidden_states, _ = self.self_attn(hidden_states=hidden_states)
487 hidden_states += residual
488
489 residual = hidden_states
490 hidden_states = self.layer_norm2(hidden_states)
491 hidden_states = self.mlp(hidden_states)
492 hidden_states += residual
493
494 return hidden_states, None
495
496
497 class SiglipEncoder(nn.Module):
498 def __init__(
499 self,
500 config: SiglipVisionConfig | SiglipTextConfig,
501 quant_config: QuantizationConfig | None = None,
502 num_hidden_layers_override: int | None = None,
503 *,
504 prefix: str = "",
505 attn_cls: type[EncoderOnlyAttention] | type[MMEncoderAttention],
506 ) -> None:
507 super().__init__()
508
509 self.config = config
510
511 if num_hidden_layers_override is None:
512 num_hidden_layers = config.num_hidden_layers
513 else:
514 num_hidden_layers = num_hidden_layers_override
515
516 self.layers = nn.ModuleList(
517 [
518 SiglipEncoderLayer(
519 config,
520 quant_config=quant_config,
521 prefix=f"{prefix}.layers.{layer_idx}",
522 attn_cls=attn_cls,
523 )
524 for layer_idx in range(num_hidden_layers)
525 ]
526 )
527
528 def forward(
529 self,
530 inputs_embeds: torch.Tensor,
531 return_all_hidden_states: bool,
532 ) -> torch.Tensor | list[torch.Tensor]:
533 hidden_states_pool = [inputs_embeds]
534 hidden_states = inputs_embeds
535
536 for encoder_layer in self.layers:
537 hidden_states, _ = encoder_layer(hidden_states)
538 if return_all_hidden_states:
539 hidden_states_pool.append(hidden_states)
540 # If we have multiple feature sample layers, we return all hidden
541 # states in order and grab the ones we need by index.
542 if return_all_hidden_states:
543 return hidden_states_pool
544 return hidden_states
545
546
547 class SiglipTextTransformer(nn.Module):
548 hf_to_vllm_mapper = WeightsMapper(
549 orig_to_new_stacked={
550 ".q_proj": (".qkv_proj", "q"),
551 ".k_proj": (".qkv_proj", "k"),
552 ".v_proj": (".qkv_proj", "v"),
553 }
554 )
555
556 def __init__(
557 self,
558 config: SiglipTextConfig,
559 quant_config: QuantizationConfig | None = None,
560 *,
561 prefix: str = "",
562 ) -> None:
563 super().__init__()
564
565 self.config = config
566 embed_dim = config.hidden_size
567
568 self.embeddings = SiglipTextEmbeddings(config)
569
570 self.encoder = SiglipEncoder(
571 config=config,
572 quant_config=quant_config,
573 prefix=f"{prefix}.encoder",
574 attn_cls=EncoderOnlyAttention,
575 )
576
577 self.final_layer_norm = LayerNorm(embed_dim, eps=config.layer_norm_eps)
578 self.head = nn.Linear(embed_dim, config.projection_size)
579
580 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
581 return self.embeddings.token_embedding(input_ids)
582
583 def forward(
584 self,
585 input_ids: torch.Tensor | None,
586 position_ids: torch.Tensor,
587 inputs_embeds: torch.Tensor | None = None,
588 ) -> torch.Tensor:
589 hidden_states = self.embeddings(input_ids, position_ids, inputs_embeds)
590
591 last_hidden_state = self.encoder(
592 inputs_embeds=hidden_states, return_all_hidden_states=False
593 )
594
595 last_hidden_state = self.final_layer_norm(last_hidden_state)
596
597 return last_hidden_state
598
599 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
600 loader = AutoWeightsLoader(self)
601 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
602
603
604 class SiglipMultiheadAttentionPoolingHead(nn.Module):
605 """Multihead Attention Pooling."""
606
607 def __init__(
608 self,
609 config: SiglipVisionConfig,
610 quant_config: QuantizationConfig | None = None,
611 prefix: str = "",
612 ) -> None:
613 super().__init__()
614
615 self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
616 # TODO(ChristopherCho): Implement vLLM version of MultiheadAttention
617 self.attention = torch.nn.MultiheadAttention(
618 config.hidden_size, config.num_attention_heads, batch_first=True
619 )
620 self.layernorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
621 self.mlp = SiglipMLP(
622 config=config,
623 quant_config=quant_config,
624 prefix=f"{prefix}.mlp",
625 )
626
627 def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
628 batch_size = hidden_state.size(0)
629
630 probe = self.probe.expand(batch_size, -1, -1)
631
632 hidden_state = self.attention(probe, hidden_state, hidden_state)[0]
633
634 residual = hidden_state
635 hidden_state = self.layernorm(hidden_state)
636 hidden_state = self.mlp(hidden_state)
637 hidden_state += residual
638
639 # Handled by resolve_visual_encoder_outputs
640 # return hidden_state[:, 0]
641 return hidden_state
642
643
644 class SiglipVisionTransformer(nn.Module):
645 hf_to_vllm_mapper = WeightsMapper(
646 orig_to_new_stacked={
647 ".q_proj": (".qkv_proj", "q"),
648 ".k_proj": (".qkv_proj", "k"),
649 ".v_proj": (".qkv_proj", "v"),
650 }
651 )
652
653 def __init__(
654 self,
655 config: SiglipVisionConfig,
656 quant_config: QuantizationConfig | None = None,
657 *,
658 num_hidden_layers_override: int | None = None,
659 require_post_norm: bool | None = None,
660 prefix: str = "",
661 use_head: bool | None = False,
662 ) -> None:
663 super().__init__()
664
665 self.config = config
666 embed_dim = config.hidden_size
667
668 self.embeddings = SiglipVisionEmbeddings(config)
669
670 self.encoder = SiglipEncoder(
671 config,
672 quant_config=quant_config,
673 num_hidden_layers_override=num_hidden_layers_override,
674 prefix=f"{prefix}.encoder",
675 attn_cls=MMEncoderAttention,
676 )
677
678 num_hidden_layers = config.num_hidden_layers
679 if len(self.encoder.layers) > config.num_hidden_layers:
680 raise ValueError(
681 f"The original encoder only has {num_hidden_layers} "
682 f"layers, but you requested {len(self.encoder.layers)} layers."
683 )
684
685 # If possible, skip post_layernorm to conserve memory
686 if require_post_norm is None:
687 require_post_norm = len(self.encoder.layers) == num_hidden_layers
688
689 self.post_layernorm: LayerNorm | None
690 if require_post_norm:
691 self.post_layernorm = LayerNorm(embed_dim, eps=config.layer_norm_eps)
692 else:
693 self.post_layernorm = None
694
695 # Fall back to the config if a bool is not provided explicitly;
696 # note that many config types, including SiglipVisionConfig,
697 # do not have vision_use_head as a defined attribute.
698 if isinstance(use_head, bool):
699 self.use_head = use_head
700 else:
701 self.use_head = (
702 True
703 if not hasattr(config, "vision_use_head")
704 else config.vision_use_head
705 )
706
707 # Only create and load the head weights if we actually need them
708 self.head = (
709 SiglipMultiheadAttentionPoolingHead(
710 config=config,
711 quant_config=quant_config,
712 prefix=f"{prefix}.head",
713 )
714 if self.use_head
715 else None
716 )
717 self.last_hs_proc = partial(self.maybe_layer_norm_and_apply_head)
718
719 drops = dict[str, None]()
720 if self.post_layernorm is None:
721 drops["post_layernorm."] = None
722 if self.head is None:
723 drops["head."] = None
724 if drops:
725 self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(
726 orig_to_new_prefix=drops
727 )
728
729 @property
730 def dtype(self):
731 return next(self.parameters()).dtype
732
733 @property
734 def device(self):
735 return next(self.parameters()).device
736
737 def forward(
738 self,
739 pixel_values: torch.Tensor,
740 *,
741 interpolate_pos_encoding: bool = False,
742 select_layers: list[int] | None = None,
743 feature_select_strategy: VisionFeatureSelectStrategy | None = None,
744 ) -> torch.Tensor:
745 hidden_states = self.embeddings(
746 pixel_values,
747 interpolate_pos_encoding=interpolate_pos_encoding,
748 )
749 # Produces either the last layer output or all of the hidden states,
750 # depending on if we have select_layers or not
751 encoder_outputs = self.encoder(
752 inputs_embeds=hidden_states,
753 return_all_hidden_states=select_layers is not None,
754 )
755
756 # In the case that we have multiple feature layers,
757 # we stack and concatenate them into a tensor.
758 # NOTE: post layer norm and the attention pooling head
759 # are handled by last_hs_proc, which runs before applying
760 # the vision feature selection strategy.
761 encoder_outputs = resolve_visual_encoder_outputs(
762 encoder_outputs,
763 None,
764 select_layers=select_layers,
765 max_possible_layers=self.config.num_hidden_layers,
766 last_hs_proc=self.last_hs_proc,
767 feature_select_strategy=feature_select_strategy,
768 )
769
770 return encoder_outputs
771
772 def maybe_layer_norm_and_apply_head(
773 self, encoder_outputs: torch.Tensor
774 ) -> torch.Tensor:
775 """Apply the post layer norm and head if they are enabled,
776 given the last hidden states tensor.
777
778 Args:
779 encoder_outputs: The last hidden states from the visual encoder.
780
781 """
782 if self.post_layernorm is not None:
783 encoder_outputs = self.post_layernorm(encoder_outputs)
784 if self.head is not None:
785 encoder_outputs = self.head(encoder_outputs)
786 return encoder_outputs
787
788 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
789 loader = AutoWeightsLoader(self)
790
791 layer_count = len(self.encoder.layers)
792
793 def _filter(ws):
794 for name, w in ws:
795 # omit layers when num_hidden_layers_override is set
796 if name.startswith("encoder.layers.") and (
797 int(name.split(".")[2]) >= layer_count
798 ):
799 continue
800 yield name, w
801
802 return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
803
804
805 class SiglipVisionModel(nn.Module):
806 def __init__(
807 self,
808 config: SiglipVisionConfig,
809 quant_config: QuantizationConfig | None = None,
810 *,
811 num_hidden_layers_override: int | None = None,
812 require_post_norm: bool | None = None,
813 prefix: str = "",
814 use_head: bool | None = False,
815 ) -> None:
816 super().__init__()
817
818 self.quant_config = quant_config
819 self.vision_model = SiglipVisionTransformer(
820 config,
821 quant_config=quant_config,
822 num_hidden_layers_override=num_hidden_layers_override,
823 require_post_norm=require_post_norm,
824 prefix=maybe_prefix(prefix, "vision_model"),
825 use_head=use_head,
826 )
827
828 def get_input_embeddings(self) -> nn.Module:
829 return self.vision_model.embeddings.patch_embedding
830
831 @property
832 def dtype(self):
833 return self.vision_model.dtype
834
835 @property
836 def device(self):
837 return self.vision_model.device
838
839 def forward(
840 self,
841 pixel_values: torch.Tensor,
842 interpolate_pos_encoding: bool = False,
843 select_layers: list[int] | None = None,
844 feature_select_strategy: VisionFeatureSelectStrategy | None = None,
845 ) -> torch.Tensor:
846 return self.vision_model(
847 pixel_values=pixel_values,
848 interpolate_pos_encoding=interpolate_pos_encoding,
849 select_layers=select_layers,
850 feature_select_strategy=feature_select_strategy,
851 )
852
853
854 # Adapted from: https://github.com/huggingface/transformers/blob/v4.54.1/src/transformers/models/siglip/modeling_siglip.py#L200
855 class SiglipTextEmbeddings(nn.Module):
856 def __init__(self, config: SiglipTextConfig):
857 super().__init__()
858 self.config = config
859
860 self.token_embedding = VocabParallelEmbedding(
861 config.vocab_size, config.hidden_size
862 )
863
864 self.position_embedding = VocabParallelEmbedding(
865 config.max_position_embeddings, config.hidden_size
866 )
867
868 self.register_buffer(
869 "position_ids",
870 torch.arange(config.max_position_embeddings).expand((1, -1)),
871 persistent=False,
872 )
873
874 def forward(
875 self,
876 input_ids: torch.Tensor | None,
877 position_ids: torch.Tensor,
878 inputs_embeds: torch.Tensor | None = None,
879 ) -> torch.Tensor:
880 if inputs_embeds is None:
881 inputs_embeds = self.token_embedding(input_ids)
882
883 position_embeddings = self.position_embedding(position_ids)
884 embeddings = inputs_embeds + position_embeddings
885
886 return embeddings
887
888
889 # Assume EOS token corresponds to CLS token in text model
890 @default_pooling_type(seq_pooling_type="CLS")
891 @MULTIMODAL_REGISTRY.register_processor(
892 SiglipMultiModalProcessor,
893 info=SiglipProcessingInfo,
894 dummy_inputs=SiglipDummyInputsBuilder,
895 )
896 class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
897 is_pooling_model = True
898
899 hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={".position_ids": None})
900
901 packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
902
903 @classmethod
904 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
905 if modality.startswith("image"):
906 return None
907
908 raise ValueError("Only image modality is supported")
909
910 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
911 super().__init__()
912
913 config: SiglipConfig = vllm_config.model_config.hf_config
914 quant_config = vllm_config.quant_config
915 self.config = config
916
917 if hasattr(config, "num_labels"):
918 config.num_labels = 0
919
920 text_config = config.text_config
921 vision_config = config.vision_config
922
923 self.text_embed_dim = text_config.hidden_size
924 self.vision_embed_dim = vision_config.hidden_size
925 self.text_projection_size = text_config.projection_size
926
927 with self._mark_language_model(vllm_config):
928 self.text_model = SiglipTextTransformer(
929 text_config,
930 quant_config=quant_config,
931 prefix=maybe_prefix(prefix, "text_model"),
932 )
933
934 with self._mark_tower_model(vllm_config, "image"):
935 self.vision_model = SiglipVisionTransformer(
936 vision_config,
937 quant_config=quant_config,
938 prefix=maybe_prefix(prefix, "vision_model"),
939 use_head=None, # Allows potential pooling head
940 )
941
942 pooler_config = vllm_config.model_config.pooler_config
943 assert pooler_config is not None
944 self.pooler_config = pooler_config
945
946 self.pooler = DispatchPooler.for_embedding(pooler_config)
947
948 # Set in embed_input_ids; consumed by forward.
949 self._has_text_tokens = True
950 self._mm_token_mask: torch.Tensor | None = None
951
952 def get_text_features(
953 self,
954 input_ids: torch.Tensor | None,
955 position_ids: torch.Tensor,
956 inputs_embeds: torch.Tensor | None = None,
957 ) -> torch.Tensor:
958 last_hidden_state = self.text_model(
959 input_ids=input_ids,
960 position_ids=position_ids,
961 inputs_embeds=inputs_embeds,
962 )
963 text_features = self.text_model.head(last_hidden_state)
964
965 # SigLIP uses reversed position_ids;
966 # flip sequences to move EOS token to first position
967 text_features = self._flip_sequences_by_position_ids(
968 text_features, position_ids
969 )
970
971 return text_features
972
973 def _flip_sequences_by_position_ids(
974 self,
975 features: torch.Tensor,
976 position_ids: torch.Tensor,
977 ) -> torch.Tensor:
978 """Flip sequences so EOS token moves to first position for CLS pooling.
979
980 SigLIP position_ids are reversed within each sequence. This method detects
981 sequence boundaries and flips each sequence individually.
982 """
983 if len(features) == 1:
984 return features
985
986 # Detect sequence boundaries where position_ids decrease
987 position_diffs = position_ids[1:] - position_ids[:-1]
988 boundary_mask = position_diffs <= 0
989
990 with gpu_sync_allowed():
991 boundary_mid_cpu = torch.where(boundary_mask.cpu())[0] + 1
992
993 zero = torch.zeros(1, dtype=boundary_mid_cpu.dtype)
994 end = torch.full((1,), len(features), dtype=boundary_mid_cpu.dtype)
995 boundary_indices_cpu = torch.cat([zero, boundary_mid_cpu, end])
996
997 # For each sequence [start, end), position i flips to: start + end - 1 - i
998 lengths_cpu = boundary_indices_cpu[1:] - boundary_indices_cpu[:-1]
999 starts_cpu = boundary_indices_cpu[:-1]
1000 ends_cpu = boundary_indices_cpu[1:]
1001
1002 # Assign sequence ID to each element
1003 sequence_ids_cpu = torch.arange(
1004 len(lengths_cpu), dtype=boundary_mid_cpu.dtype
1005 ).repeat_interleave(lengths_cpu)
1006
1007 # Calculate flipped indices for all positions at once
1008 current_positions_cpu = torch.arange(
1009 len(features), dtype=boundary_mid_cpu.dtype
1010 )
1011 flip_indices_cpu = (
1012 starts_cpu[sequence_ids_cpu] + ends_cpu[sequence_ids_cpu]
1013 ) - (1 + current_positions_cpu)
1014 flip_indices = async_tensor_h2d(flip_indices_cpu, features.device)
1015
1016 return features[flip_indices]
1017
1018 def get_image_features(
1019 self,
1020 pixel_values: torch.Tensor,
1021 feature_select_strategy: VisionFeatureSelectStrategy | None = None,
1022 ) -> torch.Tensor:
1023 if feature_select_strategy is None:
1024 pooling_type = self.pooler_config.seq_pooling_type
1025 assert pooling_type is not None
1026 feature_select_strategy = _get_vision_feature_select_strategy(pooling_type)
1027
1028 pooled_output = self.vision_model(
1029 pixel_values=pixel_values,
1030 select_layers=None,
1031 feature_select_strategy=feature_select_strategy,
1032 )
1033
1034 return pooled_output
1035
1036 def _parse_and_validate_image_input(
1037 self, **kwargs: object
1038 ) -> SiglipImagePixelInputs | None:
1039 pixel_values = kwargs.pop("pixel_values", None)
1040 if pixel_values is None:
1041 return None
1042
1043 expected_h = expected_w = self.config.vision_config.image_size
1044 return SiglipImagePixelInputs(
1045 type="pixel_values",
1046 data=pixel_values,
1047 resolve_bindings={"h": expected_h, "w": expected_w},
1048 )
1049
1050 def _process_image_inputs(self, inputs: SiglipImagePixelInputs) -> torch.Tensor:
1051 pixel_values = inputs["data"]
1052
1053 return self.get_image_features(pixel_values)
1054
1055 def _embed_text_input_ids(
1056 self,
1057 input_ids: torch.Tensor,
1058 embed_input_ids: Callable[[torch.Tensor], torch.Tensor],
1059 *,
1060 is_multimodal: torch.Tensor | None,
1061 ) -> torch.Tensor:
1062 inputs_embeds = super()._embed_text_input_ids(
1063 input_ids,
1064 embed_input_ids,
1065 is_multimodal=is_multimodal,
1066 )
1067
1068 # NOTE: inputs_embeds in model runner has size text_config.projection_size
1069 # (instead of text_config.hidden_size) to accommodate image embeddings
1070 inputs_embeds_size = self.text_projection_size
1071 if inputs_embeds.shape[1] < inputs_embeds_size:
1072 inputs_embeds = torch.cat(
1073 [
1074 inputs_embeds,
1075 inputs_embeds.new_empty(
1076 inputs_embeds.shape[0],
1077 inputs_embeds_size - inputs_embeds.shape[1],
1078 ),
1079 ],
1080 dim=1,
1081 )
1082 elif inputs_embeds.shape[1] > inputs_embeds_size:
1083 # No need to handle this case for now
1084 raise NotImplementedError
1085
1086 return inputs_embeds
1087
1088 def embed_input_ids(
1089 self,
1090 input_ids: torch.Tensor,
1091 multimodal_embeddings: MultiModalEmbeddings | None = None,
1092 *,
1093 is_multimodal: torch.Tensor | None = None,
1094 ) -> torch.Tensor:
1095 has_mm_embeddings = (
1096 multimodal_embeddings is not None and len(multimodal_embeddings) > 0
1097 )
1098 self._mm_token_mask = is_multimodal
1099 self._has_text_tokens = dual_encoder_has_text_tokens(
1100 has_mm_embeddings, is_multimodal
1101 )
1102
1103 if multimodal_embeddings is None or is_multimodal is None:
1104 return super().embed_input_ids(input_ids)
1105
1106 return super().embed_input_ids(
1107 input_ids,
1108 multimodal_embeddings=multimodal_embeddings,
1109 is_multimodal=is_multimodal,
1110 )
1111
1112 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1113 image_input = self._parse_and_validate_image_input(**kwargs)
1114 if image_input is None:
1115 return []
1116
1117 vision_embeddings = self._process_image_inputs(image_input)
1118 return vision_embeddings
1119
1120 def forward(
1121 self,
1122 input_ids: torch.Tensor | None,
1123 positions: torch.Tensor,
1124 intermediate_tensors: IntermediateTensors | None = None,
1125 inputs_embeds: torch.Tensor | None = None,
1126 **kwargs: object,
1127 ) -> torch.Tensor:
1128 if intermediate_tensors is not None:
1129 raise RuntimeError("PP is not supported for this model")
1130
1131 # Multimodal inputs (image embeddings)
1132 assert inputs_embeds is not None
1133 if not self._has_text_tokens:
1134 return inputs_embeds
1135
1136 vision_embeds = inputs_embeds
1137 # NOTE: inputs_embeds in model runner has size text_config.projection_size
1138 # (instead of text_config.hidden_size) to accommodate image embeddings
1139 hidden_size = self.text_embed_dim
1140 if inputs_embeds.shape[1] > hidden_size:
1141 inputs_embeds = inputs_embeds[:, :hidden_size]
1142 elif inputs_embeds.shape[1] < hidden_size:
1143 # No need to handle this case for now
1144 raise NotImplementedError
1145
1146 text_features = self.get_text_features(input_ids, positions, inputs_embeds)
1147 return merge_dual_encoder_text_and_vision(
1148 text_features, vision_embeds, self._mm_token_mask
1149 )
1150
1151 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
1152 loader = AutoWeightsLoader(
1153 self,
1154 ignore_unexpected_prefixes=["logit_scale.", "logit_bias."],
1155 )
1156
1157 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)