| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import importlib.metadata |
| 5 | import math |
| 6 | from collections.abc import Iterable, Mapping, Sequence |
| 7 | from dataclasses import dataclass, fields |
| 8 | from typing import Annotated, Literal |
| 9 | |
| 10 | import numpy as np |
| 11 | import torch |
| 12 | import torch.nn as nn |
| 13 | import torch.nn.functional as F |
| 14 | from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk |
| 15 | from mistral_common.protocol.instruct.messages import UserMessage |
| 16 | from mistral_common.protocol.instruct.request import ChatCompletionRequest |
| 17 | from packaging.version import Version |
| 18 | from transformers import PixtralVisionConfig |
| 19 | from transformers.models.pixtral.image_processing_pixtral import ( |
| 20 | _num_image_tokens as _get_pixtral_hf_num_image_tokens, |
| 21 | ) |
| 22 | from transformers.models.pixtral.modeling_pixtral import apply_rotary_pos_emb |
| 23 | |
| 24 | from vllm.compilation.decorators import ( |
| 25 | should_torch_compile_mm_encoder, |
| 26 | support_torch_compile, |
| 27 | ) |
| 28 | from vllm.config import VllmConfig |
| 29 | from vllm.config.multimodal import MultiModalDummyOptions |
| 30 | from vllm.distributed import divide, get_tensor_model_parallel_world_size |
| 31 | from vllm.inputs import MultiModalDataDict |
| 32 | from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn |
| 33 | from vllm.model_executor.layers.attention import MMEncoderAttention |
| 34 | from vllm.model_executor.layers.conv import Conv2dLayer |
| 35 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 36 | from vllm.model_executor.layers.linear import ( |
| 37 | MergedColumnParallelLinear, |
| 38 | QKVParallelLinear, |
| 39 | ReplicatedLinear, |
| 40 | RowParallelLinear, |
| 41 | ) |
| 42 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 43 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 44 | from vllm.model_executor.models.utils import WeightsMapper |
| 45 | from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems |
| 46 | from vllm.multimodal.inputs import ( |
| 47 | MultiModalFieldConfig, |
| 48 | MultiModalKwargsItem, |
| 49 | NestedTensors, |
| 50 | ) |
| 51 | from vllm.multimodal.parse import ( |
| 52 | ImageProcessorItems, |
| 53 | ImageSize, |
| 54 | MultiModalDataItems, |
| 55 | ) |
| 56 | from vllm.multimodal.processing import BaseDummyInputsBuilder |
| 57 | from vllm.multimodal.processing.processor import ( |
| 58 | BaseMultiModalProcessor, |
| 59 | BaseProcessingInfo, |
| 60 | HFMultiModalInputs, |
| 61 | MultiModalProcessingResult, |
| 62 | PlaceholderFeaturesInfo, |
| 63 | ProcessorInputs, |
| 64 | PromptReplacement, |
| 65 | PromptUpdate, |
| 66 | PromptUpdateDetails, |
| 67 | ) |
| 68 | from vllm.sequence import IntermediateTensors |
| 69 | from vllm.tokenizers import cached_tokenizer_from_config |
| 70 | from vllm.tokenizers.mistral import MistralTokenizer |
| 71 | from vllm.transformers_utils.processors.pixtral import ( |
| 72 | MistralCommonImageProcessor, |
| 73 | MistralCommonPixtralProcessor, |
| 74 | ) |
| 75 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 76 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 77 | |
| 78 | from .interfaces import ( |
| 79 | MultiModalEmbeddings, |
| 80 | SupportsEagle3, |
| 81 | SupportsLoRA, |
| 82 | SupportsMultiModal, |
| 83 | SupportsPP, |
| 84 | supports_eagle3, |
| 85 | ) |
| 86 | from .module_mapping import MultiModelKeys |
| 87 | from .utils import StageMissingLayer, init_vllm_registered_model, maybe_prefix |
| 88 | from .vision import ( |
| 89 | VisionEncoderInfo, |
| 90 | VisionFeatureSelectStrategy, |
| 91 | is_vit_use_data_parallel, |
| 92 | resolve_visual_encoder_outputs, |
| 93 | ) |
| 94 | |
| 95 | PATCH_MERGE = "patch_merge" |
| 96 | |
| 97 | # Transformers 5.17 renamed Pixtral's rotary embedding and switched it to axial |
| 98 | # RoPE, which takes 2D (height, width) positions instead of flattened grid ids. |
| 99 | TRANSFORMERS_VERSION = importlib.metadata.version("transformers") |
| 100 | TRANSFORMERS_WITH_AXIAL_ROPE = Version(TRANSFORMERS_VERSION) >= Version("5.17.0.dev0") |
| 101 | |
| 102 | if TRANSFORMERS_WITH_AXIAL_ROPE: |
| 103 | from transformers.models.pixtral.modeling_pixtral import ( |
| 104 | PixtralVisionRotaryEmbedding, |
| 105 | ) |
| 106 | else: |
| 107 | from transformers.models.pixtral.modeling_pixtral import ( |
| 108 | PixtralRotaryEmbedding as PixtralVisionRotaryEmbedding, |
| 109 | ) |
| 110 | from transformers.models.pixtral.modeling_pixtral import ( |
| 111 | position_ids_in_meshgrid as flat_position_ids_in_meshgrid, |
| 112 | ) |
| 113 | |
| 114 | |
| 115 | def position_ids_in_meshgrid( |
| 116 | patch_embeds_list: list[torch.Tensor], |
| 117 | max_width: int, |
| 118 | ) -> torch.Tensor: |
| 119 | if not TRANSFORMERS_WITH_AXIAL_ROPE: |
| 120 | return flat_position_ids_in_meshgrid(patch_embeds_list, max_width) |
| 121 | positions = [] |
| 122 | for patch in patch_embeds_list: |
| 123 | height, width = patch.shape[-2:] |
| 124 | h_ids, w_ids = torch.meshgrid( |
| 125 | torch.arange(height), torch.arange(width), indexing="ij" |
| 126 | ) |
| 127 | positions.append(torch.stack([h_ids.flatten(), w_ids.flatten()], dim=-1)) |
| 128 | return torch.cat(positions, dim=0) |
| 129 | |
| 130 | |
| 131 | def _make_packed_sequence_metadata( |
| 132 | sequence_lengths: list[int], |
| 133 | attn_backend: AttentionBackendEnum, |
| 134 | hidden_size: int, |
| 135 | tp_size: int, |
| 136 | device: torch.device, |
| 137 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| 138 | lengths = np.array(sequence_lengths, dtype=np.int32) |
| 139 | cu_seqlens = np.concatenate( |
| 140 | [np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)] |
| 141 | ) |
| 142 | sequence_lengths_tensor = MMEncoderAttention.maybe_compute_seq_lens( |
| 143 | attn_backend, cu_seqlens, device |
| 144 | ) |
| 145 | max_seqlen = torch.tensor( |
| 146 | MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens), |
| 147 | dtype=torch.int32, |
| 148 | ) |
| 149 | cu_seqlens_tensor = MMEncoderAttention.maybe_recompute_cu_seqlens( |
| 150 | attn_backend, |
| 151 | cu_seqlens, |
| 152 | hidden_size, |
| 153 | tp_size, |
| 154 | device, |
| 155 | ) |
| 156 | return cu_seqlens_tensor, max_seqlen, sequence_lengths_tensor |
| 157 | |
| 158 | |
| 159 | def _is_layer_none_or_staged(layer: nn.Module) -> bool: |
| 160 | return layer is None or isinstance(layer, StageMissingLayer) |
| 161 | |
| 162 | |
| 163 | class PixtralImagePixelInputs(TensorSchema): |
| 164 | """Dimensions: |
| 165 | - bn: Batch size * number of images |
| 166 | - c: Number of channels (3) |
| 167 | - h: Height of each image |
| 168 | - w: Width of each image |
| 169 | |
| 170 | The result of stacking `ImageEncoding.tokens` from each prompt. |
| 171 | """ |
| 172 | |
| 173 | type: Literal["pixel_values"] = "pixel_values" |
| 174 | |
| 175 | images: Annotated[ |
| 176 | torch.Tensor | list[torch.Tensor], |
| 177 | TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}), |
| 178 | ] |
| 179 | |
| 180 | |
| 181 | class PixtralProcessingInfo(BaseProcessingInfo): |
| 182 | def get_tokenizer(self) -> MistralTokenizer: |
| 183 | tokenizer = cached_tokenizer_from_config(self.ctx.model_config) |
| 184 | if not isinstance(tokenizer, MistralTokenizer): |
| 185 | raise ValueError("This model requires `--tokenizer-mode mistral`") |
| 186 | |
| 187 | return tokenizer |
| 188 | |
| 189 | def get_image_processor(self) -> MistralCommonImageProcessor: |
| 190 | return MistralCommonImageProcessor(self.get_tokenizer().instruct.mm_encoder) |
| 191 | |
| 192 | def get_hf_processor(self, **kwargs) -> MistralCommonPixtralProcessor: |
| 193 | return MistralCommonPixtralProcessor( |
| 194 | tokenizer=self.get_tokenizer(), |
| 195 | image_processor=self.get_image_processor(), |
| 196 | ) |
| 197 | |
| 198 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 199 | return {"image": None} |
| 200 | |
| 201 | def get_image_size_with_most_features(self) -> ImageSize: |
| 202 | image_processor = self.get_image_processor() |
| 203 | max_image_size = image_processor.mm_encoder.mm_config.max_image_size |
| 204 | |
| 205 | return ImageSize(width=max_image_size, height=max_image_size) |
| 206 | |
| 207 | |
| 208 | class PixtralDummyInputsBuilder(BaseDummyInputsBuilder[PixtralProcessingInfo]): |
| 209 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 210 | return "" |
| 211 | |
| 212 | def get_dummy_mm_data( |
| 213 | self, |
| 214 | seq_len: int, |
| 215 | mm_counts: Mapping[str, int], |
| 216 | mm_options: MultiModalDummyOptions, |
| 217 | ) -> MultiModalDataDict: |
| 218 | target_width, target_height = self.info.get_image_size_with_most_features() |
| 219 | |
| 220 | return { |
| 221 | "image": self._get_dummy_images( |
| 222 | width=target_width, |
| 223 | height=target_height, |
| 224 | num_images=mm_counts.get("image", 0), |
| 225 | overrides=mm_options.get("image"), |
| 226 | ) |
| 227 | } |
| 228 | |
| 229 | |
| 230 | class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]): |
| 231 | def get_dummy_inputs( |
| 232 | self, |
| 233 | seq_len: int, |
| 234 | mm_counts: Mapping[str, int], |
| 235 | mm_options: MultiModalDummyOptions, |
| 236 | # For test_common.py only |
| 237 | mm_data: MultiModalDataDict | None = None, |
| 238 | ) -> ProcessorInputs: |
| 239 | builder = self.dummy_inputs |
| 240 | tokenizer = self.info.get_tokenizer() |
| 241 | |
| 242 | dummy_text = builder.get_dummy_text(mm_counts) |
| 243 | dummy_mm_data = ( |
| 244 | builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) |
| 245 | if mm_data is None |
| 246 | else mm_data |
| 247 | ) |
| 248 | dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) |
| 249 | dummy_images = ( |
| 250 | [] if "image" not in dummy_mm_data else dummy_mm_items["image"].get_all() |
| 251 | ) |
| 252 | |
| 253 | request = ChatCompletionRequest( |
| 254 | messages=[ |
| 255 | UserMessage( |
| 256 | content=[ |
| 257 | TextChunk(text=dummy_text), |
| 258 | *(ImageChunk(image=image) for image in dummy_images), |
| 259 | ] |
| 260 | ), |
| 261 | ] |
| 262 | ) |
| 263 | res = tokenizer.mistral.encode_chat_completion(request) |
| 264 | dummy_tokens = res.tokens |
| 265 | |
| 266 | return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items) |
| 267 | |
| 268 | # The tokens are already inserted by the chat template, |
| 269 | # so we just double check that they exist |
| 270 | def _maybe_apply_prompt_updates( |
| 271 | self, |
| 272 | mm_items: MultiModalDataItems, |
| 273 | mm_res: MultiModalProcessingResult, |
| 274 | ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: |
| 275 | mm_item_counts = mm_items.get_all_counts() |
| 276 | self._validate_mm_kwargs(mm_res.kwargs, mm_item_counts) |
| 277 | self._validate_mm_updates(mm_res.prompt_updates, mm_item_counts) |
| 278 | |
| 279 | mm_placeholders = self._find_mm_placeholders( |
| 280 | mm_res.prompt_ids, |
| 281 | mm_res.prompt_updates, |
| 282 | ) |
| 283 | self._validate_mm_placeholders(mm_placeholders, mm_item_counts) |
| 284 | |
| 285 | return mm_res.prompt_ids, mm_placeholders |
| 286 | |
| 287 | def _get_mm_fields_config( |
| 288 | self, |
| 289 | hf_inputs: Mapping[str, NestedTensors], |
| 290 | hf_processor_mm_kwargs: Mapping[str, object], |
| 291 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 292 | return dict(images=MultiModalFieldConfig.batched("image")) |
| 293 | |
| 294 | def _get_hf_mm_inputs( |
| 295 | self, |
| 296 | mm_items: MultiModalDataItems, |
| 297 | hf_kwargs: Mapping[str, object], |
| 298 | ) -> HFMultiModalInputs: |
| 299 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 300 | |
| 301 | # Avoid padding issue |
| 302 | return hf_inputs._replace( |
| 303 | hf_kwargs=dict(hf_inputs.hf_kwargs, return_tensors=None) |
| 304 | ) |
| 305 | |
| 306 | def _get_prompt_updates( |
| 307 | self, |
| 308 | mm_items: MultiModalDataItems, |
| 309 | hf_processor_mm_kwargs: Mapping[str, object], |
| 310 | out_mm_kwargs: MultiModalKwargsItems, |
| 311 | ) -> Sequence[PromptUpdate]: |
| 312 | processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 313 | |
| 314 | image_break_id = processor.image_break_id |
| 315 | image_token_id = processor.image_token_id |
| 316 | image_end_id = processor.image_end_id |
| 317 | |
| 318 | def get_replacement(item_idx: int): |
| 319 | images = mm_items.get_items("image", ImageProcessorItems) |
| 320 | image_size = images.get_image_size(item_idx) |
| 321 | |
| 322 | _, nrows, ncols = processor.image_processor.get_number_of_image_patches( |
| 323 | image_size.height, |
| 324 | image_size.width, |
| 325 | ) |
| 326 | |
| 327 | tokens = ([image_token_id] * ncols + [image_break_id]) * nrows |
| 328 | tokens[-1] = image_end_id |
| 329 | |
| 330 | return PromptUpdateDetails.select_token_id(tokens, image_token_id) |
| 331 | |
| 332 | return [ |
| 333 | PromptReplacement( |
| 334 | modality="image", |
| 335 | target=[], # Never match the prompt (see below note) |
| 336 | replacement=get_replacement, |
| 337 | ), |
| 338 | ] |
| 339 | |
| 340 | |
| 341 | @MULTIMODAL_REGISTRY.register_processor( |
| 342 | PixtralMultiModalProcessor, |
| 343 | info=PixtralProcessingInfo, |
| 344 | dummy_inputs=PixtralDummyInputsBuilder, |
| 345 | ) |
| 346 | class PixtralForConditionalGeneration( |
| 347 | nn.Module, SupportsLoRA, SupportsEagle3, SupportsMultiModal, SupportsPP |
| 348 | ): |
| 349 | hf_to_vllm_mapper = WeightsMapper( |
| 350 | orig_to_new_prefix={ |
| 351 | "model.language_model.": "language_model.model.", |
| 352 | "model.vision_tower.": "vision_encoder.", |
| 353 | "model.multi_modal_projector.": "vision_language_adapter.", |
| 354 | }, |
| 355 | orig_to_new_substr={ |
| 356 | ".linear_1.": ".w_in.", |
| 357 | ".linear_2.": ".w_out.", |
| 358 | }, |
| 359 | ) |
| 360 | |
| 361 | packed_modules_mapping = { |
| 362 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 363 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 364 | } |
| 365 | |
| 366 | supports_tower_connector_lora = True |
| 367 | |
| 368 | @classmethod |
| 369 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 370 | if modality.startswith("image"): |
| 371 | return None |
| 372 | |
| 373 | raise ValueError("Only image modality is supported") |
| 374 | |
| 375 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 376 | super().__init__() |
| 377 | config = vllm_config.model_config.hf_config |
| 378 | multimodal_config = vllm_config.model_config.multimodal_config |
| 379 | self.config = config |
| 380 | self.multimodal_config = multimodal_config |
| 381 | |
| 382 | dataclass_fields = {field.name for field in fields(VisionEncoderArgs)} |
| 383 | vision_args = { |
| 384 | key: value |
| 385 | for key, value in self.config.vision_config.to_dict().items() |
| 386 | if key in dataclass_fields |
| 387 | } |
| 388 | |
| 389 | self.vision_args = VisionEncoderArgs(**vision_args) |
| 390 | |
| 391 | # init MistralForCausalLM |
| 392 | with self._mark_language_model(vllm_config): |
| 393 | self.language_model = init_vllm_registered_model( |
| 394 | vllm_config=vllm_config, |
| 395 | hf_config=config.text_config, |
| 396 | prefix=maybe_prefix(prefix, "language_model"), |
| 397 | ) |
| 398 | |
| 399 | with self._mark_tower_model(vllm_config, "image"): |
| 400 | self.vision_encoder = VisionTransformer( |
| 401 | self.vision_args, |
| 402 | prefix=maybe_prefix(prefix, "vision_encoder"), |
| 403 | ) |
| 404 | self.pre_mm_projector_norm = ( |
| 405 | RMSNorm(self.vision_args.hidden_size, eps=1e-5) |
| 406 | if self.vision_args.add_pre_mm_projector_layer_norm |
| 407 | else None |
| 408 | ) |
| 409 | self.patch_merger = ( |
| 410 | PatchMerger( |
| 411 | vision_encoder_dim=self.vision_args.hidden_size, |
| 412 | spatial_merge_size=self.vision_args.spatial_merge_size, |
| 413 | use_mlp_bias=False, |
| 414 | ) |
| 415 | if self.vision_args.mm_projector_id == PATCH_MERGE |
| 416 | else None |
| 417 | ) |
| 418 | self.vision_language_adapter = VisionLanguageAdapter( |
| 419 | self.vision_args, dim=config.text_config.hidden_size |
| 420 | ) |
| 421 | |
| 422 | self.make_empty_intermediate_tensors = ( |
| 423 | self.language_model.make_empty_intermediate_tensors |
| 424 | ) |
| 425 | |
| 426 | def _parse_and_validate_image_input( |
| 427 | self, **kwargs: object |
| 428 | ) -> PixtralImagePixelInputs | None: |
| 429 | images = kwargs.pop("images", None) |
| 430 | if images is None: |
| 431 | return None |
| 432 | |
| 433 | return PixtralImagePixelInputs( |
| 434 | type="pixel_values", |
| 435 | images=images, |
| 436 | ) |
| 437 | |
| 438 | def _process_image_input( |
| 439 | self, |
| 440 | image_input: PixtralImagePixelInputs, |
| 441 | ) -> tuple[torch.Tensor, ...]: |
| 442 | images = image_input["images"] |
| 443 | image_features = self.vision_encoder(images) |
| 444 | feature_sizes = [image_feature.shape[0] for image_feature in image_features] |
| 445 | image_features = torch.cat(image_features) |
| 446 | if self.pre_mm_projector_norm is not None: |
| 447 | image_features = self.pre_mm_projector_norm(image_features) |
| 448 | if self.patch_merger is not None: |
| 449 | patch_size = self.vision_args.patch_size |
| 450 | spatial_merge_size_square = self.vision_args.spatial_merge_size**2 |
| 451 | img_patch_dims = [ |
| 452 | (img.shape[1] // patch_size, img.shape[2] // patch_size) |
| 453 | for img in images |
| 454 | ] |
| 455 | feature_sizes = [ |
| 456 | feature_size // spatial_merge_size_square |
| 457 | for feature_size in feature_sizes |
| 458 | ] |
| 459 | image_features = self.patch_merger( |
| 460 | image_features, image_sizes=img_patch_dims |
| 461 | ) |
| 462 | image_embeds = self.vision_language_adapter(image_features) |
| 463 | image_embeds = torch.split(image_embeds, feature_sizes) |
| 464 | return image_embeds |
| 465 | |
| 466 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 467 | image_input = self._parse_and_validate_image_input(**kwargs) |
| 468 | if image_input is None: |
| 469 | return [] |
| 470 | |
| 471 | return self._process_image_input(image_input) |
| 472 | |
| 473 | def forward( |
| 474 | self, |
| 475 | input_ids: torch.Tensor | None, |
| 476 | positions: torch.Tensor, |
| 477 | intermediate_tensors: IntermediateTensors | None = None, |
| 478 | inputs_embeds: torch.Tensor | None = None, |
| 479 | **kwargs: object, |
| 480 | ) -> torch.Tensor | IntermediateTensors: |
| 481 | """Run forward pass for pixtral.""" |
| 482 | if intermediate_tensors is not None: |
| 483 | inputs_embeds = None |
| 484 | |
| 485 | hidden_states = self.language_model.model( |
| 486 | input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds |
| 487 | ) |
| 488 | |
| 489 | return hidden_states |
| 490 | |
| 491 | def compute_logits( |
| 492 | self, |
| 493 | hidden_states: torch.Tensor, |
| 494 | ) -> torch.Tensor | None: |
| 495 | return self.language_model.compute_logits(hidden_states) |
| 496 | |
| 497 | def _require_language_model_eagle3(self) -> None: |
| 498 | if not supports_eagle3(self.language_model): |
| 499 | raise RuntimeError( |
| 500 | f"EAGLE-3 speculative decoding requires the language model to " |
| 501 | f"support EAGLE-3, but {type(self.language_model).__name__} does not." |
| 502 | ) |
| 503 | |
| 504 | def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: |
| 505 | self._require_language_model_eagle3() |
| 506 | self.language_model.set_aux_hidden_state_layers(layers) |
| 507 | |
| 508 | def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: |
| 509 | self._require_language_model_eagle3() |
| 510 | return self.language_model.get_eagle3_aux_hidden_state_layers() |
| 511 | |
| 512 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): |
| 513 | _vision_encoder_stacked_params = [ |
| 514 | # (param_name, shard_name, shard_id) |
| 515 | # HF format |
| 516 | (".qkv_proj", ".q_proj", "q"), |
| 517 | (".qkv_proj", ".k_proj", "k"), |
| 518 | (".qkv_proj", ".v_proj", "v"), |
| 519 | (".gate_up_proj", ".gate_proj", 0), |
| 520 | (".gate_up_proj", ".up_proj", 1), |
| 521 | # Mistral native (consolidated) format |
| 522 | (".qkv_proj", ".wq", "q"), |
| 523 | (".qkv_proj", ".wk", "k"), |
| 524 | (".qkv_proj", ".wv", "v"), |
| 525 | (".gate_up_proj", ".w1", 0), |
| 526 | (".gate_up_proj", ".w3", 1), |
| 527 | ] |
| 528 | |
| 529 | # Remap Mistral native names to HF-style names |
| 530 | # used by the vLLM vision encoder modules. |
| 531 | _vision_encoder_name_remap = { |
| 532 | ".wo.": ".o_proj.", |
| 533 | ".w2.": ".down_proj.", |
| 534 | } |
| 535 | |
| 536 | def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]): |
| 537 | return weight[0].startswith(("vision_encoder", "vision_tower")) |
| 538 | |
| 539 | def is_vision_lang_adapter_weights(weight: tuple[str, torch.Tensor]): |
| 540 | return weight[0].startswith( |
| 541 | ("vision_language_adapter", "multi_modal_projector") |
| 542 | ) |
| 543 | |
| 544 | def is_patch_merger(weight: tuple[str, torch.Tensor]): |
| 545 | return weight[0].startswith("patch_merger") |
| 546 | |
| 547 | def is_pre_mm_projector_norm(weight: tuple[str, torch.Tensor]): |
| 548 | return weight[0].startswith("pre_mm_projector_norm") |
| 549 | |
| 550 | vision_encoder_dict = ( |
| 551 | dict(self.vision_encoder.named_parameters()) |
| 552 | if self.vision_encoder is not None |
| 553 | else {} |
| 554 | ) |
| 555 | patch_merger_dict = ( |
| 556 | dict(self.patch_merger.named_parameters()) |
| 557 | if self.patch_merger is not None |
| 558 | else {} |
| 559 | ) |
| 560 | pre_mm_projector_norm_dict = ( |
| 561 | dict(self.pre_mm_projector_norm.named_parameters()) |
| 562 | if self.pre_mm_projector_norm is not None |
| 563 | else {} |
| 564 | ) |
| 565 | vision_lang_adapter_dict = ( |
| 566 | dict(self.vision_language_adapter.named_parameters()) |
| 567 | if self.vision_language_adapter is not None |
| 568 | else {} |
| 569 | ) |
| 570 | |
| 571 | def llm_weights_generator(): |
| 572 | for name, w in weights: |
| 573 | if is_vision_encoder_weights((name, w)): |
| 574 | if _is_layer_none_or_staged(self.vision_encoder): |
| 575 | continue |
| 576 | trimmed_name = ".".join(name.split(".")[1:]) |
| 577 | for ( |
| 578 | param_name, |
| 579 | weight_name, |
| 580 | shard_id, |
| 581 | ) in _vision_encoder_stacked_params: |
| 582 | if weight_name in trimmed_name: |
| 583 | trimmed_name = trimmed_name.replace(weight_name, param_name) |
| 584 | param = vision_encoder_dict[trimmed_name] |
| 585 | weight_loader = param.weight_loader |
| 586 | weight_loader(param, w, shard_id) |
| 587 | break |
| 588 | else: |
| 589 | for old, new in _vision_encoder_name_remap.items(): |
| 590 | if old in trimmed_name: |
| 591 | trimmed_name = trimmed_name.replace(old, new) |
| 592 | break |
| 593 | |
| 594 | param = vision_encoder_dict.get(trimmed_name) |
| 595 | if param is not None: |
| 596 | weight_loader = getattr( |
| 597 | param, "weight_loader", default_weight_loader |
| 598 | ) |
| 599 | weight_loader(param, w) |
| 600 | elif is_patch_merger((name, w)): |
| 601 | if _is_layer_none_or_staged(self.patch_merger): |
| 602 | continue |
| 603 | trimmed_name = ".".join(name.split(".")[1:]) |
| 604 | param = patch_merger_dict[trimmed_name] |
| 605 | weight_loader = getattr( |
| 606 | param, "weight_loader", default_weight_loader |
| 607 | ) |
| 608 | weight_loader(param, w) |
| 609 | elif is_pre_mm_projector_norm((name, w)): |
| 610 | if _is_layer_none_or_staged(self.pre_mm_projector_norm): |
| 611 | continue |
| 612 | trimmed_name = ".".join(name.split(".")[1:]) |
| 613 | param = pre_mm_projector_norm_dict[trimmed_name] |
| 614 | with torch.no_grad(): |
| 615 | default_weight_loader(param, w) |
| 616 | elif is_vision_lang_adapter_weights((name, w)): |
| 617 | if _is_layer_none_or_staged(self.vision_language_adapter): |
| 618 | continue |
| 619 | trimmed_name = ".".join(name.split(".")[1:]) |
| 620 | param = vision_lang_adapter_dict.get(trimmed_name) |
| 621 | if param is not None: |
| 622 | weight_loader = getattr( |
| 623 | param, "weight_loader", default_weight_loader |
| 624 | ) |
| 625 | weight_loader(param, w) |
| 626 | else: |
| 627 | name = name.removeprefix("language_model.") |
| 628 | yield (name, w) |
| 629 | |
| 630 | self.language_model.load_weights(llm_weights_generator()) |
| 631 | |
| 632 | def get_mm_mapping(self) -> MultiModelKeys: |
| 633 | return MultiModelKeys.from_string_field( |
| 634 | language_model="language_model.", |
| 635 | connector="vision_language_adapter.", |
| 636 | tower_model="vision_encoder", |
| 637 | ) |
| 638 | |
| 639 | def get_mm_lora_token_counts( |
| 640 | self, |
| 641 | *, |
| 642 | modality: str, |
| 643 | mm_kwargs: MultiModalKwargsItem | None, |
| 644 | num_mm_embeds: int, |
| 645 | ) -> tuple[int, int | None]: |
| 646 | del modality, mm_kwargs |
| 647 | if getattr(self, "patch_merger", None) is None: |
| 648 | return num_mm_embeds, num_mm_embeds |
| 649 | merge_size = self.vision_args.spatial_merge_size |
| 650 | return num_mm_embeds * (merge_size**2), num_mm_embeds |
| 651 | |
| 652 | |
| 653 | # Vision encoder |
| 654 | @dataclass |
| 655 | class VisionEncoderArgs: |
| 656 | hidden_size: int |
| 657 | num_channels: int |
| 658 | image_size: int |
| 659 | patch_size: int |
| 660 | intermediate_size: int |
| 661 | num_hidden_layers: int |
| 662 | num_attention_heads: int |
| 663 | rope_theta: float # for rope-2D |
| 664 | image_token_id: int |
| 665 | adapter_bias: bool = True |
| 666 | spatial_merge_size: int = 1 |
| 667 | add_pre_mm_projector_layer_norm: bool = False |
| 668 | mm_projector_id: str = "" |
| 669 | |
| 670 | |
| 671 | def _reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor: |
| 672 | """freqs_cis: complex - (seq_len, head_dim / 2) |
| 673 | x: complex - (bsz, seq_len, head_dim / 2) |
| 674 | """ |
| 675 | ndim = x.ndim |
| 676 | assert ndim > 1 |
| 677 | assert freqs_cis.shape == (x.shape[1], x.shape[-1]), ( |
| 678 | freqs_cis.shape, |
| 679 | (x.shape[1], x.shape[-1]), |
| 680 | ) |
| 681 | shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] |
| 682 | return freqs_cis.view(*shape) |
| 683 | |
| 684 | |
| 685 | def precompute_freqs_cis_2d( |
| 686 | dim: int, |
| 687 | height: int, |
| 688 | width: int, |
| 689 | theta: float, |
| 690 | ) -> torch.Tensor: |
| 691 | """freqs_cis: 2D complex tensor of shape (height, width, dim // 2) |
| 692 | to be indexed by (height, width) position tuples |
| 693 | """ |
| 694 | # (dim / 2) frequency bases |
| 695 | freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) |
| 696 | |
| 697 | h = torch.arange(height, device=freqs.device) |
| 698 | w = torch.arange(width, device=freqs.device) |
| 699 | |
| 700 | freqs_h = torch.outer(h, freqs[::2]).float() |
| 701 | freqs_w = torch.outer(w, freqs[1::2]).float() |
| 702 | freqs_2d = torch.cat( |
| 703 | [ |
| 704 | freqs_h[:, None, :].repeat(1, width, 1), |
| 705 | freqs_w[None, :, :].repeat(height, 1, 1), |
| 706 | ], |
| 707 | dim=-1, |
| 708 | ) |
| 709 | return torch.polar(torch.ones_like(freqs_2d), freqs_2d) |
| 710 | |
| 711 | |
| 712 | def apply_rotary_emb_vit( |
| 713 | xq: torch.Tensor, |
| 714 | xk: torch.Tensor, |
| 715 | freqs_cis: torch.Tensor, |
| 716 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 717 | xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) |
| 718 | xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) |
| 719 | assert freqs_cis.dtype == torch.complex64 |
| 720 | freqs_cis = _reshape_for_broadcast(freqs_cis, xq_) |
| 721 | xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) |
| 722 | xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) |
| 723 | return xq_out.type_as(xq), xk_out.type_as(xk) |
| 724 | |
| 725 | |
| 726 | class FeedForward(nn.Module): |
| 727 | def __init__( |
| 728 | self, |
| 729 | hidden_size: int, |
| 730 | intermediate_size: int, |
| 731 | quant_config: QuantizationConfig | None = None, |
| 732 | bias: bool = False, |
| 733 | prefix: str = "", |
| 734 | reduce_results: bool = True, |
| 735 | disable_tp: bool = False, |
| 736 | ) -> None: |
| 737 | super().__init__() |
| 738 | |
| 739 | self.gate_up_proj = MergedColumnParallelLinear( |
| 740 | input_size=hidden_size, |
| 741 | output_sizes=[intermediate_size] * 2, |
| 742 | bias=bias, |
| 743 | quant_config=quant_config, |
| 744 | disable_tp=disable_tp, |
| 745 | prefix=f"{prefix}.w13", |
| 746 | ) |
| 747 | self.down_proj = RowParallelLinear( |
| 748 | input_size=intermediate_size, |
| 749 | output_size=hidden_size, |
| 750 | bias=bias, |
| 751 | quant_config=quant_config, |
| 752 | reduce_results=reduce_results, |
| 753 | disable_tp=disable_tp, |
| 754 | prefix=f"{prefix}.w2", |
| 755 | ) |
| 756 | |
| 757 | self.act_fn = SiluAndMul() |
| 758 | |
| 759 | def forward(self, x): |
| 760 | x, _ = self.gate_up_proj(x) |
| 761 | x = self.act_fn(x) |
| 762 | x, _ = self.down_proj(x) |
| 763 | return x |
| 764 | |
| 765 | |
| 766 | class Attention(nn.Module): |
| 767 | def __init__( |
| 768 | self, |
| 769 | args: VisionEncoderArgs, |
| 770 | quant_config: QuantizationConfig | None = None, |
| 771 | prefix: str = "", |
| 772 | disable_tp: bool = False, |
| 773 | ): |
| 774 | super().__init__() |
| 775 | self.args = args |
| 776 | assert not args.hidden_size % args.num_attention_heads |
| 777 | self.head_dim = args.hidden_size // args.num_attention_heads |
| 778 | |
| 779 | self.qkv_proj = QKVParallelLinear( |
| 780 | hidden_size=args.hidden_size, |
| 781 | head_size=self.head_dim, |
| 782 | total_num_heads=args.num_attention_heads, |
| 783 | bias=False, |
| 784 | quant_config=quant_config, |
| 785 | prefix=f"{prefix}.wqkv", |
| 786 | disable_tp=disable_tp, |
| 787 | ) |
| 788 | self.o_proj = RowParallelLinear( |
| 789 | input_size=args.hidden_size, |
| 790 | output_size=args.hidden_size, |
| 791 | bias=False, |
| 792 | quant_config=quant_config, |
| 793 | prefix=f"{prefix}.wo", |
| 794 | disable_tp=disable_tp, |
| 795 | ) |
| 796 | |
| 797 | tp_size = 1 if disable_tp else get_tensor_model_parallel_world_size() |
| 798 | self.n_heads = divide(args.num_attention_heads, tp_size) |
| 799 | self.attn = MMEncoderAttention( |
| 800 | num_heads=self.n_heads, |
| 801 | head_size=self.head_dim, |
| 802 | prefix=f"{prefix}.attn", |
| 803 | ) |
| 804 | |
| 805 | def forward( |
| 806 | self, |
| 807 | x: torch.Tensor, |
| 808 | freqs_cis: torch.Tensor, |
| 809 | cu_seqlens: torch.Tensor, |
| 810 | max_seqlen: torch.Tensor, |
| 811 | sequence_lengths: torch.Tensor | None, |
| 812 | ) -> torch.Tensor: |
| 813 | batch, patches, _ = x.shape |
| 814 | |
| 815 | qkv, _ = self.qkv_proj(x) |
| 816 | q, k, v = qkv.chunk(3, dim=-1) |
| 817 | q = q.reshape(batch, patches, self.n_heads, self.head_dim) |
| 818 | k = k.reshape(batch, patches, self.n_heads, self.head_dim) |
| 819 | v = v.reshape(batch, patches, self.n_heads, self.head_dim) |
| 820 | |
| 821 | q, k = apply_rotary_emb_vit(q, k, freqs_cis=freqs_cis) |
| 822 | out = self.attn( |
| 823 | q, |
| 824 | k, |
| 825 | v, |
| 826 | cu_seqlens=cu_seqlens, |
| 827 | max_seqlen=max_seqlen, |
| 828 | sequence_lengths=sequence_lengths, |
| 829 | ) |
| 830 | |
| 831 | out = out.reshape(batch, patches, self.n_heads * self.head_dim) |
| 832 | out, _ = self.o_proj(out) |
| 833 | return out |
| 834 | |
| 835 | |
| 836 | @support_torch_compile( |
| 837 | dynamic_arg_dims={ |
| 838 | "x": 1, |
| 839 | "freqs_cis": 0, |
| 840 | "cu_seqlens": 0, |
| 841 | "sequence_lengths": 0, |
| 842 | }, |
| 843 | enable_if=should_torch_compile_mm_encoder, |
| 844 | is_encoder=True, |
| 845 | ) |
| 846 | class TransformerBlock(nn.Module): |
| 847 | def __init__( |
| 848 | self, |
| 849 | args: VisionEncoderArgs, |
| 850 | quant_config: QuantizationConfig | None = None, |
| 851 | prefix: str = "", |
| 852 | disable_tp: bool = False, |
| 853 | ): |
| 854 | super().__init__() |
| 855 | self.attention = Attention( |
| 856 | args, |
| 857 | quant_config=quant_config, |
| 858 | prefix=f"{prefix}.attention", |
| 859 | disable_tp=disable_tp, |
| 860 | ) |
| 861 | self.feed_forward = FeedForward( |
| 862 | args.hidden_size, |
| 863 | args.intermediate_size, |
| 864 | quant_config=quant_config, |
| 865 | prefix=f"{prefix}.feed_forward", |
| 866 | disable_tp=disable_tp, |
| 867 | ) |
| 868 | self.attention_norm = RMSNorm(args.hidden_size, eps=1e-5) |
| 869 | self.ffn_norm = RMSNorm(args.hidden_size, eps=1e-5) |
| 870 | |
| 871 | def forward( |
| 872 | self, |
| 873 | x: torch.Tensor, |
| 874 | freqs_cis: torch.Tensor, |
| 875 | cu_seqlens: torch.Tensor, |
| 876 | max_seqlen: torch.Tensor, |
| 877 | sequence_lengths: torch.Tensor | None, |
| 878 | ) -> torch.Tensor: |
| 879 | r = self.attention.forward( |
| 880 | self.attention_norm(x), |
| 881 | freqs_cis=freqs_cis, |
| 882 | cu_seqlens=cu_seqlens, |
| 883 | max_seqlen=max_seqlen, |
| 884 | sequence_lengths=sequence_lengths, |
| 885 | ) |
| 886 | h = x + r |
| 887 | r = self.feed_forward.forward(self.ffn_norm(h)) |
| 888 | out = h + r |
| 889 | return out |
| 890 | |
| 891 | |
| 892 | class Transformer(nn.Module): |
| 893 | def __init__( |
| 894 | self, |
| 895 | args: VisionEncoderArgs, |
| 896 | quant_config: QuantizationConfig | None = None, |
| 897 | prefix: str = "", |
| 898 | disable_tp: bool = False, |
| 899 | ): |
| 900 | super().__init__() |
| 901 | self.layers = torch.nn.ModuleList() |
| 902 | for idx in range(args.num_hidden_layers): |
| 903 | self.layers.append( |
| 904 | TransformerBlock( |
| 905 | args, |
| 906 | quant_config=quant_config, |
| 907 | prefix=f"{prefix}.layers.{idx}", |
| 908 | disable_tp=disable_tp, |
| 909 | ) |
| 910 | ) |
| 911 | |
| 912 | def forward( |
| 913 | self, |
| 914 | x: torch.Tensor, |
| 915 | freqs_cis: torch.Tensor | None, |
| 916 | cu_seqlens: torch.Tensor, |
| 917 | max_seqlen: torch.Tensor, |
| 918 | sequence_lengths: torch.Tensor | None, |
| 919 | ) -> torch.Tensor: |
| 920 | for layer in self.layers: |
| 921 | x = layer( |
| 922 | x, |
| 923 | freqs_cis=freqs_cis, |
| 924 | cu_seqlens=cu_seqlens, |
| 925 | max_seqlen=max_seqlen, |
| 926 | sequence_lengths=sequence_lengths, |
| 927 | ) |
| 928 | return x |
| 929 | |
| 930 | |
| 931 | def position_meshgrid( |
| 932 | patch_embeds_list: list[torch.Tensor], |
| 933 | ) -> torch.Tensor: |
| 934 | positions = torch.cat( |
| 935 | [ |
| 936 | torch.stack( |
| 937 | torch.meshgrid( |
| 938 | torch.arange(p.shape[-2]), |
| 939 | torch.arange(p.shape[-1]), |
| 940 | indexing="ij", |
| 941 | ), |
| 942 | dim=-1, |
| 943 | ).reshape(-1, 2) |
| 944 | for p in patch_embeds_list |
| 945 | ] |
| 946 | ) |
| 947 | return positions |
| 948 | |
| 949 | |
| 950 | class VisionTransformer(nn.Module): |
| 951 | def __init__( |
| 952 | self, |
| 953 | args: VisionEncoderArgs, |
| 954 | quant_config: QuantizationConfig | None = None, |
| 955 | prefix: str = "", |
| 956 | ): |
| 957 | super().__init__() |
| 958 | self.args = args |
| 959 | disable_tp = is_vit_use_data_parallel() |
| 960 | self.patch_conv = Conv2dLayer( |
| 961 | in_channels=args.num_channels, |
| 962 | out_channels=args.hidden_size, |
| 963 | kernel_size=args.patch_size, |
| 964 | stride=args.patch_size, |
| 965 | bias=False, |
| 966 | ) |
| 967 | self.ln_pre = RMSNorm(args.hidden_size, eps=1e-5) |
| 968 | self.transformer = Transformer( |
| 969 | args, |
| 970 | quant_config=quant_config, |
| 971 | prefix=f"{prefix}.transformer", |
| 972 | disable_tp=disable_tp, |
| 973 | ) |
| 974 | |
| 975 | head_dim = self.args.hidden_size // self.args.num_attention_heads |
| 976 | assert head_dim % 2 == 0, "ROPE requires even head_dim" |
| 977 | self._freqs_cis: torch.Tensor | None = None |
| 978 | |
| 979 | @property |
| 980 | def max_patches_per_side(self) -> int: |
| 981 | return self.args.image_size // self.args.patch_size |
| 982 | |
| 983 | @property |
| 984 | def device(self) -> torch.types.Device: |
| 985 | return next(self.parameters()).device |
| 986 | |
| 987 | @property |
| 988 | def dtype(self) -> torch.dtype: |
| 989 | return next(self.parameters()).dtype |
| 990 | |
| 991 | @property |
| 992 | def freqs_cis(self) -> torch.Tensor: |
| 993 | if self._freqs_cis is None: |
| 994 | self._freqs_cis = precompute_freqs_cis_2d( |
| 995 | dim=self.args.hidden_size // self.args.num_attention_heads, |
| 996 | height=self.max_patches_per_side, |
| 997 | width=self.max_patches_per_side, |
| 998 | theta=self.args.rope_theta, |
| 999 | ) |
| 1000 | |
| 1001 | if self._freqs_cis.device != self.device: |
| 1002 | self._freqs_cis = self._freqs_cis.to(device=self.device) |
| 1003 | |
| 1004 | return self._freqs_cis |
| 1005 | |
| 1006 | def forward( |
| 1007 | self, |
| 1008 | images: list[torch.Tensor], |
| 1009 | ) -> torch.Tensor: |
| 1010 | """Args: |
| 1011 | images: list of N_img images of variable sizes, |
| 1012 | each of shape (C, H, W) |
| 1013 | |
| 1014 | Returns: |
| 1015 | image_features: tensor of token features for |
| 1016 | all tokens of all images of shape (N_toks, D) |
| 1017 | |
| 1018 | """ |
| 1019 | # pass images through initial convolution independently |
| 1020 | patch_embeds_list = [ |
| 1021 | self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in images |
| 1022 | ] |
| 1023 | |
| 1024 | patch_embeds = [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list] |
| 1025 | embed_sizes = [p.shape[1] for p in patch_embeds] |
| 1026 | |
| 1027 | # flatten to a single sequence |
| 1028 | patch_embeds = torch.cat(patch_embeds, dim=1) |
| 1029 | patch_embeds = self.ln_pre(patch_embeds) |
| 1030 | |
| 1031 | # positional embeddings |
| 1032 | positions = position_meshgrid(patch_embeds_list).to(self.device) |
| 1033 | freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]] |
| 1034 | |
| 1035 | attention = self.transformer.layers[0].attention.attn |
| 1036 | cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata( |
| 1037 | embed_sizes, |
| 1038 | attention.attn_backend, |
| 1039 | self.args.hidden_size, |
| 1040 | 1 if is_vit_use_data_parallel() else get_tensor_model_parallel_world_size(), |
| 1041 | patch_embeds.device, |
| 1042 | ) |
| 1043 | out = self.transformer( |
| 1044 | patch_embeds, |
| 1045 | freqs_cis=freqs_cis, |
| 1046 | cu_seqlens=cu_seqlens, |
| 1047 | max_seqlen=max_seqlen, |
| 1048 | sequence_lengths=sequence_lengths, |
| 1049 | ) |
| 1050 | |
| 1051 | # squeeze dim 0 and split into separate tensors for each image |
| 1052 | return torch.split(out.squeeze(0), embed_sizes) |
| 1053 | |
| 1054 | |
| 1055 | class VisionLanguageAdapter(nn.Module): |
| 1056 | def __init__(self, args: VisionEncoderArgs, dim: int): |
| 1057 | super().__init__() |
| 1058 | assert isinstance(args, VisionEncoderArgs) |
| 1059 | self.w_in = ReplicatedLinear( |
| 1060 | args.hidden_size, |
| 1061 | dim, |
| 1062 | bias=args.adapter_bias, |
| 1063 | return_bias=False, |
| 1064 | ) |
| 1065 | self.gelu = nn.GELU() |
| 1066 | self.w_out = ReplicatedLinear( |
| 1067 | dim, dim, bias=args.adapter_bias, return_bias=False |
| 1068 | ) |
| 1069 | |
| 1070 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 1071 | return self.w_out(self.gelu(self.w_in(x))) |
| 1072 | |
| 1073 | |
| 1074 | class PatchMerger(nn.Module): |
| 1075 | """Learned merging of spatial_merge_size ** 2 patches.""" |
| 1076 | |
| 1077 | def __init__( |
| 1078 | self, |
| 1079 | vision_encoder_dim: int, |
| 1080 | spatial_merge_size: int, |
| 1081 | use_mlp_bias: bool = False, |
| 1082 | ) -> None: |
| 1083 | super().__init__() |
| 1084 | |
| 1085 | mlp_input_dim = vision_encoder_dim * (spatial_merge_size**2) |
| 1086 | |
| 1087 | self.spatial_merge_size = spatial_merge_size |
| 1088 | self.mlp_input_dim = mlp_input_dim |
| 1089 | |
| 1090 | self.merging_layer = ReplicatedLinear( |
| 1091 | mlp_input_dim, vision_encoder_dim, bias=use_mlp_bias, return_bias=False |
| 1092 | ) |
| 1093 | |
| 1094 | def forward( |
| 1095 | self, x: torch.Tensor, image_sizes: list[tuple[int, int]] |
| 1096 | ) -> torch.Tensor: |
| 1097 | # image_sizes specified in tokens |
| 1098 | assert sum([h * w for h, w in image_sizes]) == len(x) |
| 1099 | |
| 1100 | # x is (N, vision_encoder_dim) |
| 1101 | x = self.permute(x, image_sizes) |
| 1102 | |
| 1103 | # x is (N / spatial_merge_size ** 2, |
| 1104 | # vision_encoder_dim * spatial_merge_size ** 2) |
| 1105 | x = self.merging_layer(x) |
| 1106 | |
| 1107 | # x is (N / spatial_merge_size ** 2, vision_encoder_dim) |
| 1108 | return x |
| 1109 | |
| 1110 | def permute( |
| 1111 | self, |
| 1112 | x: torch.Tensor, |
| 1113 | image_sizes: list[tuple[int, int]], |
| 1114 | ) -> torch.Tensor: |
| 1115 | """Args: |
| 1116 | x: (N, D) where N is flattened and concatenated patch tokens |
| 1117 | for all images |
| 1118 | image_sizes: list of tuple of (height, width) in tokens for |
| 1119 | each image |
| 1120 | Returns: |
| 1121 | image_features: reorders patch tokens so each grid of |
| 1122 | (spatial_merge_size, spatial_merge_size) is contiguous. |
| 1123 | now (N / spatial_merge_size ** 2, D * spatial_merge_size ** 2) |
| 1124 | |
| 1125 | """ |
| 1126 | sub_grids = get_sub_grids( |
| 1127 | x=x, image_sizes=image_sizes, spatial_merge_size=self.spatial_merge_size |
| 1128 | ) # list of [d x sub_grid_size x sub_grid_size x n_patches] |
| 1129 | permuted_tensor: list[torch.Tensor] = [] |
| 1130 | for grid in sub_grids: |
| 1131 | n_patches = grid.shape[-1] |
| 1132 | permuted_tensor.append( |
| 1133 | grid.view(-1, n_patches).t() |
| 1134 | ) # n_patches x d * sub_grid_size * sub_grid_size |
| 1135 | return torch.cat( |
| 1136 | permuted_tensor, dim=0 |
| 1137 | ) # (N / spatial_merge_size ** 2, d * spatial_merge_size ** 2) |
| 1138 | |
| 1139 | |
| 1140 | def get_sub_grids( |
| 1141 | x: torch.Tensor, |
| 1142 | image_sizes: list[tuple[int, int]], |
| 1143 | spatial_merge_size: int, |
| 1144 | ) -> list[torch.Tensor]: |
| 1145 | # image_sizes specified in tokens |
| 1146 | tokens_per_image = [h * w for h, w in image_sizes] |
| 1147 | d = x.shape[-1] |
| 1148 | all_img_sub_grids: list[torch.Tensor] = [] |
| 1149 | sub_grid_size = spatial_merge_size |
| 1150 | |
| 1151 | for image_index, image_tokens in enumerate(x.split(tokens_per_image)): |
| 1152 | # Reshape image_tokens into a 2D grid |
| 1153 | h, w = image_sizes[image_index] |
| 1154 | image_grid = image_tokens.view(h, w, d).permute(2, 0, 1)[ |
| 1155 | None, :, :, : |
| 1156 | ] # 1 x d x h x w |
| 1157 | sub_grids = torch.nn.functional.unfold( |
| 1158 | image_grid, kernel_size=sub_grid_size, stride=sub_grid_size |
| 1159 | ) |
| 1160 | sub_grids = sub_grids.view( |
| 1161 | 1, d, sub_grid_size, sub_grid_size, -1 |
| 1162 | ) # 1 x d x sub_grid_size x sub_grid_size x n_patches |
| 1163 | |
| 1164 | all_img_sub_grids.append(sub_grids[0]) |
| 1165 | |
| 1166 | return all_img_sub_grids |
| 1167 | |
| 1168 | |
| 1169 | #### HF Transformers version of Pixtral #### |
| 1170 | # Based off https://github.com/huggingface/transformers/blob/d7950bff82b18c823193d17d72188c5e46d06c83/src/transformers/models/pixtral/modeling_pixtral.py |
| 1171 | # This model follows the Llava family, meaning image embeddings are placed |
| 1172 | # instead of the `[IMG]` token placeholders. |
| 1173 | # The model uses [`PixtralVisionModel`] for its vision encoder, |
| 1174 | # and [`MistralForCausalLM`] for its language decoder. |
| 1175 | |
| 1176 | |
| 1177 | class PixtralHFEncoderInfo(VisionEncoderInfo[PixtralVisionConfig]): |
| 1178 | def get_num_image_tokens( |
| 1179 | self, |
| 1180 | *, |
| 1181 | image_width: int, |
| 1182 | image_height: int, |
| 1183 | ) -> int: |
| 1184 | ncols, nrows = self.get_patch_grid_size( |
| 1185 | image_width=image_width, |
| 1186 | image_height=image_height, |
| 1187 | ) |
| 1188 | return ncols * nrows |
| 1189 | |
| 1190 | def get_image_size(self) -> int: |
| 1191 | return self.vision_config.image_size |
| 1192 | |
| 1193 | def get_patch_size(self) -> int: |
| 1194 | # spatial_merge_size is needed for Mistral3 |
| 1195 | spatial_merge_size = getattr(self.hf_config, "spatial_merge_size", 1) |
| 1196 | return self.vision_config.patch_size * spatial_merge_size |
| 1197 | |
| 1198 | def get_patch_grid_length(self) -> int: |
| 1199 | image_size, patch_size = self.get_image_size(), self.get_patch_size() |
| 1200 | |
| 1201 | # Since interpolation is applied, the image size need not be divisible |
| 1202 | # assert image_size % patch_size == 0 |
| 1203 | return image_size // patch_size |
| 1204 | |
| 1205 | # Adapted from: https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/pixtral/image_processing_pixtral.py#L99 |
| 1206 | def get_patch_grid_size( |
| 1207 | self, |
| 1208 | *, |
| 1209 | image_width: int, |
| 1210 | image_height: int, |
| 1211 | ) -> tuple[int, int]: |
| 1212 | max_width = max_height = self.get_image_size() |
| 1213 | patch_width = patch_height = self.get_patch_size() |
| 1214 | |
| 1215 | ratio = max(image_width / max_width, image_height / max_height) |
| 1216 | |
| 1217 | if ratio > 1: |
| 1218 | image_width = int(math.floor(image_width / ratio)) |
| 1219 | image_height = int(math.floor(image_height / ratio)) |
| 1220 | |
| 1221 | nrows, ncols = _get_pixtral_hf_num_image_tokens( |
| 1222 | (image_height, image_width), |
| 1223 | (patch_height, patch_width), |
| 1224 | ) # type: ignore |
| 1225 | |
| 1226 | return ncols, nrows |
| 1227 | |
| 1228 | |
| 1229 | class PixtralHFMLP(nn.Module): |
| 1230 | def __init__( |
| 1231 | self, |
| 1232 | config: PixtralVisionConfig, |
| 1233 | quant_config: QuantizationConfig | None = None, |
| 1234 | *, |
| 1235 | prefix: str = "", |
| 1236 | ) -> None: |
| 1237 | super().__init__() |
| 1238 | |
| 1239 | use_data_parallel = is_vit_use_data_parallel() |
| 1240 | |
| 1241 | assert config.intermediate_size is not None |
| 1242 | self.gate_up_proj = MergedColumnParallelLinear( |
| 1243 | input_size=config.hidden_size, |
| 1244 | output_sizes=[config.intermediate_size] * 2, |
| 1245 | bias=False, |
| 1246 | quant_config=quant_config, |
| 1247 | prefix=f"{prefix}.gate_up_proj", |
| 1248 | disable_tp=use_data_parallel, |
| 1249 | ) |
| 1250 | self.down_proj = RowParallelLinear( |
| 1251 | input_size=config.intermediate_size, |
| 1252 | output_size=config.hidden_size, |
| 1253 | bias=False, |
| 1254 | quant_config=quant_config, |
| 1255 | prefix=f"{prefix}.down_proj", |
| 1256 | disable_tp=use_data_parallel, |
| 1257 | ) |
| 1258 | self.act_and_mul = get_act_and_mul_fn(config.hidden_act) |
| 1259 | |
| 1260 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 1261 | gate_up, _ = self.gate_up_proj(x) |
| 1262 | x = self.act_and_mul(gate_up) |
| 1263 | x, _ = self.down_proj(x) |
| 1264 | return x |
| 1265 | |
| 1266 | |
| 1267 | class PixtralHFAttention(nn.Module): |
| 1268 | def __init__( |
| 1269 | self, |
| 1270 | config: PixtralVisionConfig, |
| 1271 | quant_config: QuantizationConfig | None = None, |
| 1272 | *, |
| 1273 | prefix: str = "", |
| 1274 | ) -> None: |
| 1275 | super().__init__() |
| 1276 | |
| 1277 | self.config = config |
| 1278 | assert not config.hidden_size % config.num_attention_heads |
| 1279 | self.total_num_heads = config.num_attention_heads |
| 1280 | self.head_dim = config.hidden_size // config.num_attention_heads |
| 1281 | assert self.total_num_heads * self.head_dim == config.hidden_size |
| 1282 | |
| 1283 | use_data_parallel = is_vit_use_data_parallel() |
| 1284 | self.qkv_proj = QKVParallelLinear( |
| 1285 | hidden_size=config.hidden_size, |
| 1286 | head_size=self.head_dim, |
| 1287 | total_num_heads=self.total_num_heads, |
| 1288 | bias=False, |
| 1289 | quant_config=quant_config, |
| 1290 | prefix=f"{prefix}.qkv_proj", |
| 1291 | disable_tp=use_data_parallel, |
| 1292 | ) |
| 1293 | self.o_proj = RowParallelLinear( |
| 1294 | input_size=config.hidden_size, |
| 1295 | output_size=config.hidden_size, |
| 1296 | bias=False, |
| 1297 | quant_config=quant_config, |
| 1298 | prefix=f"{prefix}.o_proj", |
| 1299 | disable_tp=use_data_parallel, |
| 1300 | ) |
| 1301 | |
| 1302 | self.tp_size = ( |
| 1303 | 1 if use_data_parallel else get_tensor_model_parallel_world_size() |
| 1304 | ) |
| 1305 | self.n_heads = divide(config.num_attention_heads, self.tp_size) |
| 1306 | self.attn = MMEncoderAttention( |
| 1307 | num_heads=self.n_heads, |
| 1308 | head_size=self.head_dim, |
| 1309 | prefix=f"{prefix}.attn", |
| 1310 | ) |
| 1311 | |
| 1312 | def forward( |
| 1313 | self, |
| 1314 | hidden_states: torch.Tensor, |
| 1315 | position_embeddings: tuple[torch.Tensor, torch.Tensor], |
| 1316 | cu_seqlens: torch.Tensor, |
| 1317 | max_seqlen: torch.Tensor, |
| 1318 | sequence_lengths: torch.Tensor | None, |
| 1319 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 1320 | batch, patches, _ = hidden_states.size() |
| 1321 | |
| 1322 | qkv_states, _ = self.qkv_proj(hidden_states) |
| 1323 | q, k, v = qkv_states.chunk(3, dim=-1) |
| 1324 | |
| 1325 | # Transpose q and k to apply HF's Rotary Position Embedding. |
| 1326 | q = q.view(batch, patches, self.n_heads, self.head_dim).transpose(1, 2) |
| 1327 | k = k.view(batch, patches, self.n_heads, self.head_dim).transpose(1, 2) |
| 1328 | v = v.view(batch, patches, self.n_heads, self.head_dim) |
| 1329 | cos, sin = position_embeddings |
| 1330 | q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=0) |
| 1331 | out = self.attn( |
| 1332 | q.transpose(1, 2).contiguous(), |
| 1333 | k.transpose(1, 2).contiguous(), |
| 1334 | v, |
| 1335 | cu_seqlens=cu_seqlens, |
| 1336 | max_seqlen=max_seqlen, |
| 1337 | sequence_lengths=sequence_lengths, |
| 1338 | ) |
| 1339 | |
| 1340 | out = out.reshape(batch, patches, self.n_heads * self.head_dim) |
| 1341 | attn_output, _ = self.o_proj(out) |
| 1342 | |
| 1343 | return attn_output, None |
| 1344 | |
| 1345 | |
| 1346 | @support_torch_compile( |
| 1347 | dynamic_arg_dims={ |
| 1348 | "hidden_states": 1, |
| 1349 | "rotary_pos_emb_cos": 0, |
| 1350 | "rotary_pos_emb_sin": 0, |
| 1351 | "cu_seqlens": 0, |
| 1352 | "sequence_lengths": 0, |
| 1353 | }, |
| 1354 | enable_if=should_torch_compile_mm_encoder, |
| 1355 | is_encoder=True, |
| 1356 | ) |
| 1357 | class PixtralHFTransformerBlock(nn.Module): |
| 1358 | def __init__( |
| 1359 | self, |
| 1360 | config: PixtralVisionConfig, |
| 1361 | quant_config: QuantizationConfig | None = None, |
| 1362 | *, |
| 1363 | prefix: str = "", |
| 1364 | ) -> None: |
| 1365 | super().__init__() |
| 1366 | |
| 1367 | self.attention_norm = RMSNorm(config.hidden_size, eps=1e-5) |
| 1368 | self.attention = PixtralHFAttention( |
| 1369 | config, |
| 1370 | quant_config=quant_config, |
| 1371 | prefix=f"{prefix}.attention", |
| 1372 | ) |
| 1373 | self.feed_forward = PixtralHFMLP( |
| 1374 | config, |
| 1375 | quant_config=quant_config, |
| 1376 | prefix=f"{prefix}.feed_forward", |
| 1377 | ) |
| 1378 | self.ffn_norm = RMSNorm(config.hidden_size, eps=1e-5) |
| 1379 | |
| 1380 | def forward( |
| 1381 | self, |
| 1382 | hidden_states: torch.Tensor, |
| 1383 | rotary_pos_emb_cos: torch.Tensor, |
| 1384 | rotary_pos_emb_sin: torch.Tensor, |
| 1385 | cu_seqlens: torch.Tensor, |
| 1386 | max_seqlen: torch.Tensor, |
| 1387 | sequence_lengths: torch.Tensor | None, |
| 1388 | ) -> torch.Tensor: |
| 1389 | r, _ = self.attention.forward( |
| 1390 | self.attention_norm(hidden_states), |
| 1391 | position_embeddings=(rotary_pos_emb_cos, rotary_pos_emb_sin), |
| 1392 | cu_seqlens=cu_seqlens, |
| 1393 | max_seqlen=max_seqlen, |
| 1394 | sequence_lengths=sequence_lengths, |
| 1395 | ) |
| 1396 | h = hidden_states + r |
| 1397 | r = self.feed_forward.forward(self.ffn_norm(h)) |
| 1398 | out = h + r |
| 1399 | return out |
| 1400 | |
| 1401 | |
| 1402 | class PixtralHFTransformer(nn.Module): |
| 1403 | def __init__( |
| 1404 | self, |
| 1405 | config: PixtralVisionConfig, |
| 1406 | quant_config: QuantizationConfig | None = None, |
| 1407 | *, |
| 1408 | num_hidden_layers_override: int | None = None, |
| 1409 | prefix: str = "", |
| 1410 | ) -> None: |
| 1411 | super().__init__() |
| 1412 | |
| 1413 | if num_hidden_layers_override is None: |
| 1414 | num_hidden_layers = config.num_hidden_layers |
| 1415 | else: |
| 1416 | num_hidden_layers = num_hidden_layers_override |
| 1417 | |
| 1418 | self.layers = nn.ModuleList( |
| 1419 | [ |
| 1420 | PixtralHFTransformerBlock( |
| 1421 | config=config, |
| 1422 | quant_config=quant_config, |
| 1423 | prefix=f"{prefix}.layers.{layer_idx}", |
| 1424 | ) |
| 1425 | for layer_idx in range(num_hidden_layers) |
| 1426 | ] |
| 1427 | ) |
| 1428 | |
| 1429 | def forward( |
| 1430 | self, |
| 1431 | x: torch.Tensor, |
| 1432 | position_embeddings: tuple[torch.Tensor, torch.Tensor], |
| 1433 | cu_seqlens: torch.Tensor, |
| 1434 | max_seqlen: torch.Tensor, |
| 1435 | sequence_lengths: torch.Tensor | None, |
| 1436 | return_all_hidden_states: bool, |
| 1437 | ) -> torch.Tensor: |
| 1438 | hidden_states_pool = [x] |
| 1439 | rotary_pos_emb_cos, rotary_pos_emb_sin = position_embeddings |
| 1440 | |
| 1441 | for layer in self.layers: |
| 1442 | x = layer( |
| 1443 | x, |
| 1444 | rotary_pos_emb_cos, |
| 1445 | rotary_pos_emb_sin, |
| 1446 | cu_seqlens, |
| 1447 | max_seqlen, |
| 1448 | sequence_lengths, |
| 1449 | ) |
| 1450 | if return_all_hidden_states: |
| 1451 | hidden_states_pool.append(x) |
| 1452 | # If we have multiple feature sample layers, we return all hidden |
| 1453 | # states in order and grab the ones we need by index. |
| 1454 | if return_all_hidden_states: |
| 1455 | return hidden_states_pool |
| 1456 | return x |
| 1457 | |
| 1458 | |
| 1459 | def pixtral_patch_embed( |
| 1460 | images: list[torch.Tensor], |
| 1461 | weight: torch.Tensor, |
| 1462 | input_norm: nn.Module, |
| 1463 | ) -> tuple[torch.Tensor, list[torch.Tensor]]: |
| 1464 | """Normalize and project all CHW images in one matrix multiplication.""" |
| 1465 | out_channels, in_channels, patch_height, patch_width = weight.shape |
| 1466 | grid_sizes = [] |
| 1467 | for image in images: |
| 1468 | assert image.ndim == 3 and image.shape[0] == in_channels |
| 1469 | _, height, width = image.shape |
| 1470 | patch_rows = height // patch_height |
| 1471 | patch_cols = width // patch_width |
| 1472 | grid_sizes.append((patch_rows, patch_cols)) |
| 1473 | |
| 1474 | patch_counts = [rows * cols for rows, cols in grid_sizes] |
| 1475 | packed = torch.empty( |
| 1476 | (sum(patch_counts), in_channels * patch_height * patch_width), |
| 1477 | dtype=images[0].dtype, |
| 1478 | device=images[0].device, |
| 1479 | ) |
| 1480 | offset = 0 |
| 1481 | for image, (rows, cols), count in zip(images, grid_sizes, patch_counts): |
| 1482 | source = ( |
| 1483 | image[:, : rows * patch_height, : cols * patch_width] |
| 1484 | .reshape(in_channels, rows, patch_height, cols, patch_width) |
| 1485 | .permute(1, 3, 0, 2, 4) |
| 1486 | ) |
| 1487 | packed[offset : offset + count].view( |
| 1488 | rows, cols, in_channels, patch_height, patch_width |
| 1489 | ).copy_(source) |
| 1490 | offset += count |
| 1491 | |
| 1492 | normalized = input_norm(packed, weight.dtype) |
| 1493 | projected = F.linear(normalized, weight.reshape(out_channels, -1)) |
| 1494 | embeddings = [ |
| 1495 | tokens.transpose(0, 1).reshape(1, out_channels, rows, cols) |
| 1496 | for tokens, (rows, cols) in zip(projected.split(patch_counts), grid_sizes) |
| 1497 | ] |
| 1498 | return projected.unsqueeze(0), embeddings |
| 1499 | |
| 1500 | |
| 1501 | class PixtralHFVisionModel(nn.Module): |
| 1502 | def __init__( |
| 1503 | self, |
| 1504 | config: PixtralVisionConfig, |
| 1505 | quant_config: QuantizationConfig | None = None, |
| 1506 | *, |
| 1507 | input_norm: nn.Module | None = None, |
| 1508 | num_hidden_layers_override: int | None = None, |
| 1509 | require_post_norm: bool | None = None, |
| 1510 | prefix: str = "", |
| 1511 | ) -> None: |
| 1512 | super().__init__() |
| 1513 | |
| 1514 | self.config = config |
| 1515 | self.input_norm = input_norm |
| 1516 | |
| 1517 | self.patch_conv = Conv2dLayer( |
| 1518 | in_channels=config.num_channels, |
| 1519 | out_channels=config.hidden_size, |
| 1520 | kernel_size=config.patch_size, |
| 1521 | stride=config.patch_size, |
| 1522 | bias=False, |
| 1523 | ) |
| 1524 | self.ln_pre = RMSNorm(config.hidden_size, eps=1e-5) |
| 1525 | self.transformer = PixtralHFTransformer( |
| 1526 | config, |
| 1527 | quant_config=quant_config, |
| 1528 | num_hidden_layers_override=num_hidden_layers_override, |
| 1529 | prefix=f"{prefix}.transformer", |
| 1530 | ) |
| 1531 | |
| 1532 | num_hidden_layers = config.num_hidden_layers |
| 1533 | if len(self.transformer.layers) > config.num_hidden_layers: |
| 1534 | raise ValueError( |
| 1535 | f"The original encoder only has {num_hidden_layers} " |
| 1536 | f"layers, but you requested {len(self.transformer.layers)} " |
| 1537 | "layers." |
| 1538 | ) |
| 1539 | |
| 1540 | if require_post_norm is True: |
| 1541 | msg = "PixtralHFVisionModel does not have post-layernorm" |
| 1542 | raise ValueError(msg) |
| 1543 | |
| 1544 | self.dtype = next(self.parameters()).dtype |
| 1545 | self.device = next(self.parameters()).device |
| 1546 | self.patch_positional_embedding = PixtralVisionRotaryEmbedding(config).to( |
| 1547 | self.device |
| 1548 | ) |
| 1549 | |
| 1550 | def forward( |
| 1551 | self, |
| 1552 | pixel_values: list[torch.Tensor], |
| 1553 | *, |
| 1554 | select_layers: list[int] | None = None, |
| 1555 | feature_select_strategy: VisionFeatureSelectStrategy | None = None, |
| 1556 | ) -> tuple[torch.Tensor, ...]: |
| 1557 | """Args: |
| 1558 | pixel_values: Each image to be processed will be a separate tensor |
| 1559 | in pixel_values. This means it will be a list of tensors |
| 1560 | because multiple requests batched can have multiple images, |
| 1561 | each with their own shape potentially |
| 1562 | select_layers: Layer indices whose features should be |
| 1563 | concatenated and used as the visual encoder output. If none |
| 1564 | are provided, the last layer is used. |
| 1565 | |
| 1566 | Returns: |
| 1567 | image_features: tensor of token features for |
| 1568 | all tokens of all images of shape (N_toks, D) |
| 1569 | |
| 1570 | """ |
| 1571 | images = [image.to(device=self.device) for image in pixel_values] |
| 1572 | if self.input_norm is not None: |
| 1573 | patch_embeds, patch_embeds_list = pixtral_patch_embed( |
| 1574 | images, |
| 1575 | self.patch_conv.weight, |
| 1576 | self.input_norm, |
| 1577 | ) |
| 1578 | else: |
| 1579 | patch_embeds_list = [] |
| 1580 | for image in images: |
| 1581 | patch_embeds_list.append( |
| 1582 | self.patch_conv(image.unsqueeze(0).to(self.dtype)) |
| 1583 | ) |
| 1584 | patch_embeds = torch.cat( |
| 1585 | [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list], dim=1 |
| 1586 | ) |
| 1587 | |
| 1588 | embed_sizes = [p.shape[-2] * p.shape[-1] for p in patch_embeds_list] |
| 1589 | patch_embeds = self.ln_pre(patch_embeds) |
| 1590 | |
| 1591 | # positional embeddings |
| 1592 | position_ids = position_ids_in_meshgrid( |
| 1593 | patch_embeds_list, |
| 1594 | max_width=self.config.image_size // self.config.patch_size, |
| 1595 | ).to(self.device) |
| 1596 | position_embedding = self.patch_positional_embedding(patch_embeds, position_ids) |
| 1597 | |
| 1598 | attention = self.transformer.layers[0].attention.attn |
| 1599 | cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata( |
| 1600 | embed_sizes, |
| 1601 | attention.attn_backend, |
| 1602 | self.config.hidden_size, |
| 1603 | self.transformer.layers[0].attention.tp_size, |
| 1604 | patch_embeds.device, |
| 1605 | ) |
| 1606 | |
| 1607 | out = self.transformer( |
| 1608 | patch_embeds, |
| 1609 | position_embedding, |
| 1610 | cu_seqlens, |
| 1611 | max_seqlen, |
| 1612 | sequence_lengths, |
| 1613 | return_all_hidden_states=select_layers is not None, |
| 1614 | ) |
| 1615 | |
| 1616 | out = resolve_visual_encoder_outputs( |
| 1617 | out, |
| 1618 | None, |
| 1619 | select_layers=select_layers, |
| 1620 | max_possible_layers=self.config.num_hidden_layers, |
| 1621 | feature_select_strategy=feature_select_strategy, |
| 1622 | ) |
| 1623 | |
| 1624 | # squeeze dim 0 and split into separate tensors for each image |
| 1625 | return torch.split(out.squeeze(0), embed_sizes) |
| 1626 | |
| 1627 | # (TODO) Add prefix argument for filtering out weights to be loaded |
| 1628 | # ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986 |
| 1629 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1630 | stacked_params_mapping = [ |
| 1631 | # (param_name, shard_name, shard_id) |
| 1632 | (".qkv_proj", ".q_proj", "q"), |
| 1633 | (".qkv_proj", ".k_proj", "k"), |
| 1634 | (".qkv_proj", ".v_proj", "v"), |
| 1635 | (".gate_up_proj", ".gate_proj", 0), |
| 1636 | (".gate_up_proj", ".up_proj", 1), |
| 1637 | ] |
| 1638 | params_dict = dict(self.named_parameters()) |
| 1639 | loaded_params: set[str] = set() |
| 1640 | layer_count = len(self.transformer.layers) |
| 1641 | |
| 1642 | for name, loaded_weight in weights: |
| 1643 | # omit layers when num_hidden_layers_override is set |
| 1644 | if name.startswith("transformer.layers"): |
| 1645 | layer_idx = int(name.split(".")[2]) |
| 1646 | if layer_idx >= layer_count: |
| 1647 | continue |
| 1648 | |
| 1649 | for param_name, weight_name, shard_id in stacked_params_mapping: |
| 1650 | if weight_name not in name: |
| 1651 | continue |
| 1652 | name = name.replace(weight_name, param_name) |
| 1653 | param = params_dict[name] |
| 1654 | weight_loader = param.weight_loader |
| 1655 | weight_loader(param, loaded_weight, shard_id) |
| 1656 | break |
| 1657 | else: |
| 1658 | param = params_dict[name] |
| 1659 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 1660 | weight_loader(param, loaded_weight) |
| 1661 | loaded_params.add(name) |
| 1662 | return loaded_params |