qwen2_vl.py 57 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 cached_property, partial
25
from typing import (Any, Callable, Dict, Iterable, List, Literal, Mapping,
26
                    Optional, Set, 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.config import VllmConfig
43
from vllm.distributed import parallel_state
44
from vllm.distributed import utils as dist_utils
45
46
from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,
                         InputContext, token_inputs)
47
48
49
50
51
52
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.quantization import QuantizationConfig
53
54
55
from vllm.model_executor.layers.quantization.gptq import GPTQConfig
from vllm.model_executor.layers.quantization.gptq_marlin import (
    GPTQMarlinConfig)
Joe Runde's avatar
Joe Runde committed
56
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
57
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
58
from vllm.multimodal import MULTIMODAL_REGISTRY
59
from vllm.multimodal.image import cached_get_image_processor
60
from vllm.multimodal.inputs import (MultiModalData, MultiModalDataDict,
61
                                    MultiModalKwargs, NestedTensors)
62
from vllm.multimodal.utils import cached_get_tokenizer
63
from vllm.platforms import _Backend
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 (AutoWeightsLoader, WeightsMapper, get_vit_attn_backend,
                    init_vllm_registered_model, maybe_prefix)
71

zhuwenwen's avatar
zhuwenwen committed
72
73
74
75
76
import os
import re
from vllm import _custom_ops as ops
from vllm.model_executor.utils import pad_weight, gemm_bank_conf

77
78
79
80
81
logger = init_logger(__name__)

# === Vision Inputs === #


82
83
class Qwen2VLImagePixelInputs(TypedDict):
    type: Literal["pixel_values"]
84
    pixel_values: torch.Tensor
85
    """Shape:
86
87
88
89
90
91
92
93
94
    `(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.
    """


95
96
class Qwen2VLImageEmbeddingInputs(TypedDict):
    type: Literal["image_embeds"]
97
98
99
100
101
102
    image_embeds: torch.Tensor
    """Supported types:
    - List[`torch.Tensor`]: A list of tensors holding all images' features.
        Each tensor holds an image's features.
    - `torch.Tensor`: A tensor holding all images' features
        (concatenation of all images' feature tensors).
103
    
104
105
106
107
108
109
110
111
    Tensor shape: `(num_image_features, hidden_size)`
    - `num_image_features` varies based on
        the number and resolution of the images.
    - `hidden_size` must match the hidden size of language model backbone.
    """

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


116
117
118
119
Qwen2VLImageInputs = Union[Qwen2VLImagePixelInputs,
                           Qwen2VLImageEmbeddingInputs]


120
121
class Qwen2VLVideoPixelInputs(TypedDict):
    type: Literal["pixel_values_videos"]
122
    pixel_values_videos: torch.Tensor
123
124
    """Shape:
    `(num_patches,
125
126
127
128
129
      num_channels * temporal_patch_size * patch_size * patch_size)`
    """

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

131
132
133
134
    This should be in `(grid_t, grid_h, grid_w)` format.
    """


135
136
137
138
139
140
141
142
class Qwen2VLVideoEmbeddingInputs(TypedDict):
    type: Literal["video_embeds"]
    video_embeds: torch.Tensor
    """Supported types:
    - List[`torch.Tensor`]: A list of tensors holding all videos' features.
        Each tensor holds an video's features.
    - `torch.Tensor`: A tensor holding all videos' features
      (concatenation of all videos' feature tensors).
143
    
144
145
146
147
148
149
150
151
    Tensor shape: `(num_image_features, hidden_size)`
    - `num_image_features` varies based on 
        the number and resolution of the videos.
    - `hidden_size` must match the hidden size of language model backbone.
    """

    video_grid_thw: torch.Tensor
    """Shape: `(num_videos, 3)`
152
153
154
155
    This should be in `(grid_t, grid_h, grid_w)` format.
    """


156
157
158
Qwen2VLVideoInputs = Union[Qwen2VLVideoPixelInputs,
                           Qwen2VLVideoEmbeddingInputs]

159
160
161
162
163
164
165
166
167
168
169
# === 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,
170
        prefix: str = "",
171
172
173
174
    ):
        super().__init__()
        self.fc1 = ColumnParallelLinear(in_features,
                                        hidden_features,
175
176
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.fc1")
177
178
179
        self.act = act_layer()
        self.fc2 = RowParallelLinear(hidden_features,
                                     in_features,
180
181
                                     quant_config=quant_config,
                                     prefix=f"{prefix}.fc2")
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242

    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,
243
        prefix: str = "",
244
245
246
247
248
249
250
251
252
253
254
    ) -> 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,
255
256
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.qkv")
257
258
        self.proj = RowParallelLinear(input_size=projection_size,
                                      output_size=embed_dim,
259
260
                                      quant_config=quant_config,
                                      prefix=f"{prefix}.proj")
261
262

        # Detect attention implementation.
263
        self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
264
265
266
267
268
        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.")
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289

    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]

290
291
        q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
                   for x in (q, k, v))
292
293
294
295
        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)

296
        if self.attn_backend == _Backend.FLASH_ATTN:
297
298
299
300
            # from vllm_flash_attn.flash_attn_interface import (
            #   flash_attn_varlen_func)
            from flash_attn import flash_attn_varlen_func

301
            q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316

            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)
317
        elif self.attn_backend == _Backend.TORCH_SDPA:
318
            seq_length = q.size(1)
319
            q, k, v = (rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
320
321
322
323
324
325
326
327
328
329
330
331
            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 ")
332
        elif self.attn_backend == _Backend.XFORMERS:
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
            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,
359
        prefix: str = "",
360
361
362
363
364
365
366
367
368
369
370
    ) -> 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,
371
372
                                         quant_config=quant_config,
                                         prefix=f"{prefix}.attn")
373
374
375
        self.mlp = Qwen2VisionMLP(dim,
                                  mlp_hidden_dim,
                                  act_layer=act_layer,
376
377
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.mlp")
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425

    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,
426
        prefix: str = "",
427
428
429
430
431
432
433
434
435
436
    ) -> 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,
437
438
                                 quant_config=quant_config,
                                 prefix=f"{prefix}.mlp.0"),
439
440
441
442
            nn.GELU(),
            RowParallelLinear(self.hidden_size,
                              d_model,
                              bias=True,
443
444
                              quant_config=quant_config,
                              prefix=f"{prefix}.mlp.2"),
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
        ])

    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,
495
        prefix: str = "",
496
497
498
499
500
501
502
503
504
505
506
507
508
509
    ) -> 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
510
511
        self.num_heads = num_heads
        self.embed_dim = embed_dim
512
513
514
515
516
517
518
519
520
521
522
523
524

        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([
525
526
527
528
529
530
531
            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)
532
533
534
535
536
537
        ])
        self.merger = Qwen2VisionPatchMerger(
            d_model=hidden_size,
            context_dim=embed_dim,
            norm_layer=norm_layer,
            quant_config=quant_config,
538
            prefix=f"{prefix}.merger",
539
540
541
542
        )

    @property
    def dtype(self) -> torch.dtype:
543
        return self.patch_embed.proj.weight.dtype
544
545
546

    @property
    def device(self) -> torch.device:
547
        return self.patch_embed.proj.weight.device
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
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

    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

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
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
        ]
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        loaded_params: Set[str] = set()

        for name, loaded_weight in weights:
            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)

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                if name.endswith("qkv.weight"):
                    visual_num_heads = self.num_heads
                    visual_embed_dim = self.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)
                elif name.endswith("qkv.bias"):
                    visual_num_heads = self.num_heads
                    visual_embed_dim = self.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)

                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
zhuwenwen's avatar
zhuwenwen committed
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
            
        if self.use_llama_nn and self.quant_method is None:
            lay_key_words = [
                "attn.qkv.weight",
                "attn.proj.weight",
                "mlp.fc1.weight",
                "mlp.fc2.weight",
                "mlp.0.weight",
                "mlp.2.weight",
                "self_attn.qkv_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
                "mlp.down_proj.weight",
                "lm_head.weight",
            ]
            combined_words = "|".join(lay_key_words)
            
            lay_qkv_words = ["attn.qkv.weight"]   
            qkv_words = "|".join(lay_qkv_words)  
            
            lay_qkv_bias_words = ["attn.qkv.bias"]   
            qkv_bias_words = "|".join(lay_qkv_bias_words) 
            
            for layername, weight in params_dict.items():
                # if self.use_fa_pad and (re.findall(qkv_bias_words, layername)):
                #     weight.data = pad_weight(weight.data, 32)
                    
                matches = re.findall(combined_words, layername)
                if matches:   
                    if self.use_gemm_pad and gemm_bank_conf(weight.data.shape[0]):
                        weight.data = pad_weight(weight.data, 32)  
                    
                    # if self.use_fa_pad and (re.findall(qkv_words, layername)):
                    #     if not gemm_bank_conf(weight.data.shape[0]):
                    #         weight.data = pad_weight(weight.data, 32)
                        
                    _weight = torch.zeros_like(weight.data)
                    ori_shape =_weight.shape
                    
                    ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
                    weight.data.copy_(_weight)
                    
                    weight.data=weight.data.reshape(ori_shape[1],-1)
                    
690
691
        return loaded_params

692
693
694
695

# === Vision input helpers === #


696
697
698
699
700
701
702
703
704
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
705
706
707
708
709
710


def mm_input_mapper_for_qwen2_vl(
    ctx: InputContext,
    data: MultiModalData[object],
    data_type_key: str,
711
712
713
    *,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
714
) -> MultiModalKwargs:
715
    """Input mapper for Qwen2-VL."""
716
    if data_type_key == "image" and isinstance(data, dict):
717
        return MultiModalKwargs({
718
719
720
            "image_embeds": data.get("image_embeds"),
            "image_grid_thw": data.get("image_grid_thw"),
        })
721
722
723
724
725
726
    if data_type_key == "video" and isinstance(data, dict):
        return MultiModalKwargs({
            "video_embeds": data.get("video_embeds"),
            "video_grid_thw": data.get("video_grid_thw"),
        })

727
    model_config = ctx.model_config
728
729
    # Handle mm processor kwargs; we pass these at creation time
    # because preprocess() in transformers doesn't expose them
730
731
    mm_processor_kwargs = get_mm_processor_kwargs(min_pixels=min_pixels,
                                                  max_pixels=max_pixels)
732
    image_processor = cached_get_image_processor(
733
734
735
736
        model_config.model,
        trust_remote_code=model_config.trust_remote_code,
        **mm_processor_kwargs,
    )
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
    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

757
    return MultiModalKwargs(batch_data)
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808


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,
809
810
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
811
):
812
813
814
815
816
817
    # 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)

818
819
820
821
    return _get_vision_info(
        image_processor,
        height=9999999,
        width=9999999,
822
823
        min_pixels=min_pixels,
        max_pixels=max_pixels,
824
825
826
827
828
        data_type_key=data_type_key,
        mm_count=mm_count,
    )


829
830
831
832
833
def get_max_qwen2_vl_mm_tokens(ctx: InputContext,
                               data_type_key: str,
                               *,
                               min_pixels=None,
                               max_pixels=None) -> int:
834
835
836
837
    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)
838
839
    max_resized_height, max_resized_width, max_llm_image_tokens = \
        _get_max_image_info(image_processor, data_type_key=data_type_key,
840
841
                            mm_count=1, min_pixels=min_pixels,
                            max_pixels=max_pixels)
842
843
844
845
846
847
848
849
850
851
    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(
852
853
854
855
856
857
    ctx: InputContext,
    seq_len: int,
    mm_counts: Mapping[str, int],
    *,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None
858
) -> Tuple[SequenceData, Optional[MultiModalDataDict]]:
859
860
861
862
    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)
863
864
865
866

    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",
867
868
                            mm_count=num_images, min_pixels=min_pixels,
                            max_pixels=max_pixels)
869
870
871
872
873
874
875
876
877
878
    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",
879
880
                            mm_count=num_videos, min_pixels=min_pixels,
                            max_pixels=max_pixels)
881
882
    if seq_len - max_llm_video_tokens - 2 < 0:
        raise RuntimeError(
883
            f"Qwen2-VL cannot process {num_videos} videos in a prompt, "
884
885
886
887
            "please increase max_model_len or reduce video limit by "
            "--limit-mm-per-prompt.")

    hf_config = ctx.get_hf_config(Qwen2VLConfig)
888

889
    dummy_seqdata = SequenceData.from_prompt_token_counts(
890
891
892
893
894
895
        (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),
    )

896
897
898
    dummy_image = Image.new("RGB", (max_resized_width, max_resized_height),
                            color=0)

899
900
901
902
    return DummyData(dummy_seqdata, {
        "image":
        dummy_image if num_images == 1 else [dummy_image] * num_images
    })
903
904
905
906
907
908


def _get_llm_num_vision_tokens(
    mm_inputs: list,
    data_type_key: str,
    image_processor,
909
910
    min_pixels: int,
    max_pixels: int,
911
912
913
914
915
916
917
918
919
):
    """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)
920

921
922
923
924
    _, _, llm_num_vision_tokens = _get_vision_info(
        image_processor,
        height=height,
        width=width,
925
926
        min_pixels=min_pixels,
        max_pixels=max_pixels,
927
928
929
930
931
932
933
        do_resize=image_processor.do_resize,
        data_type_key=data_type_key,
        mm_count=len(mm_inputs),
    )
    return llm_num_vision_tokens


934
935
def _expand_pad_tokens(inputs: list, token_id: int, make_batched_fn: Callable,
                       data_type_key: str, image_processor: Any,
936
937
                       prompt_token_ids: List[int], min_pixels: Optional[int],
                       max_pixels: Optional[int]) -> List[int]:
938
939
940
941
942
943
944
945
946
947
    """
    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.
948
949
        min_pixels (int): min pixels to used for img processing
        max_pixels (int): max pixels to be used for img processing
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965

    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,
966
967
            min_pixels=min_pixels,
            max_pixels=max_pixels,
968
969
970
971
972
973
974
975
976
977
978
979
980
        )
        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


981
982
983
def input_processor_for_qwen2_vl(
    ctx: InputContext,
    inputs: DecoderOnlyInputs,
984
985
986
    *,
    min_pixels: Optional[int] = None,
    max_pixels: Optional[int] = None,
987
) -> DecoderOnlyInputs:
988
    multi_modal_data = inputs.get("multi_modal_data")
989
    if multi_modal_data is None:
990
        return inputs
991
992
993
994
995
996

    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
997
998
999
1000
    # 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

1001
    model_config = ctx.model_config
1002
1003
1004
1005
1006
1007
1008
    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:
1009
    #    prompt = inputs["prompt"]
1010
1011
1012
1013
1014
1015
1016
    #    inputs = processor(text=[prompt],
    #                       images=image_inputs,
    #                       videos=video_inputs,
    #                       padding=True,
    #                       return_tensors="pt")
    #    prompt_token_ids = inputs["input_ids"][0].tolist()

1017
1018
1019
1020
1021
    tokenizer = cached_get_tokenizer(
        model_config.tokenizer,
        trust_remote_code=model_config.trust_remote_code)

    prompt_token_ids = inputs["prompt_token_ids"]
1022
1023

    # Expand image pad tokens.
1024

1025
    if image_inputs is not None:
1026
1027
1028
1029
1030
1031
        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
            ]
1032
1033
1034
1035
1036
1037
1038
1039

            # ensure all image tokens have grid_thw
            assert \
                len(image_indices) == image_inputs["image_grid_thw"].size(0), \
                "image token num does not match image_grid_thw.shape"

            image_counter = 0
            pad_token_counter = 0
1040
1041
            for idx, token in enumerate(prompt_token_ids):
                if idx in image_indices:
1042
1043
1044
1045
1046
                    grid_thw = image_inputs["image_grid_thw"][image_counter]
                    grid_t, grid_h, grid_w = grid_thw
                    num_pad_tokens = (grid_t * grid_h * grid_w //
                                      image_processor.merge_size //
                                      image_processor.merge_size)
1047
1048
                    prompt_token_ids_with_image.extend([token] *
                                                       num_pad_tokens)
1049
1050
                    image_counter += 1
                    pad_token_counter += num_pad_tokens
1051
1052
                else:
                    prompt_token_ids_with_image.append(token)
1053
1054
1055
1056
1057
1058

            # ensure all embeddings are used
            assert \
                pad_token_counter == image_inputs["image_embeds"].size(0), \
                "image_embeds.shape does not match image_grid_thw"

1059
1060
1061
1062
            prompt_token_ids = prompt_token_ids_with_image
        else:
            prompt_token_ids = _expand_pad_tokens(image_inputs,
                                                  hf_config.image_token_id,
1063
1064
                                                  make_batched_images,
                                                  "image",
1065
                                                  image_processor,
1066
1067
1068
                                                  prompt_token_ids,
                                                  min_pixels=min_pixels,
                                                  max_pixels=max_pixels)
1069

1070
    if video_inputs is not None:
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
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
        if isinstance(video_inputs, dict):
            prompt_token_ids_with_video = []
            video_indices = [
                idx for idx, token in enumerate(prompt_token_ids)
                if token == hf_config.video_token_id
            ]

            # ensure all video tokens have grid_thw
            assert \
                len(video_indices) == video_inputs["video_grid_thw"].size(0), \
                "video token num does not match video_grid_thw.shape"

            video_counter = 0
            pad_token_counter = 0
            for idx, token in enumerate(prompt_token_ids):
                if idx in video_indices:
                    grid_thw = video_inputs["video_grid_thw"][video_counter]
                    grid_t, grid_h, grid_w = grid_thw
                    num_pad_tokens = (grid_t * grid_h * grid_w //
                                      image_processor.merge_size //
                                      image_processor.merge_size)
                    prompt_token_ids_with_video.extend([token] *
                                                       num_pad_tokens)
                    video_counter += 1
                    pad_token_counter += num_pad_tokens
                else:
                    prompt_token_ids_with_video.append(token)

            # ensure all embeddings are used
            assert \
                pad_token_counter == video_inputs["video_embeds"].size(0), \
                "video_embeds.shape does not match video_grid_thw"

            prompt_token_ids = prompt_token_ids_with_video
        else:
            prompt_token_ids = _expand_pad_tokens(video_inputs,
                                                  hf_config.video_token_id,
                                                  make_batched_videos,
                                                  "video",
                                                  image_processor,
                                                  prompt_token_ids,
                                                  min_pixels=min_pixels,
                                                  max_pixels=max_pixels)
1114

1115
1116
1117
1118
    prompt = inputs.get("prompt")
    if prompt is None:
        prompt = tokenizer.decode(prompt_token_ids)

1119
    return token_inputs(
1120
        prompt_token_ids=prompt_token_ids,
1121
        prompt=prompt,
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
        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)
1135
class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
                                      SupportsLoRA, SupportsPP):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }
1148

1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
    # 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 = []
1159

1160
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1161
        super().__init__()
1162
1163
1164
1165
        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
1166
1167
1168
1169
1170
1171
1172
1173
1174
        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),
1175
            quant_config=self._maybe_ignore_quant_config(quant_config),
1176
            prefix=maybe_prefix(prefix, "visual"),
1177
1178
        )

1179
1180
1181
1182
1183
        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "language_model"),
            architectures=["Qwen2ForCausalLM"],
        )
1184

1185
        self.make_empty_intermediate_tensors = (
1186
            self.language_model.make_empty_intermediate_tensors)
zhuwenwen's avatar
zhuwenwen committed
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
        
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
            
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
        self.use_gemm_pad = os.environ.get('GEMM_PAD') == '1'
        self.use_fa_pad = os.environ.get('FA_PAD') == '1'
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
1197

1198
1199
1200
1201
    @cached_property
    def sampler(self):
        if hasattr(self.language_model, "sampler"):
            return self.language_model.sampler
zhuwenwen's avatar
zhuwenwen committed
1202
        
1203
        return get_sampler()
1204

1205
1206
1207
1208
1209
1210
1211
    def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):
        # GPTQ configs do not have a list of ignored modules, however AutoGPTQ
        # seems to avoid vision encoder sections for some models.
        # See: https://huggingface.co/Qwen/Qwen2-VL-2B-Instruct-GPTQ-Int4
        if isinstance(quant_config, (GPTQConfig, GPTQMarlinConfig)):
            return None
        return quant_config
zhuwenwen's avatar
zhuwenwen committed
1212
    
1213

1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
    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)
1234
        image_embeds = kwargs.pop("image_embeds", None)
1235
1236
        image_grid_thw = kwargs.pop("image_grid_thw", None)

1237
        if pixel_values is None and image_embeds is None:
1238
1239
            return None

1240
1241
1242
1243
1244
        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")
1245

1246
1247
1248
            if not isinstance(pixel_values, (torch.Tensor, list)):
                raise ValueError("Incorrect type of image pixel values. "
                                 f"Got type: {type(pixel_values)}")
1249

1250
            return Qwen2VLImagePixelInputs(type="pixel_values",
1251
                                           pixel_values=pixel_values,
1252
1253
1254
                                           image_grid_thw=image_grid_thw)

        if image_embeds is not None:
1255
1256
            image_embeds = self._validate_and_reshape_mm_tensor(
                image_embeds, "image embeds")
1257
1258
            image_grid_thw = self._validate_and_reshape_mm_tensor(
                image_grid_thw, "image grid_thw")
1259

1260
1261
1262
1263
            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",
1264
1265
                                               image_embeds=image_embeds,
                                               image_grid_thw=image_grid_thw)
1266
1267
1268
1269

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

1273
        if pixel_values_videos is None and video_embeds is None:
1274
1275
            return None

1276
1277
1278
1279
1280
        if pixel_values_videos is not 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")
1281

1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
            return Qwen2VLVideoPixelInputs(
                type="pixel_values_videos",
                pixel_values_videos=pixel_values_videos,
                video_grid_thw=video_grid_thw,
            )

        if video_embeds is not None:
            video_embeds = self._validate_and_reshape_mm_tensor(
                video_embeds, "video embeds")
            video_grid_thw = self._validate_and_reshape_mm_tensor(
                video_grid_thw, "video grid_thw")

            if not isinstance(video_embeds, torch.Tensor):
                raise ValueError("Incorrect type of video embeddings. "
                                 f"Got type: {type(video_embeds)}")
            return Qwen2VLVideoEmbeddingInputs(type="video_embeds",
                                               video_embeds=video_embeds,
                                               video_grid_thw=video_grid_thw)
1300
1301
1302

    def _process_image_input(self,
                             image_input: Qwen2VLImageInputs) -> torch.Tensor:
1303
        if image_input["type"] == "image_embeds":
1304
            return image_input["image_embeds"].type(self.visual.dtype)
1305

1306
1307
1308
1309
1310
1311
1312
        pixel_values = image_input["pixel_values"].type(self.visual.dtype)
        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:
1313
1314
1315
        if video_input["type"] == "video_embeds":
            return video_input["video_embeds"].type(self.visual.dtype)

1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
        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

1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
    def get_multimodal_embeddings(
            self, **kwargs) -> Optional[List[Tuple[NestedTensors, str]]]:

        image_input = self._parse_and_validate_image_input(**kwargs)
        video_input = self._parse_and_validate_video_input(**kwargs)
        if image_input is None and video_input is None:
            return None

        # We make a tuple of each embedding with its modality string. This is a
        # temporary workaround for models to handle mixed modalities when
        # get_multimodal_embeddings and get_input_embeddings are called
        # separately.
        # TODO(ywang96): Add support for mixed-modality inference for v1.
        multimodal_embeddings: List[Tuple[NestedTensors, str]] = []

        if image_input is not None:
            image_embeds = self._process_image_input(image_input)
            multimodal_embeddings.append((image_embeds, "image"))
        if video_input is not None:
            video_embeds = self._process_video_input(video_input)
            multimodal_embeddings.append((video_embeds, "video"))

        return multimodal_embeddings

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: Optional[List[Tuple[NestedTensors,
                                                   str]]] = None,
    ) -> torch.Tensor:
1363
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
        if multimodal_embeddings is not None:
            for embeddings, modality in multimodal_embeddings:
                if modality == "image":
                    inputs_embeds = self._merge_multimodal_embeddings(
                        input_ids,
                        inputs_embeds,
                        embeddings,
                        placeholder_token_id=self.config.image_token_id,
                    )
                if modality == "video":
                    inputs_embeds = self._merge_multimodal_embeddings(
                        input_ids,
                        inputs_embeds,
                        embeddings,
                        placeholder_token_id=self.config.video_token_id,
                    )
        return inputs_embeds

1382
1383
1384
1385
1386
1387
1388
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
1389
        inputs_embeds: Optional[torch.Tensor] = None,
1390
        **kwargs: object,
1391
    ) -> Union[torch.Tensor, IntermediateTensors]:
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
        """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.
        """

1412
        if intermediate_tensors is not None:
1413
            inputs_embeds = None
1414

1415
1416
1417
1418
1419
1420
1421
1422
1423
        # NOTE: In v1, inputs_embeds is always generated at model runner, this
        # condition is for v0 compatibility.
        elif inputs_embeds is None:
            multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

            # We need to check for usage of mrope here in case there is
            # multimodal data.
            # TODO (ywang96): move this to model runner in V1.
            if multimodal_embeddings is not None and uses_mrope(self.config):
1424
1425
1426
1427
                assert positions.ndim == 2 and positions.size(0) == 3, (
                    "multimodal section rotary embedding requires "
                    f"(3, seq_len) positions, but got {positions.size()}")

1428
1429
            inputs_embeds = self.get_input_embeddings(input_ids,
                                                      multimodal_embeddings)
1430
1431
            input_ids = None

1432
        hidden_states = self.language_model.model(
1433
1434
1435
1436
            input_ids=input_ids,
            positions=positions,
            kv_caches=kv_caches,
            attn_metadata=attn_metadata,
1437
            intermediate_tensors=intermediate_tensors,
1438
1439
1440
1441
            inputs_embeds=inputs_embeds,
        )
        return hidden_states

1442
1443
1444
1445
1446
1447
1448
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
1449
1450
1451
1452
1453
1454

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

1457
1458
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
1459
1460
1461
1462
1463
1464
1465
1466
        hf_to_vllm_mapper = WeightsMapper(
            orig_to_new_prefix={
                "lm_head.": "language_model.lm_head.",
                "model.": "language_model.model.",
            })

        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=hf_to_vllm_mapper)