"vllm/vscode:/vscode.git/clone" did not exist on "58631d7c3f9984f979e3cd5abf20a61590b36b1f"
pixtral.py 40.6 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import math
Patrick von Platen's avatar
Patrick von Platen committed
4
from dataclasses import dataclass, fields
5
from functools import cached_property
6
from typing import Iterable, List, Mapping, Optional, Set, Tuple, Union
Patrick von Platen's avatar
Patrick von Platen committed
7
8
9
10
11
12

import torch
import torch.nn as nn
import torch.nn.functional as F
from mistral_common.protocol.instruct.messages import ImageChunk
from PIL import Image
13
from transformers import PixtralVisionConfig
14
from transformers.models.pixtral.image_processing_pixtral import (
15
    _num_image_tokens as _get_pixtral_hf_num_image_tokens)
16
from transformers.models.pixtral.modeling_pixtral import (
17
    PixtralRotaryEmbedding, apply_rotary_pos_emb, position_ids_in_meshgrid)
Patrick von Platen's avatar
Patrick von Platen committed
18
19

from vllm.attention import AttentionMetadata
20
from vllm.config import VllmConfig
21
from vllm.distributed import divide, get_tensor_model_parallel_world_size
22
23
from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,
                         InputContext, token_inputs)
24
from vllm.model_executor.layers.activation import get_act_and_mul_fn
Patrick von Platen's avatar
Patrick von Platen committed
25
from vllm.model_executor.layers.layernorm import RMSNorm
26
27
28
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
Patrick von Platen's avatar
Patrick von Platen committed
29
from vllm.model_executor.layers.quantization import QuantizationConfig
Joe Runde's avatar
Joe Runde committed
30
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
Patrick von Platen's avatar
Patrick von Platen committed
31
32
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.sampling_metadata import SamplingMetadata
33
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs
34
from vllm.multimodal.inputs import NestedTensors, PlaceholderRange
35
from vllm.multimodal.utils import (cached_get_tokenizer,
36
                                   consecutive_placeholder_ranges)
37
from vllm.sequence import IntermediateTensors, SequenceData
Patrick von Platen's avatar
Patrick von Platen committed
38

39
from .interfaces import SupportsMultiModal, SupportsPP
40
41
from .utils import (init_vllm_registered_model, maybe_prefix,
                    merge_multimodal_embeddings)
42
from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs
Patrick von Platen's avatar
Patrick von Platen committed
43

44
45
46
47
48
49
try:
    from xformers import ops as xops
    USE_XFORMERS_OPS = True
except ImportError:
    USE_XFORMERS_OPS = False

Patrick von Platen's avatar
Patrick von Platen committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

def get_max_pixtral_image_tokens(ctx: InputContext):
    tokenizer = cached_get_tokenizer(
        ctx.model_config.tokenizer,
        tokenizer_mode=ctx.model_config.tokenizer_mode)
    mm_encoder = tokenizer.instruct.mm_encoder

    max_image_size = mm_encoder.mm_config.max_image_size
    image_patch_size = mm_encoder.mm_config.image_patch_size

    return ((max_image_size // image_patch_size)**2)


def dummy_data_for_pixtral(ctx: InputContext, seq_len: int,
                           mm_counts: Mapping[str, int]):
    tokenizer = cached_get_tokenizer(
        ctx.model_config.tokenizer,
        tokenizer_mode=ctx.model_config.tokenizer_mode)

69
70
    mm_encoder = tokenizer.mistral.instruct_tokenizer.mm_encoder
    image_token_id = mm_encoder.special_ids.img
Patrick von Platen's avatar
Patrick von Platen committed
71

72
    mm_config = ctx.get_mm_config()
73
    num_images = mm_config.limit_per_prompt.get("image", 1)
Patrick von Platen's avatar
Patrick von Platen committed
74

75
76
    # dummy size
    size = 256
Patrick von Platen's avatar
Patrick von Platen committed
77
78
    image = Image.new("RGB", (size, size), color=0)

79
80
    encoding = tokenizer.instruct.mm_encoder(ImageChunk(image=image))
    image_feature_size = len(encoding.tokens)
81
    num_image_tokens = image_feature_size * num_images
82
    seq_data = SequenceData.from_prompt_token_counts(
83
84
85
        (image_token_id, num_image_tokens),
        (0, seq_len - num_image_tokens),
    )
86
87

    mm_data = {"image": num_images * [image]}
88
89
90
91
92
93
    mm_placeholders = {
        "image":
        consecutive_placeholder_ranges(num_items=num_images,
                                       item_size=image_feature_size)
    }
    return DummyData(seq_data, mm_data, mm_placeholders)
Patrick von Platen's avatar
Patrick von Platen committed
94
95
96


def input_mapper_for_pixtral(ctx: InputContext,
97
98
                             data: object) -> MultiModalKwargs:
    """Maps the input data to its MultiModalKwargs (if any).
Patrick von Platen's avatar
Patrick von Platen committed
99
100
101

    Args:
        ctx: Context of the loaded model.
102
103
        data: data potentially containing PIL images to be processed
            and mapped to `images`.
Patrick von Platen's avatar
Patrick von Platen committed
104
105

    Returns:
106
        MultiModalKwargs containing the stacked normalized images tensor or
Patrick von Platen's avatar
Patrick von Platen committed
107
108
109
110
111
112
113
114
115
        image embeddings.
    """
    model_config = ctx.model_config
    tokenizer = cached_get_tokenizer(
        model_config.tokenizer, tokenizer_mode=model_config.tokenizer_mode)

    data_list = data if isinstance(data, list) else [data]

    images = []
116
    image_tokens_list = []
Patrick von Platen's avatar
Patrick von Platen committed
117
118
119
    for image_data in data_list:
        image = ImageChunk(image=image_data)
        encoding = tokenizer.instruct.mm_encoder(image)
120
        image = torch.from_numpy(encoding.image).to(dtype=torch.float16)
Patrick von Platen's avatar
Patrick von Platen committed
121
        images.append(image)
122
        image_tokens_list.append(encoding.tokens)
Patrick von Platen's avatar
Patrick von Platen committed
123

124
125
126
127
128
    image_tokens = torch.tensor([
        token_id for image_tokens in image_tokens_list
        for token_id in image_tokens
    ])
    return MultiModalKwargs({"images": images, "image_tokens": image_tokens})
Patrick von Platen's avatar
Patrick von Platen committed
129
130


131
132
def input_processor_for_pixtral(ctx: InputContext, inputs: DecoderOnlyInputs):
    multi_modal_data = inputs.get("multi_modal_data")
133
134
    if multi_modal_data is None or "image" not in multi_modal_data:
        return inputs
Patrick von Platen's avatar
Patrick von Platen committed
135

136
137
138
139
140
    prompt_token_ids = inputs.get("prompt_token_ids")
    prompt = inputs.get("prompt")
    tokenizer = cached_get_tokenizer(
        ctx.model_config.tokenizer,
        tokenizer_mode=ctx.model_config.tokenizer_mode)
Patrick von Platen's avatar
Patrick von Platen committed
141

142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
    mm_encoder = tokenizer.mistral.instruct_tokenizer.mm_encoder
    image_token_id = mm_encoder.special_ids.img
    image_break_id = mm_encoder.special_ids.img_break
    image_end_id = mm_encoder.special_ids.img_end

    if image_token_id not in inputs['prompt_token_ids']:
        raise ValueError(
            f"You've passed {inputs=} without {image_token_id=}"
            " Make sure to process your input via mistral_common's"
            " tokenizer or pass a chat completion request. For more"
            " For more info, see: "
            "https://github.com/vllm-project/vllm/issues/8411.")

    # Get precise tracking of placeholder positions
    placeholder_ranges = []
    curr_offset = -1
    curr_length = 0
    for i in range(len(prompt_token_ids)):
        if prompt_token_ids[i] in (image_token_id, image_break_id):
            if curr_offset < 0:
                curr_offset = i
            curr_length += 1
        elif prompt_token_ids[i] == image_end_id:
            curr_length += 1
            placeholder_ranges.append(
                PlaceholderRange(offset=curr_offset, length=curr_length))
            curr_offset = -1
            curr_length = 0
        else:
            pass
    return token_inputs(prompt=prompt,
                        prompt_token_ids=prompt_token_ids,
                        multi_modal_data=multi_modal_data,
                        multi_modal_placeholders={"image": placeholder_ranges})
Patrick von Platen's avatar
Patrick von Platen committed
176
177
178
179
180


@MULTIMODAL_REGISTRY.register_image_input_mapper(input_mapper_for_pixtral)
@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_pixtral_image_tokens)
@INPUT_REGISTRY.register_dummy_data(dummy_data_for_pixtral)
181
@INPUT_REGISTRY.register_input_processor(input_processor_for_pixtral)
182
183
class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,
                                      SupportsPP):
Patrick von Platen's avatar
Patrick von Platen committed
184

185
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Patrick von Platen's avatar
Patrick von Platen committed
186
        super().__init__()
187
188
        config = vllm_config.model_config.hf_config
        multimodal_config = vllm_config.model_config.multimodal_config
Patrick von Platen's avatar
Patrick von Platen committed
189
190
191
192
193
194
195
196
197
198
        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
        }

199
200
201
202
203
204
205
        if not ("image_break_token_id" in vision_args
                and "image_end_token_id" in vision_args):
            raise ValueError(
                "'image_break_token_id' and 'image_end_token_id' not found "
                "in the vision_encoder arguments. Please download the latest "
                "version of 'params.json' from the model repository.")

Patrick von Platen's avatar
Patrick von Platen committed
206
207
208
209
        self.vision_args = VisionEncoderArgs(**vision_args)

        # init MistralForCausalLM
        self.language_model = init_vllm_registered_model(
210
            vllm_config=vllm_config,
211
212
213
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
        )
Patrick von Platen's avatar
Patrick von Platen committed
214
215
216
217
218

        self.vision_encoder = VisionTransformer(self.vision_args)
        self.vision_language_adapter = VisionLanguageAdapter(
            self.vision_args, dim=config.text_config.hidden_size)

219
220
221
222
223
224
225
226
        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

    @cached_property
    def sampler(self):
        if hasattr(self.language_model, "sampler"):
            return self.language_model.sampler

Joe Runde's avatar
Joe Runde committed
227
        return get_sampler()
228

229
    def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:
230
231
        image_input, image_tokens = self._parse_and_validate_image_input(
            **kwargs)
232
233
        if image_input is None:
            return None
234

235
        vision_embeddings = self._process_image_input(image_input)
236
237
238
239
240
241
242
243
244

        # NOTE: We patch the outputs of the vision encoder with embeddings
        # from `[IMG_BREAK]` and `[IMG_END]` tokens.
        image_embeds = self.language_model.get_input_embeddings(image_tokens)
        image_token_mask = image_tokens == self.vision_args.image_token_id
        image_embeds[image_token_mask] = vision_embeddings

        # NOTE: Image embeddings are split into separate tensors for each image
        # by the indices of `[IMG_END]` token.
245
246
        image_end_mask = image_tokens == self.vision_args.image_end_token_id
        split_indices = torch.where(image_end_mask)[0] + 1
247
248
249
250
        if len(split_indices) <= 1:
            # Do not split, return as tensor of shape [1, fs, hs]
            return image_embeds.unsqueeze(0)

251
252
253
254
255
        # If the last split index is the last index in image_tokens, we
        # ignore it to avoid empty split tensor
        if split_indices[-1] == len(image_tokens):
            split_indices = split_indices[:-1]

256
257
        image_embeds = image_embeds.tensor_split(split_indices.cpu())
        return image_embeds
258
259
260
261
262
263
264
265
266

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: Optional[NestedTensors] = None,
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
        if multimodal_embeddings is not None:
            inputs_embeds = merge_multimodal_embeddings(
267
                input_ids, inputs_embeds, multimodal_embeddings, [
268
                    self.vision_args.image_token_id,
269
270
                    self.vision_args.image_break_token_id,
                    self.vision_args.image_end_token_id,
271
                ])
272
273
        return inputs_embeds

Patrick von Platen's avatar
Patrick von Platen committed
274
275
276
277
278
279
280
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
281
        inputs_embeds: Optional[torch.Tensor] = None,
Patrick von Platen's avatar
Patrick von Platen committed
282
        **kwargs: object,
283
    ) -> Union[torch.Tensor, IntermediateTensors]:
Patrick von Platen's avatar
Patrick von Platen committed
284
285
        """Run forward pass for pixtral.
        """
286
287
        if intermediate_tensors is not None:
            inputs_embeds = None
Patrick von Platen's avatar
Patrick von Platen committed
288

289
290
291
292
293
294
295
        # NOTE: In v1, inputs_embeds is always generated at model runner, this
        # condition is for v0 compatibility.
        elif inputs_embeds is None:
            vision_embeddings = self.get_multimodal_embeddings(**kwargs)
            inputs_embeds = self.get_input_embeddings(input_ids,
                                                      vision_embeddings)
            input_ids = None
Patrick von Platen's avatar
Patrick von Platen committed
296
297
298
299
300

        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
301
                                                  intermediate_tensors,
Patrick von Platen's avatar
Patrick von Platen committed
302
303
304
305
306
307
308
                                                  inputs_embeds=inputs_embeds)

        return hidden_states

    def _parse_and_validate_image_input(
        self,
        images: Optional[Union[List[List[torch.Tensor]], List[torch.Tensor],
309
310
                               torch.Tensor]] = None,
        image_tokens: Optional[torch.Tensor] = None,
311
    ) -> Tuple[Optional[List[torch.Tensor]], Optional[torch.Tensor]]:
Patrick von Platen's avatar
Patrick von Platen committed
312
        if images is None:
313
            return None, None
Patrick von Platen's avatar
Patrick von Platen committed
314
315

        if isinstance(images, torch.Tensor):
316
317
318
319
            # if passed as batch take all images
            N, B, C, W, H = images.shape
            images = images.reshape(N * B, C, W, H)
            images = [images[i] for i in range(images.size(0))]
Patrick von Platen's avatar
Patrick von Platen committed
320
        elif isinstance(images, list):
321
322
323
324
325
326
327
328
329
330
            # if passed as list flatten lists of tensors
            flatten_images = []
            for imgs_per_req in images:
                imgs_per_req = [
                    imgs_per_req[i] for i in range(imgs_per_req.size(0))
                ] if isinstance(imgs_per_req, torch.Tensor) else imgs_per_req

                flatten_images.extend(imgs_per_req)

            images = flatten_images
Patrick von Platen's avatar
Patrick von Platen committed
331

332
333
334
335
336
337
338
339
340
341
        if isinstance(image_tokens, torch.Tensor):
            # image_tokens are batched
            image_tokens = image_tokens.flatten()
        elif isinstance(image_tokens, list):
            # image_tokens are of different lengths thus passed as a list
            image_tokens = torch.cat(image_tokens)

        assert image_tokens.dim() == 1

        return images, image_tokens
Patrick von Platen's avatar
Patrick von Platen committed
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369

    def _process_image_input(self,
                             image_input: List[torch.Tensor]) -> torch.Tensor:
        return self.vision_language_adapter(self.vision_encoder(image_input))

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)

    def sample(
        self,
        logits: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        return self.language_model.sample(logits, sampling_metadata)

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):

        def is_vision_encoder_weights(weight: Tuple[str, torch.Tensor]):
            return weight[0].startswith("vision_encoder")

        def is_vision_lang_adapter_weights(weight: Tuple[str, torch.Tensor]):
            return weight[0].startswith("vision_language_adapter")

370
        # Get references to parameters for direct loading
Patrick von Platen's avatar
Patrick von Platen committed
371
        vision_encoder_dict = dict(self.vision_encoder.named_parameters())
372
        vision_lang_adapter_dict = dict(
Patrick von Platen's avatar
Patrick von Platen committed
373
            self.vision_language_adapter.named_parameters())
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396

        def llm_weights_generator():
            # Single pass over weights
            for name, w in weights:
                if is_vision_encoder_weights((name, w)):
                    # Load vision encoder weights directly
                    trimmed_name = '.'.join(name.split(".")[1:])
                    param = vision_encoder_dict[trimmed_name]
                    with torch.no_grad():
                        default_weight_loader(param, w)
                elif is_vision_lang_adapter_weights((name, w)):
                    # Load vision-language adapter weights directly
                    trimmed_name = '.'.join(name.split(".")[1:])
                    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
397
398
399
400
401
402
403
404
405
406
407
408
409
410


# 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
411
412
    image_break_token_id: int
    image_end_token_id: int
413
    adapter_bias: bool = True
Patrick von Platen's avatar
Patrick von Platen committed
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511


def _reshape_for_broadcast(freqs_cis: torch.Tensor,
                           x: torch.Tensor) -> torch.Tensor:
    """
    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]),
    )
    shape = [
        d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)
    ]
    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
    freqs = 1.0 / (theta**(torch.arange(0, dim, 2).float() / dim))

    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,
) -> Tuple[torch.Tensor, torch.Tensor]:
    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
        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)

    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,
512
        mask: torch.Tensor,
Patrick von Platen's avatar
Patrick von Platen committed
513
514
515
516
517
518
519
520
521
522
        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)
523
        out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)
Patrick von Platen's avatar
Patrick von Platen committed
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
        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,
540
        mask: torch.Tensor,
Patrick von Platen's avatar
Patrick von Platen committed
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
        freqs_cis: torch.Tensor,
    ) -> torch.Tensor:
        r = self.attention.forward(self.attention_norm(x),
                                   mask=mask,
                                   freqs_cis=freqs_cis)
        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,
563
        mask: torch.Tensor,
Patrick von Platen's avatar
Patrick von Platen committed
564
565
566
567
568
569
570
        freqs_cis: Optional[torch.Tensor],
    ) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, mask=mask, freqs_cis=freqs_cis)
        return x


571
def position_meshgrid(patch_embeds_list: List[torch.Tensor], ) -> torch.Tensor:
Patrick von Platen's avatar
Patrick von Platen committed
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
    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
    ])
    return positions


class VisionTransformer(nn.Module):

    def __init__(self, args: VisionEncoderArgs):
        super().__init__()
        self.args = args
        self.patch_conv = nn.Conv2d(
            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"
        self._freqs_cis: Optional[torch.Tensor] = None

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

    @property
609
    def device(self) -> torch.types.Device:
Patrick von Platen's avatar
Patrick von Platen committed
610
611
612
        return next(self.parameters()).device

    @property
613
    def dtype(self) -> torch.dtype:
Patrick von Platen's avatar
Patrick von Platen committed
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
        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,
        images: List[torch.Tensor],
    ) -> torch.Tensor:
        """
        Args:
            images: list of N_img images of variable sizes, 
                each of shape (C, H, W)
        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 = [
            self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in images
        ]

        # flatten to a single sequence
        patch_embeds = torch.cat(
            [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list], dim=1)
        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
658
659
660
661
662
663
        if USE_XFORMERS_OPS:
            mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(
                [p.shape[-2] * p.shape[-1] for p in patch_embeds_list], )
        else:
            raise ImportError("Xformers is required for Pixtral inference "
                              "with the Mistral format")
Patrick von Platen's avatar
Patrick von Platen committed
664
665
666
667
668
669
670
671
672
673
674
675
676
677
        out = self.transformer(patch_embeds, mask=mask, freqs_cis=freqs_cis)

        # remove batch dimension of the single sequence
        return out.squeeze(0)


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,
678
            bias=args.adapter_bias,
Patrick von Platen's avatar
Patrick von Platen committed
679
680
        )
        self.gelu = nn.GELU()
681
        self.w_out = nn.Linear(dim, dim, bias=args.adapter_bias)
Patrick von Platen's avatar
Patrick von Platen committed
682
683
684

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_out(self.gelu(self.w_in(x)))
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701


#### 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.


def get_pixtral_hf_patch_grid_length(*, image_size: int,
                                     patch_size: int) -> int:
    # Since interpolation is applied, the image size need not be divisible
    # assert image_size % patch_size == 0
    return image_size // patch_size


702
703
704
705
706
707
708
709
710
711
712
713
def get_pixtral_hf_image_feature_size(
    *,
    image_size: int,
    patch_size: int,
) -> int:
    grid_length = get_pixtral_hf_patch_grid_length(
        image_size=image_size,
        patch_size=patch_size,
    )

    # Consider the image_break_token
    return (grid_length + 1) * grid_length
714
715
716


def get_max_pixtral_hf_image_tokens(hf_config: PixtralVisionConfig) -> int:
717
718
719
720
721
722
723
    grid_length = get_pixtral_hf_patch_grid_length(
        image_size=hf_config.image_size,
        patch_size=hf_config.patch_size,
    )

    # Consider the image_break_token
    return (grid_length + 1) * grid_length
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742


def dummy_image_for_pixtral_hf(
    hf_config: PixtralVisionConfig,
    num_images: int,
    *,
    image_width_override: Optional[int] = None,
    image_height_override: Optional[int] = None,
):
    width = height = hf_config.image_size
    if image_width_override is not None:
        width = image_width_override
    if image_height_override is not None:
        height = image_height_override

    image = Image.new("RGB", (width, height), color=0)
    return {"image": image if num_images == 1 else [image] * num_images}


743
744
745
746
747
748
749
750
751
752
# Adapted from transformers.models.pixtral.image_processing_pixtral.get_resize_output_image_size # noqa: E501
# https://github.com/huggingface/transformers/blob/2bd4d5897dc73e8b172832070a6f9e567a0df017/src/transformers/models/pixtral/image_processing_pixtral.py#L180
def get_pixtral_hf_image_feature_grid_size(
    hf_config: PixtralVisionConfig,
    *,
    image_width: int,
    image_height: int,
) -> tuple[int, int]:
    max_width = max_height = hf_config.image_size
    patch_width = patch_height = hf_config.patch_size
753
754
755
756

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

    if ratio > 1:
757
758
        image_width = int(math.ceil(image_width / ratio))
        image_height = int(math.ceil(image_height / ratio))
759

760
    nrows, ncols = _get_pixtral_hf_num_image_tokens(
761
762
        (image_height, image_width),
        (patch_height, patch_width),
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
    )  # type: ignore

    return ncols, nrows


class PixtralHFEncoderInfo(VisionEncoderInfo[PixtralVisionConfig]):

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
        return get_pixtral_hf_image_feature_size(
            image_size=self.vision_config.image_size,
778
            patch_size=self.vision_config.patch_size,
779
780
781
782
783
        )

    def get_max_image_tokens(self) -> int:
        return get_max_pixtral_hf_image_tokens(self.vision_config)

784
785
786
787
788
789
790
    def get_image_size(self) -> int:
        return self.vision_config.image_size

    def get_patch_size(self) -> int:
        return self.vision_config.patch_size

    def get_patch_grid_length(self) -> int:
791
792
793
794
        return get_pixtral_hf_patch_grid_length(
            image_size=self.vision_config.image_size,
            patch_size=self.vision_config.patch_size,
        )
795
796
797
798


class PixtralHFMLP(nn.Module):

799
800
801
802
803
804
805
    def __init__(
        self,
        config: PixtralVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        *,
        prefix: str = "",
    ) -> None:
806
        super().__init__()
807

808
        assert config.intermediate_size is not None
809
810
811
812
813
814
815
816
817
818
819
820
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=config.hidden_size,
            output_sizes=[config.intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj")
        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")
        self.act_and_mul = get_act_and_mul_fn(config.hidden_act)
821
822

    def forward(self, x: torch.Tensor) -> torch.Tensor:
823
824
825
826
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_and_mul(gate_up)
        x, _ = self.down_proj(x)
        return x
827
828
829
830


class PixtralHFAttention(nn.Module):

831
832
833
834
835
836
837
    def __init__(
        self,
        config: PixtralVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        *,
        prefix: str = "",
    ) -> None:
838
        super().__init__()
839

840
841
        self.config = config
        assert not config.hidden_size % config.num_attention_heads
842
843
844
        self.total_num_heads = config.num_attention_heads
        tp_size = get_tensor_model_parallel_world_size()
        self.n_heads = divide(config.num_attention_heads, tp_size)
845
846
        self.head_dim = config.hidden_size // config.num_attention_heads

847
848
849
        self.qkv_proj = QKVParallelLinear(
            hidden_size=config.hidden_size,
            head_size=self.head_dim,
850
            total_num_heads=self.total_num_heads,
851
852
853
854
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
855
        assert self.total_num_heads * self.head_dim == config.hidden_size
856
857
858
859
860
861
862
        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",
        )
863
864
865
866

    def forward(
        self,
        hidden_states: torch.Tensor,
867
        attention_mask: torch.Tensor,
868
869
        position_embeddings: torch.Tensor,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
870
        batch, patches, _ = hidden_states.size()
871

872
873
        qkv_states, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv_states.chunk(3, dim=-1)
874

875
876
877
        # 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)
878
        v = v.view(batch, patches, self.n_heads, self.head_dim)
879
        cos, sin = position_embeddings
880
        q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=0)
881

882
883
884
885
886
887
888
889
890
891
        if USE_XFORMERS_OPS:
            # Transpose q and k back for attention
            q = q.transpose(1, 2).contiguous()
            k = k.transpose(1, 2).contiguous()

            out = xops.memory_efficient_attention(q,
                                                  k,
                                                  v,
                                                  attn_bias=attention_mask)
        else:
892
            v = v.transpose(1, 2)
893
894
895
            out = nn.functional.scaled_dot_product_attention(
                q, k, v, attn_mask=attention_mask)
            out = out.transpose(1, 2)
896

897
898
        out = out.view(batch, patches, self.n_heads * self.head_dim)
        attn_output, _ = self.o_proj(out)
899

900
        return attn_output, None
901
902
903
904


class PixtralHFTransformerBlock(nn.Module):

905
906
907
908
909
910
911
    def __init__(
        self,
        config: PixtralVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        *,
        prefix: str = "",
    ) -> None:
912
        super().__init__()
913

914
        self.attention_norm = RMSNorm(config.hidden_size, eps=1e-5)
915
916
917
918
919
920
        self.attention = PixtralHFAttention(config,
                                            quant_config=quant_config,
                                            prefix=f"{prefix}.attention")
        self.feed_forward = PixtralHFMLP(config,
                                         quant_config=quant_config,
                                         prefix=f"{prefix}.feed_forward")
921
922
923
924
925
        self.ffn_norm = RMSNorm(config.hidden_size, eps=1e-5)

    def forward(
        self,
        hidden_states: torch.Tensor,
926
        attention_mask: torch.Tensor,
927
928
        position_embeddings: torch.Tensor,
    ) -> torch.Tensor:
929
930
931
        r, _ = self.attention.forward(self.attention_norm(hidden_states),
                                      attention_mask=attention_mask,
                                      position_embeddings=position_embeddings)
932
933
934
935
936
937
938
939
        h = hidden_states + r
        r = self.feed_forward.forward(self.ffn_norm(h))
        out = h + r
        return out


class PixtralHFTransformer(nn.Module):

940
941
942
943
944
945
946
947
    def __init__(
        self,
        config: PixtralVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        *,
        num_hidden_layers_override: Optional[int] = None,
        prefix: str = "",
    ) -> None:
948
        super().__init__()
949
950
951
952
953
954
955
956
957
958
959
960

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

        self.layers = nn.ModuleList([
            PixtralHFTransformerBlock(config=config,
                                      quant_config=quant_config,
                                      prefix=f"{prefix}.layers.{layer_idx}")
            for layer_idx in range(num_hidden_layers)
        ])
961
962
963
964

    def forward(
        self,
        x: torch.Tensor,
965
        attention_mask: torch.Tensor,
966
        position_embeddings: torch.Tensor,
967
        return_all_hidden_states: bool,
968
    ) -> torch.Tensor:
969
970
        hidden_states_pool = []

971
972
        for layer in self.layers:
            x = layer(x, attention_mask, position_embeddings)
973
974
975
976
977
978
            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
979
980
981
982
983
        return x


class PixtralHFVisionModel(nn.Module):

984
985
986
987
988
989
990
991
992
    def __init__(
        self,
        config: PixtralVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        *,
        num_hidden_layers_override: Optional[int] = None,
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
    ) -> None:
993
994
995
        super().__init__()

        self.config = config
996

997
998
999
1000
1001
1002
1003
1004
        self.patch_conv = nn.Conv2d(
            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)
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
        self.transformer = PixtralHFTransformer(
            config,
            quant_config,
            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)} "
                "layers.")

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

1023
1024
1025
1026
1027
1028
1029
1030
        self.dtype = next(self.parameters()).dtype
        self.device = next(self.parameters()).device
        self.patch_positional_embedding = PixtralRotaryEmbedding(
            config, self.device)

    def forward(
        self,
        pixel_values: List[torch.Tensor],
1031
        feature_sample_layers: Optional[list[int]] = None,
1032
1033
1034
    ) -> torch.Tensor:
        """
        Args:
1035
1036
1037
1038
            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
1039
1040
1041
            feature_sample_layers: Layer indices whose features should be
                concatenated and used as the visual encoder output. If none
                are provided, the last layer is used.
1042

1043
1044
1045
1046
1047
1048
        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 = [
1049
            self.patch_conv(img.unsqueeze(0).to(self.dtype))
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
            for img in pixel_values
        ]

        # flatten to a single sequence
        patch_embeds = torch.cat(
            [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list], dim=1)
        patch_embeds = self.ln_pre(patch_embeds)

        # positional embeddings
        position_ids = position_ids_in_meshgrid(
            patch_embeds_list,
            max_width=self.config.image_size // self.config.patch_size).to(
                self.device)
        position_embedding = self.patch_positional_embedding(
            patch_embeds, position_ids)
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075

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

1076
1077
1078
1079
1080
1081
1082
1083
1084
        return_all_hidden_states = feature_sample_layers is not None
        out = self.transformer(
            patch_embeds,
            attention_mask,
            position_embedding,
            return_all_hidden_states=return_all_hidden_states)

        out = resolve_visual_encoder_outputs(out, feature_sample_layers, None,
                                             self.config.num_hidden_layers)
1085
1086
1087
1088
1089

        return out

    # (TODO) Add prefix argument for filtering out weights to be loaded
    #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
1090
1091
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
1092
1093
1094
1095
1096
1097
1098
1099
        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),
        ]
1100
        params_dict = dict(self.named_parameters())
1101
        loaded_params: Set[str] = set()
1102
        layer_count = len(self.transformer.layers)
1103
1104

        for name, loaded_weight in weights:
1105
1106
1107
1108
1109
1110
            # 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

1111
1112
1113
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
1114
1115
                name = name.replace(weight_name, param_name)
                param = params_dict[name]
1116
1117
1118
1119
1120
1121
1122
1123
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
1124
1125
            loaded_params.add(name)
        return loaded_params