"vllm/vscode:/vscode.git/clone" did not exist on "f776086d8c0fa4d136f9ccba274b8adb2daf1b55"
pixtral.py 49 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import math
5
from collections.abc import Iterable, Mapping, Sequence
Patrick von Platen's avatar
Patrick von Platen committed
6
from dataclasses import dataclass, fields
7
from functools import cached_property
8
from typing import Annotated, Literal
Patrick von Platen's avatar
Patrick von Platen committed
9
10
11
12

import torch
import torch.nn as nn
import torch.nn.functional as F
13
14
from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk
from mistral_common.protocol.instruct.messages import UserMessage
15
from mistral_common.protocol.instruct.request import ChatCompletionRequest
16
from mistral_common.tokens.tokenizers.multimodal import ImageEncoder
Patrick von Platen's avatar
Patrick von Platen committed
17
from PIL import Image
18
from transformers import BatchFeature, PixtralVisionConfig, TensorType
19
from transformers.image_utils import ImageInput
20
from transformers.models.pixtral.image_processing_pixtral import (
21
22
    _num_image_tokens as _get_pixtral_hf_num_image_tokens,
)
23
from transformers.models.pixtral.modeling_pixtral import (
24
25
26
27
    PixtralRotaryEmbedding,
    apply_rotary_pos_emb,
    position_ids_in_meshgrid,
)
28
from transformers.tokenization_utils_base import TextInput
Patrick von Platen's avatar
Patrick von Platen committed
29

30
from vllm.config import VllmConfig
31
from vllm.config.multimodal import BaseDummyOptions, MultiModalConfig
32
from vllm.distributed import divide, get_tensor_model_parallel_world_size
33
from vllm.model_executor.layers.activation import get_act_and_mul_fn
34
from vllm.model_executor.layers.conv import Conv2dLayer
Patrick von Platen's avatar
Patrick von Platen committed
35
from vllm.model_executor.layers.layernorm import RMSNorm
36
37
38
39
40
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
Patrick von Platen's avatar
Patrick von Platen committed
41
42
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
44
45
46
47
48
49
50
from vllm.multimodal.inputs import (
    MultiModalDataDict,
    MultiModalFieldConfig,
    MultiModalUUIDDict,
    NestedTensors,
)
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
51
52
from vllm.multimodal.processing import BaseDummyInputsBuilder, ProcessorInputs
from vllm.multimodal.processing.processor import (
53
54
55
56
57
58
59
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    MultiModalProcessingInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
)
60
from vllm.platforms import current_platform
61
from vllm.sequence import IntermediateTensors
62
63
from vllm.tokenizers import cached_tokenizer_from_config
from vllm.tokenizers.mistral import MistralTokenizer
64
from vllm.utils.tensor_schema import TensorSchema, TensorShape
Patrick von Platen's avatar
Patrick von Platen committed
65

66
67
68
69
70
71
72
from .interfaces import (
    MultiModalEmbeddings,
    SupportsLoRA,
    SupportsMultiModal,
    SupportsPP,
)
from .module_mapping import MultiModelKeys
73
from .utils import init_vllm_registered_model, maybe_prefix
74
75
76
77
78
from .vision import (
    VisionEncoderInfo,
    VisionFeatureSelectStrategy,
    resolve_visual_encoder_outputs,
)
Patrick von Platen's avatar
Patrick von Platen committed
79

80
try:
81
    # Note: vLLM does not install xformers by default.
82
    from xformers import ops as xops
83
84

    if current_platform.is_cuda() and current_platform.has_device_capability(100):
85
86
87
88
        # Xformers FA is not compatible with B200
        USE_XFORMERS_OPS = False
    else:
        USE_XFORMERS_OPS = True
89
90
91
except ImportError:
    USE_XFORMERS_OPS = False

Patrick von Platen's avatar
Patrick von Platen committed
92
93
PATCH_MERGE = "patch_merge"

Patrick von Platen's avatar
Patrick von Platen committed
94

95
class PixtralImagePixelInputs(TensorSchema):
96
    """
97
98
99
100
101
    Dimensions:
        - bn: Batch size * number of images
        - c: Number of channels (3)
        - h: Height of each image
        - w: Width of each image
102

103
    The result of stacking `ImageEncoding.tokens` from each prompt.
104
    """
105

106
107
    type: Literal["pixel_values"] = "pixel_values"

108
    images: Annotated[
109
        torch.Tensor | list[torch.Tensor],
110
111
        TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}),
    ]
Patrick von Platen's avatar
Patrick von Platen committed
112
113


114
115
116
class PixtralProcessorAdapter:
    """
    Provide a HF-compatible interface for
117
    `mistral_common.tokens.tokenizers.multimodal.ImageEncoder`.
118
    """
Patrick von Platen's avatar
Patrick von Platen committed
119

120
121
    def __init__(self, tokenizer: MistralTokenizer) -> None:
        super().__init__()
Patrick von Platen's avatar
Patrick von Platen committed
122

123
        self.tokenizer = tokenizer
Patrick von Platen's avatar
Patrick von Platen committed
124

125
126
127
128
129
    @property
    def image_processor(self) -> ImageEncoder:
        image_encoder = self.tokenizer.instruct.mm_encoder
        assert isinstance(image_encoder, ImageEncoder)
        return image_encoder
130

131
132
133
    @cached_property
    def image_break_id(self) -> int:
        return self.image_processor.special_ids.img_break
Patrick von Platen's avatar
Patrick von Platen committed
134

135
136
137
    @cached_property
    def image_token_id(self) -> int:
        return self.image_processor.special_ids.img
Patrick von Platen's avatar
Patrick von Platen committed
138

139
140
141
    @cached_property
    def image_end_id(self) -> int:
        return self.image_processor.special_ids.img_end
Patrick von Platen's avatar
Patrick von Platen committed
142

143
144
145
    @cached_property
    def image_size(self) -> int:
        return self.image_processor.mm_config.max_image_size
Patrick von Platen's avatar
Patrick von Platen committed
146

147
148
149
150
151
152
    @cached_property
    def patch_size(self) -> int:
        return self.image_processor.mm_config.image_patch_size

    def __call__(
        self,
153
154
155
        text: TextInput | list[TextInput] | None = None,
        images: ImageInput | list[ImageInput] | None = None,
        return_tensors: str | TensorType | None = None,
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
        **kwargs,
    ) -> Mapping[str, NestedTensors]:
        if text is None:
            text = []
        if not isinstance(text, list):
            text = [text]
        if images is None:
            images = []
        if not isinstance(images, list):
            images = [images]

        if not images:
            input_ids = self.tokenizer(text).input_ids

            return {"input_ids": torch.tensor(input_ids)}

        # Allow dummy text, which is used for profiling as well as token inputs
        if any(len(t) > 0 for t in text):
            raise ValueError(
                "You've passed text inputs instead of token inputs. "
                "Make sure to process your input via `mistral_common`'s "
                "tokenizer or pass a chat completion request. "
                "For more info, see: "
179
180
                "https://github.com/vllm-project/vllm/issues/8411."
            )
181
182
183
184
185
186
187
188
189
190
191
192

        images_processed = list[torch.Tensor]()
        images_tokens = list[torch.Tensor]()

        for image in images:
            image_inputs = self.image_processor(ImageChunk(image=image))
            image_processed = torch.tensor(image_inputs.image)
            image_tokens = torch.tensor(image_inputs.tokens)

            images_processed.append(image_processed)
            images_tokens.append(image_tokens)

193
194
195
196
197
198
        return BatchFeature(
            {
                "input_ids": torch.cat(images_tokens)[None].expand(len(text), -1),
                "images": images_processed,
            }
        )
199
200
201
202


class PixtralProcessingInfo(BaseProcessingInfo):
    def get_tokenizer(self) -> MistralTokenizer:
203
        tokenizer = cached_tokenizer_from_config(self.ctx.model_config)
204
205
206
207
208
209
210
211
        if not isinstance(tokenizer, MistralTokenizer):
            raise ValueError("This model requires `--tokenizer-mode mistral`")

        return tokenizer

    def get_hf_processor(self) -> PixtralProcessorAdapter:
        return PixtralProcessorAdapter(self.get_tokenizer())

212
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
213
214
215
216
        return {"image": None}

    def get_vision_config(
        self,
217
        processor: PixtralProcessorAdapter | None = None,
218
219
220
221
222
223
224
225
226
227
228
229
230
231
    ):
        if processor is None:
            processor = self.get_hf_processor()

        return PixtralVisionConfig(
            image_size=processor.image_size,
            patch_size=processor.patch_size,
        )

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
232
        processor: PixtralProcessorAdapter | None = None,
233
234
235
236
237
    ) -> int:
        if processor is None:
            processor = self.get_hf_processor()

        ncols, nrows = processor.image_processor._image_to_num_tokens(
238
239
            Image.new("RGB", (image_width, image_height))
        )
240

241
        return ncols * nrows
242
243
244
245
246
247
248
249
250

    def get_image_size_with_most_features(self) -> ImageSize:
        image_processor = self.get_hf_processor().image_processor
        max_image_size = image_processor.mm_config.max_image_size

        return ImageSize(width=max_image_size, height=max_image_size)


class PixtralDummyInputsBuilder(BaseDummyInputsBuilder[PixtralProcessingInfo]):
251
252
253
254
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        return ""

    def get_dummy_mm_data(
255
256
257
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
258
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
259
    ) -> MultiModalDataDict:
260
261
        num_images = mm_counts.get("image", 0)

262
        target_width, target_height = self.info.get_image_size_with_most_features()
263

264
265
        image_overrides = mm_options.get("image") if mm_options else None

266
        return {
267
268
269
270
271
272
            "image": self._get_dummy_images(
                width=target_width,
                height=target_height,
                num_images=num_images,
                overrides=image_overrides,
            )
273
274
        }

275
276
277
278
    def get_dummy_processor_inputs(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
279
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
280
281
282
283
    ) -> ProcessorInputs:
        tokenizer = self.info.get_tokenizer()

        dummy_text = self.get_dummy_text(mm_counts)
284
        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
285
        dummy_images = dummy_mm_data.get("image", [])
286
        tokenization_kwargs = {"truncation": False}
287

288
289
290
291
292
293
294
295
296
297
        request = ChatCompletionRequest(
            messages=[
                UserMessage(
                    content=[
                        TextChunk(text=dummy_text),
                        *(ImageChunk(image=image) for image in dummy_images),
                    ]
                ),
            ]
        )
298
299
300
        res = tokenizer.mistral.encode_chat_completion(request)
        dummy_tokens = res.tokens

301
302
303
304
305
        return ProcessorInputs(
            prompt=dummy_tokens,
            mm_data=dummy_mm_data,
            tokenization_kwargs=tokenization_kwargs,
        )
306

Patrick von Platen's avatar
Patrick von Platen committed
307

308
class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]):
309
310
311
312
313
    def _get_mm_fields_config(
        self,
        hf_inputs: Mapping[str, NestedTensors],
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
314
        return dict(images=MultiModalFieldConfig.batched("image"))
315
316
317
318
319

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
320
        out_mm_kwargs: MultiModalKwargsItems,
321
322
323
324
325
326
327
328
329
330
331
332
    ) -> Sequence[PromptUpdate]:
        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

        image_break_id = processor.image_break_id
        image_token_id = processor.image_token_id
        image_end_id = processor.image_end_id

        def get_replacement(item_idx: int):
            images = mm_items.get_items("image", ImageProcessorItems)
            image_size = images.get_image_size(item_idx)

            ncols, nrows = processor.image_processor._image_to_num_tokens(
333
334
                Image.new("RGB", (image_size.width, image_size.height))
            )
335
336
337
338

            tokens = ([image_token_id] * ncols + [image_break_id]) * nrows
            tokens[-1] = image_end_id

339
            return PromptUpdateDetails.select_token_id(tokens, image_token_id)
340
341
342
343
344
345
346
347
348
349
350

        return [
            PromptReplacement(
                modality="image",
                target="",  # Never match the prompt (see below note)
                replacement=get_replacement,
            ),
        ]

    def _cached_apply_hf_processor(
        self,
351
        prompt: str | list[int],
352
353
        mm_data_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
354
        tokenization_kwargs: Mapping[str, object],
355
        mm_uuids: MultiModalUUIDDict | None = None,
356
357
    ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
        prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(
358
359
360
            prompt=prompt,
            mm_data_items=mm_data_items,
            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
361
            tokenization_kwargs=tokenization_kwargs,
362
            mm_uuids=mm_uuids,
363
364
365
        )

        # NOTE: The tokens are already inserted by the chat template
366
        return prompt_ids, mm_info, True
Patrick von Platen's avatar
Patrick von Platen committed
367

368

369
370
371
372
373
@MULTIMODAL_REGISTRY.register_processor(
    PixtralMultiModalProcessor,
    info=PixtralProcessingInfo,
    dummy_inputs=PixtralDummyInputsBuilder,
)
374
375
376
class PixtralForConditionalGeneration(
    nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP
):
377
    @classmethod
378
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
379
380
381
382
383
        if modality.startswith("image"):
            return None

        raise ValueError("Only image modality is supported")

384
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Patrick von Platen's avatar
Patrick von Platen committed
385
        super().__init__()
386
387
        config = vllm_config.model_config.hf_config
        multimodal_config = vllm_config.model_config.multimodal_config
Patrick von Platen's avatar
Patrick von Platen committed
388
389
390
391
392
393
394
395
396
397
398
399
400
        self.config = config
        self.multimodal_config = multimodal_config

        dataclass_fields = {field.name for field in fields(VisionEncoderArgs)}
        vision_args = {
            key: value
            for key, value in self.config.vision_config.to_dict().items()
            if key in dataclass_fields
        }

        self.vision_args = VisionEncoderArgs(**vision_args)

        # init MistralForCausalLM
401
402
403
404
405
406
        with self._mark_language_model(vllm_config):
            self.language_model = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.text_config,
                prefix=maybe_prefix(prefix, "language_model"),
            )
Patrick von Platen's avatar
Patrick von Platen committed
407

408
        with self._mark_tower_model(vllm_config, "image"):
409
410
411
412
413
            self.vision_encoder = VisionTransformer(self.vision_args)
            self.pre_mm_projector_norm = (
                RMSNorm(self.vision_args.hidden_size, eps=1e-5)
                if self.vision_args.add_pre_mm_projector_layer_norm
                else None
Patrick von Platen's avatar
Patrick von Platen committed
414
            )
415
416
417
418
419
420
421
422
423
424
425
426
            self.patch_merger = (
                PatchMerger(
                    vision_encoder_dim=self.vision_args.hidden_size,
                    spatial_merge_size=self.vision_args.spatial_merge_size,
                    use_mlp_bias=False,
                )
                if self.vision_args.mm_projector_id == PATCH_MERGE
                else None
            )
            self.vision_language_adapter = VisionLanguageAdapter(
                self.vision_args, dim=config.text_config.hidden_size
            )
Patrick von Platen's avatar
Patrick von Platen committed
427

428
        self.make_empty_intermediate_tensors = (
429
430
            self.language_model.make_empty_intermediate_tensors
        )
431

432
    def _parse_and_validate_image_input(
433
        self, **kwargs: object
434
    ) -> PixtralImagePixelInputs | None:
435
436
437
438
439
440
        images = kwargs.pop("images", None)
        if images is None:
            return None

        return PixtralImagePixelInputs(
            type="pixel_values",
441
            images=images,
442
443
444
445
446
447
448
449
        )

    def _process_image_input(
        self,
        image_input: PixtralImagePixelInputs,
    ) -> tuple[torch.Tensor, ...]:
        images = image_input["images"]
        image_features = self.vision_encoder(images)
450
        feature_sizes = [image_feature.shape[0] for image_feature in image_features]
Patrick von Platen's avatar
Patrick von Platen committed
451
        image_features = torch.cat(image_features)
452
        if self.pre_mm_projector_norm is not None:
Patrick von Platen's avatar
Patrick von Platen committed
453
            image_features = self.pre_mm_projector_norm(image_features)
454
        if self.patch_merger is not None:
Patrick von Platen's avatar
Patrick von Platen committed
455
456
            patch_size = self.vision_args.patch_size
            spatial_merge_size_square = self.vision_args.spatial_merge_size**2
457
458
459
460
            img_patch_dims = [
                (img.shape[1] // patch_size, img.shape[2] // patch_size)
                for img in images
            ]
Patrick von Platen's avatar
Patrick von Platen committed
461
462
463
464
            feature_sizes = [
                feature_size // spatial_merge_size_square
                for feature_size in feature_sizes
            ]
465
466
467
            image_features = self.patch_merger(
                image_features, image_sizes=img_patch_dims
            )
Patrick von Platen's avatar
Patrick von Platen committed
468
        image_embeds = self.vision_language_adapter(image_features)
469
470
471
        image_embeds = torch.split(image_embeds, feature_sizes)
        return image_embeds

472
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
473
        image_input = self._parse_and_validate_image_input(**kwargs)
474
        if image_input is None:
475
            return []
476

477
        return self._process_image_input(image_input)
478

Patrick von Platen's avatar
Patrick von Platen committed
479
480
481
482
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
483
484
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Patrick von Platen's avatar
Patrick von Platen committed
485
        **kwargs: object,
486
    ) -> torch.Tensor | IntermediateTensors:
487
        """Run forward pass for pixtral."""
488
489
        if intermediate_tensors is not None:
            inputs_embeds = None
Patrick von Platen's avatar
Patrick von Platen committed
490

491
492
493
        hidden_states = self.language_model.model(
            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
        )
Patrick von Platen's avatar
Patrick von Platen committed
494
495
496
497
498
499

        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
500
    ) -> torch.Tensor | None:
501
        return self.language_model.compute_logits(hidden_states)
Patrick von Platen's avatar
Patrick von Platen committed
502

503
504
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]):
Patrick von Platen's avatar
Patrick von Platen committed
505
506
            return weight[0].startswith("vision_encoder")

507
        def is_vision_lang_adapter_weights(weight: tuple[str, torch.Tensor]):
Patrick von Platen's avatar
Patrick von Platen committed
508
509
            return weight[0].startswith("vision_language_adapter")

510
        def is_patch_merger(weight: tuple[str, torch.Tensor]):
Patrick von Platen's avatar
Patrick von Platen committed
511
512
            return weight[0].startswith("patch_merger")

513
        def is_pre_mm_projector_norm(weight: tuple[str, torch.Tensor]):
Patrick von Platen's avatar
Patrick von Platen committed
514
515
            return weight[0].startswith("pre_mm_projector_norm")

516
        # Get references to parameters for direct loading
517
518
519
520
521
        vision_encoder_dict = (
            dict(self.vision_encoder.named_parameters())
            if self.vision_encoder is not None
            else {}
        )
522
523
        patch_merger_dict = (
            dict(self.patch_merger.named_parameters())
524
525
            if self.patch_merger is not None
            else {}
526
527
528
        )
        pre_mm_projector_norm_dict = (
            dict(self.pre_mm_projector_norm.named_parameters())
529
530
531
532
533
534
535
            if self.pre_mm_projector_norm is not None
            else {}
        )
        vision_lang_adapter_dict = (
            dict(self.vision_language_adapter.named_parameters())
            if self.vision_language_adapter is not None
            else {}
536
        )
537
538
539
540
541

        def llm_weights_generator():
            # Single pass over weights
            for name, w in weights:
                if is_vision_encoder_weights((name, w)):
542
543
                    if self.vision_encoder is None:
                        continue
544
                    # Load vision encoder weights directly
545
                    trimmed_name = ".".join(name.split(".")[1:])
546
547
548
                    param = vision_encoder_dict[trimmed_name]
                    with torch.no_grad():
                        default_weight_loader(param, w)
Patrick von Platen's avatar
Patrick von Platen committed
549
                elif is_patch_merger((name, w)):
550
551
                    if self.patch_merger is None:
                        continue
Patrick von Platen's avatar
Patrick von Platen committed
552
                    # Load vision patch merger weights directly
553
                    trimmed_name = ".".join(name.split(".")[1:])
Patrick von Platen's avatar
Patrick von Platen committed
554
555
556
557
                    param = patch_merger_dict[trimmed_name]
                    with torch.no_grad():
                        default_weight_loader(param, w)
                elif is_pre_mm_projector_norm((name, w)):
558
559
                    if self.pre_mm_projector_norm is None:
                        continue
Patrick von Platen's avatar
Patrick von Platen committed
560
                    # Load vision pre_mm_projector_norm weights directly
561
                    trimmed_name = ".".join(name.split(".")[1:])
Patrick von Platen's avatar
Patrick von Platen committed
562
563
564
                    param = pre_mm_projector_norm_dict[trimmed_name]
                    with torch.no_grad():
                        default_weight_loader(param, w)
565
                elif is_vision_lang_adapter_weights((name, w)):
566
567
                    if self.vision_language_adapter is None:
                        continue
568
                    # Load vision-language adapter weights directly
569
                    trimmed_name = ".".join(name.split(".")[1:])
570
571
572
573
574
575
576
577
578
579
                    param = vision_lang_adapter_dict[trimmed_name]
                    with torch.no_grad():
                        default_weight_loader(param, w)
                else:
                    # LLM weights: yield them to be loaded
                    # by language_model.load_weights
                    yield (name, w)

        # Now we call the language model load with the generator
        self.language_model.load_weights(llm_weights_generator())
Patrick von Platen's avatar
Patrick von Platen committed
580

581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
    def get_mm_mapping(self) -> MultiModelKeys:
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector="vision_language_adapter",
            tower_model="vision_encoder",
        )

    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
        if getattr(self, "patch_merger", None) is None:
            return num_image_tokens
        merge_size = self.vision_args.spatial_merge_size
        return num_image_tokens * (merge_size**2)

    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
        if getattr(self, "patch_merger", None) is None:
            return num_vision_tokens
        merge_size = self.vision_args.spatial_merge_size
        return num_vision_tokens // (merge_size**2)

Patrick von Platen's avatar
Patrick von Platen committed
600
601
602
603
604
605
606
607
608
609
610
611
612

# Vision encoder
@dataclass
class VisionEncoderArgs:
    hidden_size: int
    num_channels: int
    image_size: int
    patch_size: int
    intermediate_size: int
    num_hidden_layers: int
    num_attention_heads: int
    rope_theta: float  # for rope-2D
    image_token_id: int
613
    adapter_bias: bool = True
Patrick von Platen's avatar
Patrick von Platen committed
614
615
616
    spatial_merge_size: int = 1
    add_pre_mm_projector_layer_norm: bool = False
    mm_projector_id: str = ""
Patrick von Platen's avatar
Patrick von Platen committed
617
618


619
def _reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
Patrick von Platen's avatar
Patrick von Platen committed
620
621
622
623
624
625
626
627
628
629
    """
    freqs_cis: complex - (seq_len, head_dim / 2)
    x: complex - (bsz, seq_len, head_dim / 2)
    """
    ndim = x.ndim
    assert ndim > 1
    assert freqs_cis.shape == (x.shape[1], x.shape[-1]), (
        freqs_cis.shape,
        (x.shape[1], x.shape[-1]),
    )
630
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
Patrick von Platen's avatar
Patrick von Platen committed
631
632
633
634
635
636
637
638
639
640
641
642
643
644
    return freqs_cis.view(*shape)


def precompute_freqs_cis_2d(
    dim: int,
    height: int,
    width: int,
    theta: float,
) -> torch.Tensor:
    """
    freqs_cis: 2D complex tensor of shape (height, width, dim // 2)
        to be indexed by (height, width) position tuples
    """
    # (dim / 2) frequency bases
645
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
Patrick von Platen's avatar
Patrick von Platen committed
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665

    h = torch.arange(height, device=freqs.device)
    w = torch.arange(width, device=freqs.device)

    freqs_h = torch.outer(h, freqs[::2]).float()
    freqs_w = torch.outer(w, freqs[1::2]).float()
    freqs_2d = torch.cat(
        [
            freqs_h[:, None, :].repeat(1, width, 1),
            freqs_w[None, :, :].repeat(height, 1, 1),
        ],
        dim=-1,
    )
    return torch.polar(torch.ones_like(freqs_2d), freqs_2d)


def apply_rotary_emb_vit(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
666
) -> tuple[torch.Tensor, torch.Tensor]:
Patrick von Platen's avatar
Patrick von Platen committed
667
668
669
670
671
672
673
674
675
676
677
678
679
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    assert freqs_cis.dtype == torch.complex64
    freqs_cis = _reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)


class FeedForward(nn.Module):
    def __init__(self, args: VisionEncoderArgs):
        super().__init__()
        assert args.intermediate_size is not None
680
681
682
        self.w1 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
        self.w2 = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
        self.w3 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
Patrick von Platen's avatar
Patrick von Platen committed
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class Attention(nn.Module):
    def __init__(self, args: VisionEncoderArgs):
        super().__init__()
        self.args = args
        assert not args.hidden_size % args.num_attention_heads
        self.n_heads = args.num_attention_heads
        self.head_dim = args.hidden_size // args.num_attention_heads

        self.wq = nn.Linear(args.hidden_size, args.hidden_size, bias=False)
        self.wk = nn.Linear(args.hidden_size, args.hidden_size, bias=False)
        self.wv = nn.Linear(args.hidden_size, args.hidden_size, bias=False)
        self.wo = nn.Linear(args.hidden_size, args.hidden_size, bias=False)

    def forward(
        self,
        x: torch.Tensor,
704
        mask: torch.Tensor,
Patrick von Platen's avatar
Patrick von Platen committed
705
706
707
708
709
710
711
712
713
714
        freqs_cis: torch.Tensor,
    ) -> torch.Tensor:
        batch, patches, _ = x.shape

        q, k, v = self.wq(x), self.wk(x), self.wv(x)
        q = q.reshape(batch, patches, self.n_heads, self.head_dim)
        k = k.reshape(batch, patches, self.n_heads, self.head_dim)
        v = v.reshape(batch, patches, self.n_heads, self.head_dim)

        q, k = apply_rotary_emb_vit(q, k, freqs_cis=freqs_cis)
715
716
717
718
719
720
721

        if USE_XFORMERS_OPS:
            out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)
        else:
            q = q.transpose(1, 2)
            k = k.transpose(1, 2)
            v = v.transpose(1, 2)
722
            out = nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
723
724
            out = out.transpose(1, 2)

Patrick von Platen's avatar
Patrick von Platen committed
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
        out = out.reshape(batch, patches, self.n_heads * self.head_dim)
        return self.wo(out)


class TransformerBlock(nn.Module):
    def __init__(self, args: VisionEncoderArgs):
        super().__init__()
        self.attention = Attention(args)
        self.feed_forward = FeedForward(args)
        self.attention_norm = RMSNorm(args.hidden_size, eps=1e-5)
        self.ffn_norm = RMSNorm(args.hidden_size, eps=1e-5)

    def forward(
        self,
        x: torch.Tensor,
740
        mask: torch.Tensor,
Patrick von Platen's avatar
Patrick von Platen committed
741
742
        freqs_cis: torch.Tensor,
    ) -> torch.Tensor:
743
744
745
        r = self.attention.forward(
            self.attention_norm(x), mask=mask, freqs_cis=freqs_cis
        )
Patrick von Platen's avatar
Patrick von Platen committed
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
        h = x + r
        r = self.feed_forward.forward(self.ffn_norm(h))
        out = h + r
        return out


class Transformer(nn.Module):
    def __init__(self, args: VisionEncoderArgs):
        super().__init__()
        self.layers = torch.nn.ModuleList()
        for _ in range(args.num_hidden_layers):
            self.layers.append(TransformerBlock(args))

    def forward(
        self,
        x: torch.Tensor,
762
        mask: torch.Tensor,
763
        freqs_cis: torch.Tensor | None,
Patrick von Platen's avatar
Patrick von Platen committed
764
765
766
767
768
769
    ) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, mask=mask, freqs_cis=freqs_cis)
        return x


770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
def position_meshgrid(
    patch_embeds_list: list[torch.Tensor],
) -> torch.Tensor:
    positions = torch.cat(
        [
            torch.stack(
                torch.meshgrid(
                    torch.arange(p.shape[-2]),
                    torch.arange(p.shape[-1]),
                    indexing="ij",
                ),
                dim=-1,
            ).reshape(-1, 2)
            for p in patch_embeds_list
        ]
    )
Patrick von Platen's avatar
Patrick von Platen committed
786
787
788
789
790
791
792
    return positions


class VisionTransformer(nn.Module):
    def __init__(self, args: VisionEncoderArgs):
        super().__init__()
        self.args = args
793
        self.patch_conv = Conv2dLayer(
Patrick von Platen's avatar
Patrick von Platen committed
794
795
796
797
798
799
800
801
802
803
804
            in_channels=args.num_channels,
            out_channels=args.hidden_size,
            kernel_size=args.patch_size,
            stride=args.patch_size,
            bias=False,
        )
        self.ln_pre = RMSNorm(args.hidden_size, eps=1e-5)
        self.transformer = Transformer(args)

        head_dim = self.args.hidden_size // self.args.num_attention_heads
        assert head_dim % 2 == 0, "ROPE requires even head_dim"
805
        self._freqs_cis: torch.Tensor | None = None
Patrick von Platen's avatar
Patrick von Platen committed
806
807
808
809
810
811

    @property
    def max_patches_per_side(self) -> int:
        return self.args.image_size // self.args.patch_size

    @property
812
    def device(self) -> torch.types.Device:
Patrick von Platen's avatar
Patrick von Platen committed
813
814
815
        return next(self.parameters()).device

    @property
816
    def dtype(self) -> torch.dtype:
Patrick von Platen's avatar
Patrick von Platen committed
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
        return next(self.parameters()).dtype

    @property
    def freqs_cis(self) -> torch.Tensor:
        if self._freqs_cis is None:
            self._freqs_cis = precompute_freqs_cis_2d(
                dim=self.args.hidden_size // self.args.num_attention_heads,
                height=self.max_patches_per_side,
                width=self.max_patches_per_side,
                theta=self.args.rope_theta,
            )

        if self._freqs_cis.device != self.device:
            self._freqs_cis = self._freqs_cis.to(device=self.device)

        return self._freqs_cis

    def forward(
        self,
836
        images: list[torch.Tensor],
Patrick von Platen's avatar
Patrick von Platen committed
837
838
839
    ) -> torch.Tensor:
        """
        Args:
840
            images: list of N_img images of variable sizes,
Patrick von Platen's avatar
Patrick von Platen committed
841
842
                each of shape (C, H, W)
        Returns:
843
            image_features: tensor of token features for
Patrick von Platen's avatar
Patrick von Platen committed
844
845
846
847
848
849
850
                all tokens of all images of shape (N_toks, D)
        """
        # pass images through initial convolution independently
        patch_embeds_list = [
            self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in images
        ]

851
        patch_embeds = [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list]
852
853
        embed_sizes = [p.shape[1] for p in patch_embeds]

Patrick von Platen's avatar
Patrick von Platen committed
854
        # flatten to a single sequence
855
        patch_embeds = torch.cat(patch_embeds, dim=1)
Patrick von Platen's avatar
Patrick von Platen committed
856
857
858
859
860
861
862
        patch_embeds = self.ln_pre(patch_embeds)

        # positional embeddings
        positions = position_meshgrid(patch_embeds_list).to(self.device)
        freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]

        # pass through Transformer with a block diagonal mask delimiting images
863
864
        if USE_XFORMERS_OPS:
            mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(
865
866
                [p.shape[-2] * p.shape[-1] for p in patch_embeds_list],
            )
867
        else:
868
            from transformers.models.pixtral.modeling_pixtral import (
869
870
871
                generate_block_attention_mask,
            )

872
            mask = generate_block_attention_mask(
873
874
                [p.shape[-2] * p.shape[-1] for p in patch_embeds_list], patch_embeds
            )
Patrick von Platen's avatar
Patrick von Platen committed
875
876
        out = self.transformer(patch_embeds, mask=mask, freqs_cis=freqs_cis)

877
878
        # squeeze dim 0 and split into separate tensors for each image
        return torch.split(out.squeeze(0), embed_sizes)
Patrick von Platen's avatar
Patrick von Platen committed
879
880
881
882
883
884
885
886
887


class VisionLanguageAdapter(nn.Module):
    def __init__(self, args: VisionEncoderArgs, dim: int):
        super().__init__()
        assert isinstance(args, VisionEncoderArgs)
        self.w_in = nn.Linear(
            args.hidden_size,
            dim,
888
            bias=args.adapter_bias,
Patrick von Platen's avatar
Patrick von Platen committed
889
890
        )
        self.gelu = nn.GELU()
891
        self.w_out = nn.Linear(dim, dim, bias=args.adapter_bias)
Patrick von Platen's avatar
Patrick von Platen committed
892
893
894

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_out(self.gelu(self.w_in(x)))
895
896


Patrick von Platen's avatar
Patrick von Platen committed
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
class PatchMerger(nn.Module):
    """
    Learned merging of spatial_merge_size ** 2 patches
    """

    def __init__(
        self,
        vision_encoder_dim: int,
        spatial_merge_size: int,
        use_mlp_bias: bool = False,
    ) -> None:
        super().__init__()

        mlp_input_dim = vision_encoder_dim * (spatial_merge_size**2)

        self.spatial_merge_size = spatial_merge_size
        self.mlp_input_dim = mlp_input_dim

        self.merging_layer = nn.Linear(
            mlp_input_dim,
            vision_encoder_dim,
            bias=use_mlp_bias,
        )

921
922
923
    def forward(
        self, x: torch.Tensor, image_sizes: list[tuple[int, int]]
    ) -> torch.Tensor:
Patrick von Platen's avatar
Patrick von Platen committed
924
925
926
927
928
929
        # image_sizes specified in tokens
        assert sum([h * w for h, w in image_sizes]) == len(x)

        # x is (N, vision_encoder_dim)
        x = self.permute(x, image_sizes)

930
931
        # x is (N / spatial_merge_size ** 2,
        #       vision_encoder_dim * spatial_merge_size ** 2)
Patrick von Platen's avatar
Patrick von Platen committed
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
        x = self.merging_layer(x)

        # x is (N / spatial_merge_size ** 2, vision_encoder_dim)
        return x

    def permute(
        self,
        x: torch.Tensor,
        image_sizes: list[tuple[int, int]],
    ) -> torch.Tensor:
        """
        Args:
            x: (N, D) where N is flattened and concatenated patch tokens
                for all images
            image_sizes: list of tuple of (height, width) in tokens for
                each image
        Returns:
            image_features: reorders patch tokens so each grid of
                (spatial_merge_size, spatial_merge_size) is contiguous.
                now (N / spatial_merge_size ** 2, D * spatial_merge_size ** 2)
        """

        sub_grids = get_sub_grids(
955
            x=x, image_sizes=image_sizes, spatial_merge_size=self.spatial_merge_size
Patrick von Platen's avatar
Patrick von Platen committed
956
957
958
959
        )  # list of [d x sub_grid_size x sub_grid_size x n_patches]
        permuted_tensor: list[torch.Tensor] = []
        for grid in sub_grids:
            n_patches = grid.shape[-1]
960
961
962
            permuted_tensor.append(
                grid.view(-1, n_patches).t()
            )  # n_patches x d * sub_grid_size * sub_grid_size
Patrick von Platen's avatar
Patrick von Platen committed
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
        return torch.cat(
            permuted_tensor, dim=0
        )  # (N / spatial_merge_size ** 2, d * spatial_merge_size ** 2)


def get_sub_grids(
    x: torch.Tensor,
    image_sizes: list[tuple[int, int]],
    spatial_merge_size: int,
) -> list[torch.Tensor]:
    # image_sizes specified in tokens
    tokens_per_image = [h * w for h, w in image_sizes]
    d = x.shape[-1]
    all_img_sub_grids: list[torch.Tensor] = []
    sub_grid_size = spatial_merge_size

    for image_index, image_tokens in enumerate(x.split(tokens_per_image)):
        # Reshape image_tokens into a 2D grid
        h, w = image_sizes[image_index]
982
983
984
985
986
987
        image_grid = image_tokens.view(h, w, d).permute(2, 0, 1)[
            None, :, :, :
        ]  # 1 x d x h x w
        sub_grids = torch.nn.functional.unfold(
            image_grid, kernel_size=sub_grid_size, stride=sub_grid_size
        )
Patrick von Platen's avatar
Patrick von Platen committed
988
        sub_grids = sub_grids.view(
989
990
            1, d, sub_grid_size, sub_grid_size, -1
        )  # 1 x d x sub_grid_size x sub_grid_size x n_patches
Patrick von Platen's avatar
Patrick von Platen committed
991
992
993
994
995
996

        all_img_sub_grids.append(sub_grids[0])

    return all_img_sub_grids


997
998
999
1000
1001
1002
1003
1004
#### HF Transformers version of Pixtral ####
# Based off https://github.com/huggingface/transformers/blob/d7950bff82b18c823193d17d72188c5e46d06c83/src/transformers/models/pixtral/modeling_pixtral.py
# This model follows the Llava family, meaning image embeddings are placed
# instead of the `[IMG]` token placeholders.
# The model uses [`PixtralVisionModel`] for its vision encoder,
# and [`MistralForCausalLM`] for its language decoder.


1005
1006
1007
1008
1009
1010
1011
class PixtralHFEncoderInfo(VisionEncoderInfo[PixtralVisionConfig]):
    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
1012
1013
1014
        ncols, nrows = self.get_patch_grid_size(
            image_width=image_width,
            image_height=image_height,
1015
        )
1016
        return ncols * nrows
1017

1018
1019
1020
1021
    def get_image_size(self) -> int:
        return self.vision_config.image_size

    def get_patch_size(self) -> int:
1022
1023
1024
        # spatial_merge_size is needed for Mistral3
        spatial_merge_size = getattr(self.hf_config, "spatial_merge_size", 1)
        return self.vision_config.patch_size * spatial_merge_size
1025
1026

    def get_patch_grid_length(self) -> int:
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
        image_size, patch_size = self.get_image_size(), self.get_patch_size()

        # Since interpolation is applied, the image size need not be divisible
        # assert image_size % patch_size == 0
        return image_size // patch_size

    # Adapted from: https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/pixtral/image_processing_pixtral.py#L99
    def get_patch_grid_size(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> tuple[int, int]:
        max_width = max_height = self.get_image_size()
        patch_width = patch_height = self.get_patch_size()

        ratio = max(image_width / max_width, image_height / max_height)

        if ratio > 1:
1046
1047
            image_width = int(math.floor(image_width / ratio))
            image_height = int(math.floor(image_height / ratio))
1048
1049
1050
1051
1052
1053
1054

        nrows, ncols = _get_pixtral_hf_num_image_tokens(
            (image_height, image_width),
            (patch_height, patch_width),
        )  # type: ignore

        return ncols, nrows
1055
1056
1057


class PixtralHFMLP(nn.Module):
1058
1059
1060
    def __init__(
        self,
        config: PixtralVisionConfig,
1061
        quant_config: QuantizationConfig | None = None,
1062
        multimodal_config: MultiModalConfig | None = None,
1063
1064
1065
        *,
        prefix: str = "",
    ) -> None:
1066
        super().__init__()
1067

1068
1069
1070
1071
1072
1073
        use_data_parallel = (
            multimodal_config.mm_encoder_tp_mode == "data"
            if multimodal_config
            else False
        )

1074
        assert config.intermediate_size is not None
1075
1076
1077
1078
1079
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=config.hidden_size,
            output_sizes=[config.intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
1080
            prefix=f"{prefix}.gate_up_proj",
1081
            disable_tp=use_data_parallel,
1082
1083
1084
1085
1086
1087
1088
        )
        self.down_proj = RowParallelLinear(
            input_size=config.intermediate_size,
            output_size=config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
1089
            disable_tp=use_data_parallel,
1090
        )
1091
        self.act_and_mul = get_act_and_mul_fn(config.hidden_act)
1092
1093

    def forward(self, x: torch.Tensor) -> torch.Tensor:
1094
1095
1096
1097
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_and_mul(gate_up)
        x, _ = self.down_proj(x)
        return x
1098
1099
1100


class PixtralHFAttention(nn.Module):
1101
1102
1103
    def __init__(
        self,
        config: PixtralVisionConfig,
1104
        quant_config: QuantizationConfig | None = None,
1105
        multimodal_config: MultiModalConfig | None = None,
1106
1107
1108
        *,
        prefix: str = "",
    ) -> None:
1109
        super().__init__()
1110

1111
1112
        self.config = config
        assert not config.hidden_size % config.num_attention_heads
1113
        self.total_num_heads = config.num_attention_heads
1114
        self.head_dim = config.hidden_size // config.num_attention_heads
1115
        assert self.total_num_heads * self.head_dim == config.hidden_size
1116

1117
1118
1119
1120
1121
        use_data_parallel = (
            multimodal_config.mm_encoder_tp_mode == "data"
            if multimodal_config
            else False
        )
1122
1123
1124
        self.qkv_proj = QKVParallelLinear(
            hidden_size=config.hidden_size,
            head_size=self.head_dim,
1125
            total_num_heads=self.total_num_heads,
1126
1127
1128
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
1129
            disable_tp=use_data_parallel,
1130
1131
1132
1133
1134
1135
1136
        )
        self.o_proj = RowParallelLinear(
            input_size=config.hidden_size,
            output_size=config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
1137
1138
1139
1140
1141
            disable_tp=use_data_parallel,
        )

        self.tp_size = (
            1 if use_data_parallel else get_tensor_model_parallel_world_size()
1142
        )
1143
        self.n_heads = divide(config.num_attention_heads, self.tp_size)
1144
1145
1146
1147

    def forward(
        self,
        hidden_states: torch.Tensor,
1148
        attention_mask: torch.Tensor,
1149
        position_embeddings: torch.Tensor,
1150
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
1151
        batch, patches, _ = hidden_states.size()
1152

1153
1154
        qkv_states, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv_states.chunk(3, dim=-1)
1155

1156
1157
1158
        # Transpose q and k to apply HF's Rotary Position Embedding
        q = q.view(batch, patches, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, patches, self.n_heads, self.head_dim).transpose(1, 2)
1159
        v = v.view(batch, patches, self.n_heads, self.head_dim)
1160
        cos, sin = position_embeddings
1161
        q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=0)
1162

1163
1164
1165
1166
        if USE_XFORMERS_OPS:
            # Transpose q and k back for attention
            q = q.transpose(1, 2).contiguous()
            k = k.transpose(1, 2).contiguous()
1167
            out = xops.memory_efficient_attention(q, k, v, attn_bias=attention_mask)
1168
        else:
1169
            v = v.transpose(1, 2)
1170
            out = nn.functional.scaled_dot_product_attention(
1171
1172
                q, k, v, attn_mask=attention_mask
            )
1173
            out = out.transpose(1, 2)
1174

1175
        out = out.reshape(batch, patches, self.n_heads * self.head_dim)
1176
        attn_output, _ = self.o_proj(out)
1177

1178
        return attn_output, None
1179
1180
1181


class PixtralHFTransformerBlock(nn.Module):
1182
1183
1184
    def __init__(
        self,
        config: PixtralVisionConfig,
1185
        quant_config: QuantizationConfig | None = None,
1186
        multimodal_config: MultiModalConfig | None = None,
1187
1188
1189
        *,
        prefix: str = "",
    ) -> None:
1190
        super().__init__()
1191

1192
        self.attention_norm = RMSNorm(config.hidden_size, eps=1e-5)
1193
        self.attention = PixtralHFAttention(
1194
1195
1196
1197
            config,
            quant_config=quant_config,
            multimodal_config=multimodal_config,
            prefix=f"{prefix}.attention",
1198
1199
        )
        self.feed_forward = PixtralHFMLP(
1200
1201
1202
1203
            config,
            quant_config=quant_config,
            multimodal_config=multimodal_config,
            prefix=f"{prefix}.feed_forward",
1204
        )
1205
1206
1207
1208
1209
        self.ffn_norm = RMSNorm(config.hidden_size, eps=1e-5)

    def forward(
        self,
        hidden_states: torch.Tensor,
1210
        attention_mask: torch.Tensor,
1211
1212
        position_embeddings: torch.Tensor,
    ) -> torch.Tensor:
1213
1214
1215
1216
1217
        r, _ = self.attention.forward(
            self.attention_norm(hidden_states),
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
        )
1218
1219
1220
1221
1222
1223
1224
        h = hidden_states + r
        r = self.feed_forward.forward(self.ffn_norm(h))
        out = h + r
        return out


class PixtralHFTransformer(nn.Module):
1225
1226
1227
    def __init__(
        self,
        config: PixtralVisionConfig,
1228
        quant_config: QuantizationConfig | None = None,
1229
        multimodal_config: MultiModalConfig | None = None,
1230
        *,
1231
        num_hidden_layers_override: int | None = None,
1232
1233
        prefix: str = "",
    ) -> None:
1234
        super().__init__()
1235
1236
1237
1238
1239
1240

        if num_hidden_layers_override is None:
            num_hidden_layers = config.num_hidden_layers
        else:
            num_hidden_layers = num_hidden_layers_override

1241
1242
1243
1244
1245
        self.layers = nn.ModuleList(
            [
                PixtralHFTransformerBlock(
                    config=config,
                    quant_config=quant_config,
1246
                    multimodal_config=multimodal_config,
1247
1248
1249
1250
1251
                    prefix=f"{prefix}.layers.{layer_idx}",
                )
                for layer_idx in range(num_hidden_layers)
            ]
        )
1252
1253
1254
1255

    def forward(
        self,
        x: torch.Tensor,
1256
        attention_mask: torch.Tensor,
1257
        position_embeddings: torch.Tensor,
1258
        return_all_hidden_states: bool,
1259
    ) -> torch.Tensor:
1260
        hidden_states_pool = [x]
1261

1262
1263
        for layer in self.layers:
            x = layer(x, attention_mask, position_embeddings)
1264
1265
1266
1267
1268
1269
            if return_all_hidden_states:
                hidden_states_pool.append(x)
        # If we have multiple feature sample layers, we return all hidden
        # states in order and grab the ones we need by index.
        if return_all_hidden_states:
            return hidden_states_pool
1270
1271
1272
1273
        return x


class PixtralHFVisionModel(nn.Module):
1274
1275
1276
    def __init__(
        self,
        config: PixtralVisionConfig,
1277
        quant_config: QuantizationConfig | None = None,
1278
        multimodal_config: MultiModalConfig | None = None,
1279
        *,
1280
1281
        num_hidden_layers_override: int | None = None,
        require_post_norm: bool | None = None,
1282
1283
        prefix: str = "",
    ) -> None:
1284
1285
1286
        super().__init__()

        self.config = config
1287

1288
        self.patch_conv = Conv2dLayer(
1289
1290
1291
1292
1293
1294
1295
            in_channels=config.num_channels,
            out_channels=config.hidden_size,
            kernel_size=config.patch_size,
            stride=config.patch_size,
            bias=False,
        )
        self.ln_pre = RMSNorm(config.hidden_size, eps=1e-5)
1296
1297
        self.transformer = PixtralHFTransformer(
            config,
1298
1299
            quant_config=quant_config,
            multimodal_config=multimodal_config,
1300
1301
1302
1303
1304
1305
1306
1307
1308
            num_hidden_layers_override=num_hidden_layers_override,
            prefix=f"{prefix}.transformer",
        )

        num_hidden_layers = config.num_hidden_layers
        if len(self.transformer.layers) > config.num_hidden_layers:
            raise ValueError(
                f"The original encoder only has {num_hidden_layers} "
                f"layers, but you requested {len(self.transformer.layers)} "
1309
1310
                "layers."
            )
1311
1312
1313
1314
1315

        if require_post_norm is True:
            msg = "PixtralHFVisionModel does not have post-layernorm"
            raise ValueError(msg)

1316
1317
        self.dtype = next(self.parameters()).dtype
        self.device = next(self.parameters()).device
1318
        self.patch_positional_embedding = PixtralRotaryEmbedding(config, self.device)
1319
1320
1321

    def forward(
        self,
1322
        pixel_values: list[torch.Tensor],
1323
        *,
1324
1325
        select_layers: list[int] | None = None,
        feature_select_strategy: VisionFeatureSelectStrategy | None = None,
1326
    ) -> tuple[torch.Tensor, ...]:
1327
1328
        """
        Args:
1329
1330
1331
1332
            pixel_values: Each image to be processed will be a separate tensor
                in pixel_values. This means it will be a list of tensors
                because multiple requests batched can have multiple images,
                each with their own shape potentially
1333
            select_layers: Layer indices whose features should be
1334
1335
                concatenated and used as the visual encoder output. If none
                are provided, the last layer is used.
1336

1337
1338
1339
1340
1341
1342
        Returns:
            image_features: tensor of token features for
                all tokens of all images of shape (N_toks, D)
        """
        # pass images through initial convolution independently
        patch_embeds_list = [
1343
            self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in pixel_values
1344
1345
        ]

1346
        patch_embeds = [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list]
1347
1348
        embed_sizes = [p.shape[1] for p in patch_embeds]

1349
        # flatten to a single sequence
1350
        patch_embeds = torch.cat(patch_embeds, dim=1)
1351
1352
1353
1354
1355
        patch_embeds = self.ln_pre(patch_embeds)

        # positional embeddings
        position_ids = position_ids_in_meshgrid(
            patch_embeds_list,
1356
1357
1358
            max_width=self.config.image_size // self.config.patch_size,
        ).to(self.device)
        position_embedding = self.patch_positional_embedding(patch_embeds, position_ids)
1359
1360
1361

        if USE_XFORMERS_OPS:
            attention_mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(
1362
1363
                [p.shape[-2] * p.shape[-1] for p in patch_embeds_list],
            )
1364
1365
        else:
            from transformers.models.pixtral.modeling_pixtral import (
1366
1367
1368
                generate_block_attention_mask,
            )

1369
            attention_mask = generate_block_attention_mask(
1370
1371
                [p.shape[-2] * p.shape[-1] for p in patch_embeds_list], patch_embeds
            )
1372

1373
1374
1375
1376
        out = self.transformer(
            patch_embeds,
            attention_mask,
            position_embedding,
1377
1378
            return_all_hidden_states=select_layers is not None,
        )
1379

1380
1381
1382
1383
1384
1385
1386
        out = resolve_visual_encoder_outputs(
            out,
            None,
            select_layers=select_layers,
            max_possible_layers=self.config.num_hidden_layers,
            feature_select_strategy=feature_select_strategy,
        )
1387

1388
        # squeeze dim 0 and split into separate tensors for each image
1389
        return torch.split(out.squeeze(0), embed_sizes)
1390
1391
1392

    # (TODO) Add prefix argument for filtering out weights to be loaded
    #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
1393
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1394
1395
1396
1397
1398
1399
1400
1401
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
        ]
1402
        params_dict = dict(self.named_parameters())
1403
        loaded_params: set[str] = set()
1404
        layer_count = len(self.transformer.layers)
1405
1406

        for name, loaded_weight in weights:
1407
1408
1409
1410
1411
1412
            # omit layers when num_hidden_layers_override is set
            if name.startswith("transformer.layers"):
                layer_idx = int(name.split(".")[2])
                if layer_idx >= layer_count:
                    continue

1413
            for param_name, weight_name, shard_id in stacked_params_mapping:
1414
1415
                if weight_name not in name:
                    continue
1416
1417
                name = name.replace(weight_name, param_name)
                param = params_dict[name]
1418
1419
1420
1421
1422
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                param = params_dict[name]
1423
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
1424
                weight_loader(param, loaded_weight)
1425
1426
            loaded_params.add(name)
        return loaded_params