qwen2_vl.py 48.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# Adapted from
# https://github.com/huggingface/transformers/blob/19e6e80e10118f855137b90740936c0b11ac397f/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py
# Copyright 2024 The Qwen team.
# Copyright 2023 The vLLM team.
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Inference-only Qwen2-VL model compatible with HuggingFace weights."""
24
from functools import partial
25
26
from typing import (Any, Callable, Dict, Iterable, List, Literal, Mapping,
                    Optional, Tuple, Type, TypedDict, Union)
27
28
29
30
31
32
33
34
35

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from PIL import Image
from transformers.image_utils import (get_image_size,
                                      infer_channel_dimension_format,
                                      to_numpy_array)
36
37
from transformers.models.qwen2_vl.configuration_qwen2_vl import (
    Qwen2VLConfig, Qwen2VLVisionConfig)
38
39
40
41
from transformers.models.qwen2_vl.image_processing_qwen2_vl import (
    make_batched_images, make_batched_videos, smart_resize)

from vllm.attention import AttentionMetadata
42
from vllm.attention.selector import _Backend
43
from vllm.config import CacheConfig, LoRAConfig, MultiModalConfig
44
from vllm.distributed import get_pp_group, parallel_state
45
from vllm.distributed import utils as dist_utils
46
47
from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,
                         InputContext, token_inputs)
48
49
50
51
52
53
54
from vllm.logger import init_logger
from vllm.model_executor import SamplingMetadata
from vllm.model_executor.layers.activation import QuickGELU
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
Joe Runde's avatar
Joe Runde committed
55
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
56
57
58
59
60
61
62
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.qwen2 import Qwen2Model
from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalDataDict,
                             MultiModalInputs)
from vllm.multimodal.base import MultiModalData
from vllm.multimodal.image import cached_get_image_processor
63
from vllm.multimodal.utils import cached_get_tokenizer
64
from vllm.sequence import IntermediateTensors, SequenceData
65
from vllm.transformers_utils.config import uses_mrope
66
from vllm.transformers_utils.processor import cached_get_processor
67

68
from .interfaces import SupportsLoRA, SupportsMultiModal, SupportsPP
69
70
from .utils import (PPMissingLayer, get_vit_attn_backend,
                    is_pp_missing_parameter,
71
72
                    make_empty_intermediate_tensors_factory)

73
74
75
76
77
logger = init_logger(__name__)

# === Vision Inputs === #


78
79
80
class Qwen2VLImagePixelInputs(TypedDict):
    type: Literal["pixel_values"]
    data: torch.Tensor
81
    """Shape:
82
83
84
85
86
87
88
89
90
    `(num_patches, num_channels * patch_size * patch_size)`
    """

    image_grid_thw: torch.Tensor
    """Shape: `(num_images, 3)`
    This should be in `(grid_t, grid_h, grid_w)` format.
    """


91
92
93
94
95
96
97
98
99
100
101
102
class Qwen2VLImageEmbeddingInputs(TypedDict):
    type: Literal["image_embeds"]
    data: torch.Tensor
    """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
    `hidden_size` must match the hidden size of language model backbone.
    """


Qwen2VLImageInputs = Union[Qwen2VLImagePixelInputs,
                           Qwen2VLImageEmbeddingInputs]


103
104
class Qwen2VLVideoInputs(TypedDict):
    pixel_values_videos: torch.Tensor
105
106
    """Shape:
    `(num_patches,
107
108
109
110
111
      num_channels * temporal_patch_size * patch_size * patch_size)`
    """

    video_grid_thw: torch.Tensor
    """Shape: `(num_videos, 3)`
112

113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
    This should be in `(grid_t, grid_h, grid_w)` format.
    """


# === Vision Encoder === #


class Qwen2VisionMLP(nn.Module):

    def __init__(
        self,
        in_features: int,
        hidden_features: int = None,
        act_layer: Type[nn.Module] = QuickGELU,
        quant_config: Optional[QuantizationConfig] = None,
128
        prefix: str = "",
129
130
131
132
    ):
        super().__init__()
        self.fc1 = ColumnParallelLinear(in_features,
                                        hidden_features,
133
134
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.fc1")
135
136
137
        self.act = act_layer()
        self.fc2 = RowParallelLinear(hidden_features,
                                     in_features,
138
139
                                     quant_config=quant_config,
                                     prefix=f"{prefix}.fc2")
140
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
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_parallel, _ = self.fc1(x)
        x_parallel = self.act(x_parallel)
        x, _ = self.fc2(x_parallel)
        return x


def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
    if not interleaved:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)
    else:
        x1, x2 = x[..., ::2], x[..., 1::2]
        return rearrange(torch.stack((-x2, x1), dim=-1),
                         "... d two -> ... (d two)",
                         two=2)


def apply_rotary_emb_torch(x: torch.Tensor,
                           cos: torch.Tensor,
                           sin: torch.Tensor,
                           interleaved: bool = False) -> torch.Tensor:
    """
    x: (batch_size, seqlen, nheads, headdim)
    cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
    """
    ro_dim = cos.shape[-1] * 2
    assert ro_dim <= x.shape[-1]
    cos = repeat(
        cos,
        "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
    sin = repeat(
        sin,
        "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
    return torch.cat(
        [
            x[..., :ro_dim] * cos +
            rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]
        ],
        dim=-1,
    )


def apply_rotary_pos_emb_vision(t: torch.Tensor,
                                freqs: torch.Tensor) -> torch.Tensor:
    t_ = t.float()
    cos = freqs.cos()
    sin = freqs.sin()
    output = apply_rotary_emb_torch(t_, cos, sin).type_as(t)
    return output


class Qwen2VisionAttention(nn.Module):

    def __init__(
        self,
        embed_dim: Optional[int] = None,
        num_heads: Optional[int] = None,
        projection_size: Optional[int] = None,
        quant_config: Optional[QuantizationConfig] = None,
201
        prefix: str = "",
202
203
204
205
206
207
208
209
210
211
212
    ) -> None:
        super().__init__()
        # Per attention head and per partition values.
        world_size = parallel_state.get_tensor_model_parallel_world_size()
        self.hidden_size_per_attention_head = dist_utils.divide(
            projection_size, num_heads)
        self.num_attention_heads_per_partition = dist_utils.divide(
            num_heads, world_size)

        self.qkv = ColumnParallelLinear(input_size=embed_dim,
                                        output_size=3 * projection_size,
213
214
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.qkv")
215
216
        self.proj = RowParallelLinear(input_size=projection_size,
                                      output_size=embed_dim,
217
218
                                      quant_config=quant_config,
                                      prefix=f"{prefix}.proj")
219
220

        # Detect attention implementation.
221
222
223
224
225
226
        self.attn_backend: _Backend = get_vit_attn_backend()
        if self.attn_backend not in {
                _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS
        }:
            raise RuntimeError(
                f"Qwen2-VL does not support {self.attn_backend} backend now.")
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247

    def forward(
        self,
        x: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: torch.Tensor = None,
    ) -> torch.Tensor:
        # [s, b, c] --> [s, b, head * 3 * head_dim]
        x, _ = self.qkv(x)

        # [s, b, head * 3 * head_dim] --> [s, b, head, 3 * head_dim]
        new_x_shape = x.size()[:-1] + (
            self.num_attention_heads_per_partition,
            3 * self.hidden_size_per_attention_head,
        )
        x = x.view(*new_x_shape)

        # [s, b, head, 3 * head_dim] --> 3 [s, b, head, head_dim]
        q, k, v = dist_utils.split_tensor_along_last_dim(x, 3)
        batch_size = q.shape[1]

248
249
        q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
                   for x in (q, k, v))
250
251
252
253
        if rotary_pos_emb is not None:
            q = apply_rotary_pos_emb_vision(q, rotary_pos_emb)
            k = apply_rotary_pos_emb_vision(k, rotary_pos_emb)

254
        if self.attn_backend == _Backend.FLASH_ATTN:
255
256
257
258
            # from vllm_flash_attn.flash_attn_interface import (
            #   flash_attn_varlen_func)
            from flash_attn import flash_attn_varlen_func

259
            q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274

            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
            output = flash_attn_varlen_func(q,
                                            k,
                                            v,
                                            cu_seqlens_q=cu_seqlens,
                                            cu_seqlens_k=cu_seqlens,
                                            max_seqlen_q=max_seqlen,
                                            max_seqlen_k=max_seqlen,
                                            dropout_p=0,
                                            causal=False)

            context_layer = rearrange(output,
                                      "(b s) ... -> b s ...",
                                      b=batch_size)
275
        elif self.attn_backend == _Backend.TORCH_SDPA:
276
            seq_length = q.size(1)
277
            q, k, v = (rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
278
279
280
281
282
283
284
285
286
287
288
289
            attention_mask = torch.zeros([1, seq_length, seq_length],
                                         device=q.device,
                                         dtype=torch.bool)
            for i in range(1, len(cu_seqlens)):
                attention_mask[..., cu_seqlens[i - 1]:cu_seqlens[i],
                               cu_seqlens[i - 1]:cu_seqlens[i]] = True
            output = F.scaled_dot_product_attention(q,
                                                    k,
                                                    v,
                                                    attention_mask,
                                                    dropout_p=0.0)
            context_layer = rearrange(output, "b h s d -> b s h d ")
290
        elif self.attn_backend == _Backend.XFORMERS:
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
            from xformers import ops as xops
            from xformers.ops.fmha.attn_bias import BlockDiagonalMask

            seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
            attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
                                                       kv_seqlen=None)

            context_layer = xops.memory_efficient_attention_forward(
                q, k, v, attn_bias=attn_bias, p=0, scale=None)
        context_layer = rearrange(context_layer,
                                  "b s h d -> s b (h d)").contiguous()

        output, _ = self.proj(context_layer)
        return output


class Qwen2VisionBlock(nn.Module):

    def __init__(
        self,
        dim: int,
        num_heads: int,
        mlp_ratio: float,
        act_layer: Type[nn.Module] = QuickGELU,
        norm_layer: Type[nn.Module] = None,
        quant_config: Optional[QuantizationConfig] = None,
317
        prefix: str = "",
318
319
320
321
322
323
324
325
326
327
328
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = partial(nn.LayerNorm, eps=1e-6)
        self.norm1 = norm_layer(dim)
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)

        self.attn = Qwen2VisionAttention(embed_dim=dim,
                                         num_heads=num_heads,
                                         projection_size=dim,
329
330
                                         quant_config=quant_config,
                                         prefix=f"{prefix}.attn")
331
332
333
        self.mlp = Qwen2VisionMLP(dim,
                                  mlp_hidden_dim,
                                  act_layer=act_layer,
334
335
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.mlp")
336
337
338
339
340
341
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
370
371
372
373
374
375
376
377
378
379
380
381
382
383

    def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor,
                rotary_pos_emb: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x),
                          cu_seqlens=cu_seqlens,
                          rotary_pos_emb=rotary_pos_emb)
        x = x + self.mlp(self.norm2(x))
        return x


class Qwen2VisionPatchEmbed(nn.Module):

    def __init__(
        self,
        patch_size: int = 14,
        temporal_patch_size: int = 2,
        in_chans: int = 3,
        embed_dim: int = 1152,
    ) -> None:
        super().__init__()
        self.patch_size = patch_size
        self.temporal_patch_size = temporal_patch_size
        self.embed_dim = embed_dim

        kernel_size = [temporal_patch_size, patch_size, patch_size]
        self.proj = nn.Conv3d(in_chans,
                              embed_dim,
                              kernel_size=kernel_size,
                              stride=kernel_size,
                              bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        L, C = x.shape
        x = x.view(L, -1, self.temporal_patch_size, self.patch_size,
                   self.patch_size)
        x = self.proj(x).view(L, self.embed_dim)
        return x


class Qwen2VisionPatchMerger(nn.Module):

    def __init__(
        self,
        d_model: int,
        context_dim: int,
        norm_layer: Type[nn.Module] = None,
        spatial_merge_size: int = 2,
        quant_config: Optional[QuantizationConfig] = None,
384
        prefix: str = "",
385
386
387
388
389
390
391
392
393
394
    ) -> None:
        super().__init__()
        self.hidden_size = context_dim * (spatial_merge_size**2)
        if norm_layer is None:
            norm_layer = partial(nn.LayerNorm, eps=1e-6)
        self.ln_q = norm_layer(context_dim)
        self.mlp = nn.ModuleList([
            ColumnParallelLinear(self.hidden_size,
                                 self.hidden_size,
                                 bias=True,
395
396
                                 quant_config=quant_config,
                                 prefix=f"{prefix}.mlp.0"),
397
398
399
400
            nn.GELU(),
            RowParallelLinear(self.hidden_size,
                              d_model,
                              bias=True,
401
402
                              quant_config=quant_config,
                              prefix=f"{prefix}.mlp.2"),
403
404
405
406
407
408
409
410
411
412
413
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
        ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.ln_q(x)
        x = x.view(-1, self.hidden_size)

        mlp_fc1, mlp_act, mlp_fc2 = self.mlp
        x_parallel, _ = mlp_fc1(x)
        x_parallel = mlp_act(x_parallel)
        out, _ = mlp_fc2(x_parallel)
        return out


class Qwen2VisionRotaryEmbedding(nn.Module):

    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        self.dim = dim
        self.theta = theta
        inv_freq = 1.0 / (theta
                          **(torch.arange(0, dim, 2, dtype=torch.float) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._seq_len_cached = 0
        self._freqs_cached = None

    def update_freqs_cache(self, seqlen: int) -> None:
        if seqlen > self._seq_len_cached:
            seqlen *= 2
            self._seq_len_cached = seqlen
            self.inv_freq = 1.0 / (self.theta**(torch.arange(
                0, self.dim, 2, dtype=torch.float, device=self.inv_freq.device)
                                                / self.dim))
            seq = torch.arange(seqlen,
                               device=self.inv_freq.device,
                               dtype=self.inv_freq.dtype)
            freqs = torch.outer(seq, self.inv_freq)
            self._freqs_cached = freqs

    def forward(self, seqlen: int) -> torch.Tensor:
        self.update_freqs_cache(seqlen)
        return self._freqs_cached[:seqlen]


class Qwen2VisionTransformer(nn.Module):

    def __init__(
        self,
        vision_config: Qwen2VLVisionConfig,
        norm_eps: float = 1e-6,
        quant_config: Optional[QuantizationConfig] = None,
453
        prefix: str = "",
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
    ) -> None:
        super().__init__()

        patch_size: int = vision_config.patch_size
        temporal_patch_size: int = vision_config.temporal_patch_size
        spatial_merge_size: int = vision_config.spatial_merge_size
        in_chans: int = vision_config.in_chans
        hidden_size: int = vision_config.hidden_size
        embed_dim: int = vision_config.embed_dim
        depth: int = vision_config.depth
        num_heads: int = vision_config.num_heads
        mlp_ratio: float = vision_config.mlp_ratio

        self.spatial_merge_size = spatial_merge_size

        self.patch_embed = Qwen2VisionPatchEmbed(
            patch_size=patch_size,
            temporal_patch_size=temporal_patch_size,
            in_chans=in_chans,
            embed_dim=embed_dim,
        )

        norm_layer = partial(nn.LayerNorm, eps=norm_eps)
        head_dim = embed_dim // num_heads
        self.rotary_pos_emb = Qwen2VisionRotaryEmbedding(head_dim // 2)

        self.blocks = nn.ModuleList([
481
482
483
484
485
486
487
            Qwen2VisionBlock(dim=embed_dim,
                             num_heads=num_heads,
                             mlp_ratio=mlp_ratio,
                             norm_layer=norm_layer,
                             quant_config=quant_config,
                             prefix=f"{prefix}.blocks.{layer_idx}")
            for layer_idx in range(depth)
488
489
490
491
492
493
        ])
        self.merger = Qwen2VisionPatchMerger(
            d_model=hidden_size,
            context_dim=embed_dim,
            norm_layer=norm_layer,
            quant_config=quant_config,
494
            prefix=f"{prefix}.merger",
495
496
497
498
        )

    @property
    def dtype(self) -> torch.dtype:
499
        return self.patch_embed.proj.weight.dtype
500
501
502

    @property
    def device(self) -> torch.device:
503
        return self.patch_embed.proj.weight.device
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560

    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
        pos_ids = []
        for t, h, w in grid_thw:
            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
            hpos_ids = hpos_ids.reshape(
                h // self.spatial_merge_size,
                self.spatial_merge_size,
                w // self.spatial_merge_size,
                self.spatial_merge_size,
            ).permute(0, 2, 1, 3).flatten()
            wpos_ids = wpos_ids.reshape(
                h // self.spatial_merge_size,
                self.spatial_merge_size,
                w // self.spatial_merge_size,
                self.spatial_merge_size,
            ).permute(0, 2, 1, 3).flatten()
            pos_ids.append(
                torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
        pos_ids = torch.cat(pos_ids, dim=0)
        max_grid_size = grid_thw[:, 1:].max()
        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
        rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
        return rotary_pos_emb

    def forward(
        self,
        x: torch.Tensor,
        grid_thw: torch.Tensor,
    ) -> torch.Tensor:
        # patchify
        x = x.to(device=self.device, dtype=self.dtype)
        x = self.patch_embed(x)

        # compute position embedding
        rotary_pos_emb = self.rot_pos_emb(grid_thw)

        # compute cu_seqlens
        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2],
                                             grid_thw[:, 0]).cumsum(
                                                 dim=0, dtype=torch.int32)
        cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)

        # transformers
        x = x.unsqueeze(1)
        for blk in self.blocks:
            x = blk(x, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb)

        # adapter
        x = self.merger(x)
        return x


# === Vision input helpers === #


561
562
563
564
565
566
567
568
569
570
571
def get_mm_processor_kwargs(
        min_pixels: Optional[int] = None,
        max_pixels: Optional[int] = None) -> Dict[str, int]:
    mm_processor_kwargs = {}
    if min_pixels:
        mm_processor_kwargs["min_pixels"] = min_pixels
    if max_pixels:
        mm_processor_kwargs["max_pixels"] = max_pixels
    return mm_processor_kwargs


572
573
574
575
def mm_input_mapper_for_qwen2_vl(
    ctx: InputContext,
    data: MultiModalData[object],
    data_type_key: str,
576
577
578
    *,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
579
580
) -> MultiModalInputs:
    """Input mapper for Qwen2-VL."""
581
582
583
584
585
    if data_type_key == "image" and isinstance(data, dict):
        return MultiModalInputs({
            "image_embeds": data.get("image_embeds"),
            "image_grid_thw": data.get("image_grid_thw"),
        })
586
    model_config = ctx.model_config
587
588
    # Handle mm processor kwargs; we pass these at creation time
    # because preprocess() in transformers doesn't expose them
589
590
    mm_processor_kwargs = get_mm_processor_kwargs(min_pixels=min_pixels,
                                                  max_pixels=max_pixels)
591
    image_processor = cached_get_image_processor(
592
593
594
595
        model_config.model,
        trust_remote_code=model_config.trust_remote_code,
        **mm_processor_kwargs,
    )
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
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
658
659
660
661
662
663
664
665
666
667
    if image_processor is None:
        raise RuntimeError("No HuggingFace processor is available "
                           "to process the image object")

    images = None
    videos = None
    if data_type_key == "image":
        images = data
    else:
        assert data_type_key == "video"
        videos = data

    try:
        batch_data = image_processor \
            .preprocess(images=images, videos=videos, return_tensors="pt") \
            .data
    except Exception:
        logger.error("Failed to process image (%s)", data)
        raise

    return MultiModalInputs(batch_data)


image_input_mapper_for_qwen2_vl = partial(mm_input_mapper_for_qwen2_vl,
                                          data_type_key="image")
video_input_mapper_for_qwen2_vl = partial(mm_input_mapper_for_qwen2_vl,
                                          data_type_key="video")


def _get_vision_info(
    image_processor,
    height: int,
    width: int,
    min_pixels: int,
    max_pixels: int,
    do_resize: bool = True,
    data_type_key: str = "image",
    mm_count: int = 1,
):
    """Get information (resized height / width and number of vision tokens)
    of input image / video frame."""

    if do_resize:
        resized_height, resized_width = smart_resize(
            height=height,
            width=width,
            factor=image_processor.patch_size * image_processor.merge_size,
            min_pixels=min_pixels,
            max_pixels=max_pixels,
        )
    else:
        resized_height, resized_width = height, width

    if data_type_key == "image":
        grid_t = mm_count
    else:
        assert data_type_key == "video"
        grid_t = max(mm_count // image_processor.temporal_patch_size, 1)

    grid_h = resized_height // image_processor.patch_size
    grid_w = resized_width // image_processor.patch_size
    vision_tokens = grid_t * grid_h * grid_w
    llm_num_vision_tokens = (vision_tokens // image_processor.merge_size //
                             image_processor.merge_size)

    return resized_height, resized_width, llm_num_vision_tokens


def _get_max_image_info(
    image_processor,
    data_type_key: str = "image",
    mm_count: int = 1,
668
669
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
670
):
671
672
673
674
675
676
    # Limit min / max pixels unless they're explicitly provided
    if min_pixels is None:
        min_pixels = max(image_processor.min_pixels, 28 * 28)
    if max_pixels is None:
        max_pixels = min(image_processor.max_pixels, 1280 * 28 * 28)

677
678
679
680
    return _get_vision_info(
        image_processor,
        height=9999999,
        width=9999999,
681
682
        min_pixels=min_pixels,
        max_pixels=max_pixels,
683
684
685
686
687
        data_type_key=data_type_key,
        mm_count=mm_count,
    )


688
689
690
691
692
def get_max_qwen2_vl_mm_tokens(ctx: InputContext,
                               data_type_key: str,
                               *,
                               min_pixels=None,
                               max_pixels=None) -> int:
693
694
695
696
    mm_processor_kwargs = get_mm_processor_kwargs(min_pixels=min_pixels,
                                                  max_pixels=max_pixels)
    image_processor = cached_get_image_processor(ctx.model_config.model,
                                                 **mm_processor_kwargs)
697
698
    max_resized_height, max_resized_width, max_llm_image_tokens = \
        _get_max_image_info(image_processor, data_type_key=data_type_key,
699
700
                            mm_count=1, min_pixels=min_pixels,
                            max_pixels=max_pixels)
701
702
703
704
705
706
707
708
709
710
    return max_llm_image_tokens


get_max_qwen2_vl_image_tokens = partial(get_max_qwen2_vl_mm_tokens,
                                        data_type_key="image")
get_max_qwen2_vl_video_tokens = partial(get_max_qwen2_vl_mm_tokens,
                                        data_type_key="video")


def dummy_data_for_qwen2_vl(
711
712
713
714
715
716
    ctx: InputContext,
    seq_len: int,
    mm_counts: Mapping[str, int],
    *,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None
717
) -> Tuple[SequenceData, Optional[MultiModalDataDict]]:
718
719
720
721
    mm_processor_kwargs = get_mm_processor_kwargs(min_pixels=min_pixels,
                                                  max_pixels=max_pixels)
    image_processor = cached_get_image_processor(ctx.model_config.model,
                                                 **mm_processor_kwargs)
722
723
724
725

    num_images = mm_counts["image"]
    max_resized_height, max_resized_width, max_llm_image_tokens = \
        _get_max_image_info(image_processor, data_type_key="image",
726
727
                            mm_count=num_images, min_pixels=min_pixels,
                            max_pixels=max_pixels)
728
729
730
731
732
733
734
735
736
737
    if seq_len - max_llm_image_tokens - 2 < 0:
        raise RuntimeError(
            f"Qwen2-VL cannot process {num_images} images in a prompt, "
            "please increase max_model_len or reduce image limit by "
            "--limit-mm-per-prompt.")

    # Check video counts.
    num_videos = mm_counts["video"]
    max_resized_height, max_resized_width, max_llm_video_tokens = \
        _get_max_image_info(image_processor, data_type_key="video",
738
739
                            mm_count=num_videos, min_pixels=min_pixels,
                            max_pixels=max_pixels)
740
741
    if seq_len - max_llm_video_tokens - 2 < 0:
        raise RuntimeError(
742
            f"Qwen2-VL cannot process {num_videos} videos in a prompt, "
743
744
745
746
            "please increase max_model_len or reduce video limit by "
            "--limit-mm-per-prompt.")

    hf_config = ctx.get_hf_config(Qwen2VLConfig)
747

748
    dummy_seqdata = SequenceData.from_prompt_token_counts(
749
750
751
752
753
754
        (hf_config.vision_start_token_id, 1),
        (hf_config.image_token_id, max_llm_image_tokens),
        (hf_config.vision_end_token_id, 1),
        (0, seq_len - max_llm_image_tokens - 2),
    )

755
756
757
    dummy_image = Image.new("RGB", (max_resized_width, max_resized_height),
                            color=0)

758
759
760
761
    return DummyData(dummy_seqdata, {
        "image":
        dummy_image if num_images == 1 else [dummy_image] * num_images
    })
762
763
764
765
766
767


def _get_llm_num_vision_tokens(
    mm_inputs: list,
    data_type_key: str,
    image_processor,
768
769
    min_pixels: int,
    max_pixels: int,
770
771
772
773
774
775
776
777
778
):
    """Get number of vision tokens of multimodal inputs.

    This method is derived from `transformers.models.qwen2_vl.
    image_processing_qwen2_vl.Qwen2VLImageProcessor._preprocess`.
    """
    image = to_numpy_array(mm_inputs[0])
    input_data_format = infer_channel_dimension_format(image)
    height, width = get_image_size(image, channel_dim=input_data_format)
779

780
781
782
783
    _, _, llm_num_vision_tokens = _get_vision_info(
        image_processor,
        height=height,
        width=width,
784
785
        min_pixels=min_pixels,
        max_pixels=max_pixels,
786
787
788
789
790
791
792
        do_resize=image_processor.do_resize,
        data_type_key=data_type_key,
        mm_count=len(mm_inputs),
    )
    return llm_num_vision_tokens


793
794
def _expand_pad_tokens(inputs: list, token_id: int, make_batched_fn: Callable,
                       data_type_key: str, image_processor: Any,
795
796
                       prompt_token_ids: List[int], min_pixels: Optional[int],
                       max_pixels: Optional[int]) -> List[int]:
797
798
799
800
801
802
803
804
805
806
    """
    Expand pad tokens for multi-modal inputs (e.g., images or videos).

    Args:
        inputs (list): The multi-modal inputs (e.g., images or videos).
        token_id (int): The token ID used to represent the multi-modal input.
        make_batched_fn (Callable): A function to batch the inputs.
        data_type_key (str): The type of the multi-modal input.
        image_processor (Any): The image processor used to process the inputs.
        prompt_token_ids (List[int]): The list of token IDs in the prompt.
807
808
        min_pixels (int): min pixels to used for img processing
        max_pixels (int): max pixels to be used for img processing
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824

    Returns:
        List[int]: The list of token IDs for the multi-modal inputs.
    """
    indices = [
        idx for idx, token in enumerate(prompt_token_ids) if token == token_id
    ]
    inputs = make_batched_fn(inputs)
    assert len(indices) == len(inputs)

    prompt_token_ids_with_data = []
    for cnt, data in enumerate(inputs):
        num_tokens = _get_llm_num_vision_tokens(
            [data] if data_type_key == "image" else data,
            data_type_key=data_type_key,
            image_processor=image_processor,
825
826
            min_pixels=min_pixels,
            max_pixels=max_pixels,
827
828
829
830
831
832
833
834
835
836
837
838
839
        )
        if cnt == 0:
            end_idx = indices[cnt]
            non_data_tokens = prompt_token_ids[:end_idx]
        else:
            non_data_tokens = prompt_token_ids[indices[cnt - 1] +
                                               1:indices[cnt]]
        prompt_token_ids_with_data.extend(non_data_tokens)
        prompt_token_ids_with_data.extend(token_id for _ in range(num_tokens))
    prompt_token_ids_with_data.extend(prompt_token_ids[indices[-1] + 1:])
    return prompt_token_ids_with_data


840
841
842
def input_processor_for_qwen2_vl(
    ctx: InputContext,
    inputs: DecoderOnlyInputs,
843
844
845
    *,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
846
) -> DecoderOnlyInputs:
847
    multi_modal_data = inputs.get("multi_modal_data")
848
    if multi_modal_data is None:
849
        return inputs
850
851
852
853
854
855

    image_inputs = multi_modal_data.get("image", None)
    video_inputs = multi_modal_data.get("video", None)

    processor = cached_get_processor(ctx.model_config.model)
    image_processor = processor.image_processor
856
857
858
859
    # Apply processor kwarg overrides for image processor options
    min_pixels = min_pixels if min_pixels else image_processor.min_pixels
    max_pixels = max_pixels if max_pixels else image_processor.max_pixels

860
    model_config = ctx.model_config
861
862
863
864
865
866
867
    hf_config = ctx.get_hf_config(Qwen2VLConfig)

    # To avoid redundant processing of vision objects (resize, rescale, etc.),
    # we extract code of calculating number of vision tokens from
    # `transformers.models.qwen2_vl.processing_qwen2_vl.Qwen2VLProcessor`.
    #
    # The following code is equivalent to:
868
    #    prompt = inputs["prompt"]
869
870
871
872
873
874
875
    #    inputs = processor(text=[prompt],
    #                       images=image_inputs,
    #                       videos=video_inputs,
    #                       padding=True,
    #                       return_tensors="pt")
    #    prompt_token_ids = inputs["input_ids"][0].tolist()

876
877
878
879
880
    tokenizer = cached_get_tokenizer(
        model_config.tokenizer,
        trust_remote_code=model_config.trust_remote_code)

    prompt_token_ids = inputs["prompt_token_ids"]
881
882

    # Expand image pad tokens.
883

884
    if image_inputs is not None:
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
        if isinstance(image_inputs, dict):
            prompt_token_ids_with_image = []
            image_indices = [
                idx for idx, token in enumerate(prompt_token_ids)
                if token == hf_config.image_token_id
            ]
            image_cnt = len(image_indices)
            embed_dim = image_inputs.get('image_embeds').size(0)
            assert embed_dim % image_cnt == 0
            num_pad_tokens = embed_dim // image_cnt
            for idx, token in enumerate(prompt_token_ids):
                if idx in image_indices:
                    prompt_token_ids_with_image.extend([token] *
                                                       num_pad_tokens)
                else:
                    prompt_token_ids_with_image.append(token)
            prompt_token_ids = prompt_token_ids_with_image
        else:
            prompt_token_ids = _expand_pad_tokens(image_inputs,
                                                  hf_config.image_token_id,
905
906
                                                  make_batched_images,
                                                  "image",
907
                                                  image_processor,
908
909
910
                                                  prompt_token_ids,
                                                  min_pixels=min_pixels,
                                                  max_pixels=max_pixels)
911

912
    if video_inputs is not None:
913
914
        prompt_token_ids = _expand_pad_tokens(video_inputs,
                                              hf_config.video_token_id,
915
916
                                              make_batched_videos,
                                              "video",
917
                                              image_processor,
918
919
920
                                              prompt_token_ids,
                                              min_pixels=min_pixels,
                                              max_pixels=max_pixels)
921

922
923
924
925
    prompt = inputs.get("prompt")
    if prompt is None:
        prompt = tokenizer.decode(prompt_token_ids)

926
    return token_inputs(
927
        prompt_token_ids=prompt_token_ids,
928
        prompt=prompt,
929
930
931
932
933
934
935
936
937
938
939
940
941
        multi_modal_data=multi_modal_data,
    )


@MULTIMODAL_REGISTRY.register_image_input_mapper(
    image_input_mapper_for_qwen2_vl)
@MULTIMODAL_REGISTRY.register_input_mapper("video",
                                           video_input_mapper_for_qwen2_vl)
@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_qwen2_vl_image_tokens)
@MULTIMODAL_REGISTRY.register_max_multimodal_tokens(
    "video", get_max_qwen2_vl_video_tokens)
@INPUT_REGISTRY.register_dummy_data(dummy_data_for_qwen2_vl)
@INPUT_REGISTRY.register_input_processor(input_processor_for_qwen2_vl)
942
class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
                                      SupportsLoRA, SupportsPP):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    # TODO Support LoRA for the visual encoder in the future.
    supported_lora_modules = [
        "qkv_proj",
        "o_proj",
        "gate_up_proj",
        "down_proj",
    ]
    embedding_modules = {}
    embedding_padding_modules = []
966
967
968
969
970

    def __init__(self,
                 config: Qwen2VLConfig,
                 multimodal_config: MultiModalConfig,
                 cache_config: Optional[CacheConfig] = None,
971
972
973
                 quant_config: Optional[QuantizationConfig] = None,
                 lora_config: Optional[LoRAConfig] = None) -> None:

974
975
976
977
978
979
980
981
982
983
984
        super().__init__()

        assert not cache_config.enable_prefix_caching, \
            "Qwen2-VL currently does not support prefix caching"

        self.config = config
        self.multimodal_config = multimodal_config

        self.visual = Qwen2VisionTransformer(
            config.vision_config,
            norm_eps=getattr(config, "rms_norm_eps", 1e-6),
985
986
            quant_config=quant_config,
            prefix="visual",
987
988
        )

989
990
991
992
        self.model = Qwen2Model(config,
                                cache_config,
                                quant_config,
                                prefix="model")
993

994
995
996
997
998
999
        if get_pp_group().is_last_rank:
            if config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
                self.lm_head = ParallelLMHead(config.vocab_size,
                                              config.hidden_size,
1000
1001
                                              quant_config=quant_config,
                                              prefix="lm_head")
1002
        else:
1003
            self.lm_head = PPMissingLayer()
1004
1005

        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
1006
        self.sampler = get_sampler()
1007
1008
1009
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030

    def _validate_and_reshape_mm_tensor(self,
                                        mm_input: Union[torch.Tensor,
                                                        List[torch.Tensor]],
                                        name: str) -> torch.Tensor:
        if not isinstance(mm_input, (torch.Tensor, list)):
            raise ValueError(f"Incorrect type of {name}. "
                             f"Got type: {type(mm_input)}")
        if isinstance(mm_input, torch.Tensor):
            if mm_input.ndim == 2:
                return mm_input
            if mm_input.ndim != 3:
                raise ValueError(f"{name} should be 2D or batched 3D tensor. "
                                 f"Got ndim: {mm_input.ndim}")
            return torch.concat(list(mm_input))
        else:
            return torch.concat(mm_input)

    def _parse_and_validate_image_input(
            self, **kwargs: object) -> Optional[Qwen2VLImageInputs]:
        pixel_values = kwargs.pop("pixel_values", None)
1031
        image_embeds = kwargs.pop("image_embeds", None)
1032
1033
        image_grid_thw = kwargs.pop("image_grid_thw", None)

1034
        if pixel_values is None and image_embeds is None:
1035
1036
            return None

1037
1038
1039
1040
1041
        if pixel_values is not None:
            pixel_values = self._validate_and_reshape_mm_tensor(
                pixel_values, "image pixel values")
            image_grid_thw = self._validate_and_reshape_mm_tensor(
                image_grid_thw, "image grid_thw")
1042

1043
1044
1045
            if not isinstance(pixel_values, (torch.Tensor, list)):
                raise ValueError("Incorrect type of image pixel values. "
                                 f"Got type: {type(pixel_values)}")
1046

1047
1048
1049
1050
1051
            return Qwen2VLImagePixelInputs(type="pixel_values",
                                           data=pixel_values,
                                           image_grid_thw=image_grid_thw)

        if image_embeds is not None:
1052
1053
1054
            image_embeds = self._validate_and_reshape_mm_tensor(
                image_embeds, "image embeds")

1055
1056
1057
1058
1059
            if not isinstance(image_embeds, torch.Tensor):
                raise ValueError("Incorrect type of image embeddings. "
                                 f"Got type: {type(image_embeds)}")
            return Qwen2VLImageEmbeddingInputs(type="image_embeds",
                                               data=image_embeds)
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080

    def _parse_and_validate_video_input(
            self, **kwargs: object) -> Optional[Qwen2VLVideoInputs]:
        pixel_values_videos = kwargs.pop("pixel_values_videos", None)
        video_grid_thw = kwargs.pop("video_grid_thw", None)

        if pixel_values_videos is None:
            return None

        pixel_values_videos = self._validate_and_reshape_mm_tensor(
            pixel_values_videos, "video pixel values")
        video_grid_thw = self._validate_and_reshape_mm_tensor(
            video_grid_thw, "video grid_thw")

        return Qwen2VLVideoInputs(
            pixel_values_videos=pixel_values_videos,
            video_grid_thw=video_grid_thw,
        )

    def _process_image_input(self,
                             image_input: Qwen2VLImageInputs) -> torch.Tensor:
1081
1082
1083
1084
        if image_input["type"] == "image_embeds":
            return image_input["data"].type(self.visual.dtype)

        pixel_values = image_input["data"].type(self.visual.dtype)
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
        image_embeds = self.visual(pixel_values,
                                   grid_thw=image_input["image_grid_thw"])
        return image_embeds

    def _process_video_input(self,
                             video_input: Qwen2VLVideoInputs) -> torch.Tensor:
        pixel_values_videos = video_input["pixel_values_videos"].type(
            self.visual.dtype)
        video_embeds = self.visual(pixel_values_videos,
                                   grid_thw=video_input["video_grid_thw"])
        return video_embeds

    def _merge_multimodal_embeddings(
        self,
        input_ids: torch.Tensor,
        inputs_embeds: torch.Tensor,
        multimodal_embeddings: torch.Tensor,
        placeholder_token_id: int,
    ) -> torch.Tensor:
        mask = (input_ids == placeholder_token_id)
        inputs_embeds[mask, :] = multimodal_embeddings
        return inputs_embeds

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        **kwargs: object,
1116
    ) -> Union[torch.Tensor, IntermediateTensors]:
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
        """Run forward pass for Qwen2-VL.

        Args:
            input_ids: Flattened (concatenated) input_ids corresponding to a
                batch.
            positions: Flattened (concatenated) position ids corresponding to a
                batch.
                **NOTE**: If mrope is enabled (default setting for Qwen2-VL
                opensource models), the shape will be `(3, seq_len)`,
                otherwise it will be `(seq_len,).
            pixel_values: Pixel values to be fed to a model.
                `None` if no images are passed.
            image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in LLM.
                `None` if no images are passed.
            pixel_values_videos: Pixel values of videos to be fed to a model.
                `None` if no videos are passed.
            video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in LLM.
                `None` if no videos are passed.
        """
1136
1137
        if intermediate_tensors is not None:
            input_ids = None
1138
1139
            inputs_embeds = None
        else:
1140
1141
            image_input = self._parse_and_validate_image_input(**kwargs)
            video_input = self._parse_and_validate_video_input(**kwargs)
1142

1143
1144
1145
            if image_input is None and video_input is None:
                inputs_embeds = None
            else:
1146
                if uses_mrope(self.config):
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
                    assert positions.ndim == 2 and positions.size(0) == 3, (
                        "multimodal section rotary embedding requires "
                        f"(3, seq_len) positions, but got {positions.size()}")

                inputs_embeds = self.model.embed_tokens(input_ids)

                if image_input is not None:
                    image_embeds = self._process_image_input(image_input)
                    inputs_embeds = self._merge_multimodal_embeddings(
                        input_ids,
                        inputs_embeds,
                        image_embeds,
                        placeholder_token_id=self.config.image_token_id,
                    )

                if video_input is not None:
                    video_embeds = self._process_video_input(video_input)
                    inputs_embeds = self._merge_multimodal_embeddings(
                        input_ids,
                        inputs_embeds,
                        video_embeds,
                        placeholder_token_id=self.config.video_token_id,
                    )

                input_ids = None
1172
1173
1174
1175
1176
1177

        hidden_states = self.model(
            input_ids=input_ids,
            positions=positions,
            kv_caches=kv_caches,
            attn_metadata=attn_metadata,
1178
            intermediate_tensors=intermediate_tensors,
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
            inputs_embeds=inputs_embeds,
        )
        return hidden_states

    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        logits = self.logits_processor(self.lm_head, hidden_states,
                                       sampling_metadata)
        return logits

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

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        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", "up_proj", 1),
            ("gate_up_proj", "gate_proj", 0),
        ]
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
            if self.config.tie_word_embeddings and "lm_head.weight" in name:
                continue
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
1216
1217
1218
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
1219
1220
                if is_pp_missing_parameter(name, self):
                    continue
1221
1222
1223
1224
1225
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
1226
                if "visual" in name and name.endswith("qkv.weight"):
1227
1228
1229
1230
1231
1232
1233
1234
                    visual_num_heads = self.config.vision_config.num_heads
                    visual_embed_dim = self.config.vision_config.embed_dim
                    head_size = visual_embed_dim // visual_num_heads
                    loaded_weight = loaded_weight.view(3, visual_num_heads,
                                                       head_size,
                                                       visual_embed_dim)
                    loaded_weight = loaded_weight.transpose(0, 1)
                    loaded_weight = loaded_weight.reshape(-1, visual_embed_dim)
1235
                elif "visual" in name and name.endswith("qkv.bias"):
1236
1237
1238
1239
1240
1241
1242
1243
                    visual_num_heads = self.config.vision_config.num_heads
                    visual_embed_dim = self.config.vision_config.embed_dim
                    head_size = visual_embed_dim // visual_num_heads
                    loaded_weight = loaded_weight.view(3, visual_num_heads,
                                                       head_size)
                    loaded_weight = loaded_weight.transpose(0, 1)
                    loaded_weight = loaded_weight.reshape(-1)
                try:
1244
1245
1246
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue
1247
1248
                    if is_pp_missing_parameter(name, self):
                        continue
1249
1250
                    param = params_dict[name]
                except KeyError:
1251
                    raise ValueError(f"Unexpected weight: {name}") from None
1252
1253
1254
1255

                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)