blip.py 16.9 KB
Newer Older
1
2
"""Minimal implementation of BlipVisionModel intended to be only used 
within a vision language model."""
3
from typing import Iterable, Optional, Set, Tuple, Union
4
5
6

import torch
import torch.nn as nn
7
import torch.nn.functional as F
8
9
10
from PIL import Image
from transformers import Blip2VisionConfig, BlipVisionConfig

11
from vllm.attention.selector import _Backend
12
from vllm.config import ModelConfig
13
from vllm.distributed import divide, get_tensor_model_parallel_world_size
14
from vllm.inputs import DecoderOnlyInputs, token_inputs
15
16
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17
                                               QKVParallelLinear,
18
19
                                               RowParallelLinear)
from vllm.model_executor.layers.quantization import QuantizationConfig
20
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
22
from vllm.multimodal.utils import (cached_get_tokenizer,
                                   repeat_and_pad_placeholder_tokens)
23
from vllm.sequence import SequenceData
24

25
from .utils import get_vit_attn_backend
26

27
28
29
30
31
32
33
34
35
36
37
38
39

def get_blip_patch_grid_length(*, image_size: int, patch_size: int) -> int:
    assert image_size % patch_size == 0
    return image_size // patch_size


def get_blip_num_patches(*, image_size: int, patch_size: int) -> int:
    grid_length = get_blip_patch_grid_length(image_size=image_size,
                                             patch_size=patch_size)
    return grid_length * grid_length


def get_blip_image_feature_size(
40
        hf_config: Union[BlipVisionConfig, Blip2VisionConfig]) -> int:
41
42
43
44
45
    return get_blip_num_patches(image_size=hf_config.image_size,
                                patch_size=hf_config.patch_size)


def get_max_blip_image_tokens(
46
        hf_config: Union[BlipVisionConfig, Blip2VisionConfig]) -> int:
47
48
49
50
51
52
    return get_blip_image_feature_size(hf_config)


def dummy_seq_data_for_blip(
    hf_config: Union[BlipVisionConfig, Blip2VisionConfig],
    seq_len: int,
53
    num_images: int,
54
55
56
57
58
59
60
61
62
    *,
    image_token_id: int,
    image_feature_size_override: Optional[int] = None,
):
    if image_feature_size_override is None:
        image_feature_size = get_blip_image_feature_size(hf_config)
    else:
        image_feature_size = image_feature_size_override

63
    return SequenceData.from_prompt_token_counts(
64
65
66
        (image_token_id, image_feature_size * num_images),
        (0, seq_len - image_feature_size * num_images),
    )
67
68
69
70


def dummy_image_for_blip(
    hf_config: Union[BlipVisionConfig, Blip2VisionConfig],
71
    num_images: int,
72
73
74
75
76
77
78
79
80
81
82
    *,
    image_width_override: Optional[int] = None,
    image_height_override: Optional[int] = None,
):
    width = height = hf_config.image_size
    if image_width_override is not None:
        width = image_width_override
    if image_height_override is not None:
        height = image_height_override

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


def input_processor_for_blip(
    model_config: ModelConfig,
    hf_config: Union[BlipVisionConfig, Blip2VisionConfig],
89
    inputs: DecoderOnlyInputs,
90
91
92
93
    *,
    image_token_id: int,
    image_feature_size_override: Optional[int] = None,
):
94
    multi_modal_data = inputs.get("multi_modal_data")
95
    if multi_modal_data is None or "image" not in multi_modal_data:
96
        return inputs
97

98
99
100
101
102
    if "multi_modal_placeholders" in inputs and "image" in inputs[
            "multi_modal_placeholders"]:
        # The inputs already have placeholders.
        return inputs

103
104
105
106
107
108
109
    tokenizer = cached_get_tokenizer(model_config.tokenizer)

    if image_feature_size_override is None:
        image_feature_size = get_blip_image_feature_size(hf_config)
    else:
        image_feature_size = image_feature_size_override

110
    new_prompt, new_token_ids, ranges = repeat_and_pad_placeholder_tokens(
111
        tokenizer,
112
113
        inputs.get("prompt"),
        inputs["prompt_token_ids"],
114
        placeholder_token_id=image_token_id,
115
116
117
118
        repeat_count=image_feature_size,
    )

    # NOTE: Create a defensive copy of the original inputs
119
120
    return token_inputs(prompt_token_ids=new_token_ids,
                        prompt=new_prompt,
121
122
                        multi_modal_data=multi_modal_data,
                        multi_modal_placeholders={"image": ranges})
123
124
125
126
127


# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/blip/modeling_blip.py#L164 # noqa
class BlipVisionEmbeddings(nn.Module):

128
    def __init__(self, config: Union[BlipVisionConfig, Blip2VisionConfig]):
129
130
131
132
133
134
135
136
137
138
139
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
        super().__init__()

        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.class_embedding = nn.Parameter(torch.randn(1, 1, self.embed_dim))

        self.patch_embedding = nn.Conv2d(
            in_channels=3,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
        )

        self.num_patches = get_blip_num_patches(image_size=self.image_size,
                                                patch_size=self.patch_size)
        self.num_positions = self.num_patches + 1

        self.position_embedding = nn.Parameter(
            torch.randn(1, self.num_positions, self.embed_dim))

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        batch_size = pixel_values.shape[0]
        target_dtype = self.patch_embedding.weight.dtype
        patch_embeds = self.patch_embedding(pixel_values.to(
            dtype=target_dtype))  # shape = [*, width, grid, grid]
        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

        class_embeds = self.class_embedding.expand(batch_size, 1, -1)
        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)

        position_embeds = self.position_embedding.to(target_dtype)
        embeddings = embeddings + position_embeds[:, :embeddings.size(1), :]

        return embeddings


168
class BlipAttention(nn.Module):
169
170
171
172
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(
        self,
173
        config: Union[BlipVisionConfig, Blip2VisionConfig],
174
        quant_config: Optional[QuantizationConfig] = None,
175
176
        prefix: str = "",
    ) -> None:
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                "embed_dim must be divisible by num_heads "
                f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
                f" {self.num_heads}).")
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_dropout

        self.qkv = QKVParallelLinear(
            self.embed_dim,
            self.head_dim,
            self.num_heads,
            bias=config.qkv_bias,
            quant_config=quant_config,
196
            prefix=f"{prefix}.qkv",
197
198
199
200
201
        )
        self.projection = RowParallelLinear(
            self.embed_dim,
            self.embed_dim,
            quant_config=quant_config,
202
            prefix=f"{prefix}.projection",
203
204
205
206
207
        )

        self.tp_size = get_tensor_model_parallel_world_size()
        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)

208
209
210
211
212
213
        # Detect attention implementation.
        self.attn_backend = get_vit_attn_backend(support_fa=False)
        if self.attn_backend not in {_Backend.TORCH_SDPA, _Backend.XFORMERS}:
            raise RuntimeError(
                f"BLIP does not support {self.attn_backend} backend now.")

214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads,
                           self.head_dim).transpose(1, 2).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
    ):
        """Input shape: Batch x Time x Channel"""
        bsz, tgt_len, _ = hidden_states.size()

        qkv_states, _ = self.qkv(hidden_states)
        query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
        query_states = query_states.view(bsz, tgt_len,
                                         self.num_heads_per_partition,
                                         self.head_dim)
        key_states = key_states.view(bsz, tgt_len,
                                     self.num_heads_per_partition,
                                     self.head_dim)
        value_states = value_states.view(bsz, tgt_len,
                                         self.num_heads_per_partition,
                                         self.head_dim)

237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
        if self.attn_backend == _Backend.XFORMERS:
            from xformers import ops as xops

            out = xops.memory_efficient_attention_forward(query_states,
                                                          key_states,
                                                          value_states,
                                                          p=self.dropout,
                                                          scale=self.scale)
        elif self.attn_backend == _Backend.TORCH_SDPA:
            query_states, key_states, value_states = (x.transpose(1, 2)
                                                      for x in (query_states,
                                                                key_states,
                                                                value_states))
            out = F.scaled_dot_product_attention(query_states,
                                                 key_states,
                                                 value_states,
                                                 dropout_p=self.dropout,
                                                 scale=self.scale)
            out = out.transpose(1, 2)

257
258
259
        out = out.view(bsz, tgt_len, -1)
        attn_output, _ = self.projection(out)

260
        return attn_output, None
261
262


263
264
class BlipMLP(nn.Module):

265
266
267
268
269
270
    def __init__(
        self,
        config: BlipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
271
272
273
274
275
276
277
278
        super().__init__()

        self.config = config

        self.activation_fn = get_act_fn(config.hidden_act)
        self.fc1 = ColumnParallelLinear(config.hidden_size,
                                        config.intermediate_size,
                                        bias=True,
279
280
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.fc1")
281
282
283
        self.fc2 = RowParallelLinear(config.intermediate_size,
                                     config.hidden_size,
                                     bias=True,
284
285
                                     quant_config=quant_config,
                                     prefix=f"{prefix}.fc2")
286
287
288
289
290
291
292
293
294
295
296

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states, _ = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states, _ = self.fc2(hidden_states)

        return hidden_states


class BlipEncoderLayer(nn.Module):

297
298
299
300
301
302
    def __init__(
        self,
        config: BlipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
303
304
        super().__init__()

305
        # fallback to sdpa attention if tp unavailable
306
307
308
309
310
        self.self_attn = BlipAttention(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
311
312
        self.layer_norm1 = nn.LayerNorm(config.hidden_size,
                                        eps=config.layer_norm_eps)
313
314
315
        self.mlp = BlipMLP(config,
                           quant_config=quant_config,
                           prefix=f"{prefix}.mlp")
316
317
318
319
320
321
322
        self.layer_norm2 = nn.LayerNorm(config.hidden_size,
                                        eps=config.layer_norm_eps)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residual = hidden_states

        hidden_states = self.layer_norm1(hidden_states)
323
        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states


class BlipEncoder(nn.Module):
    """
    Transformer encoder consisting of `config.num_hidden_layers` self 
    attention layers. Each layer is a [`BlipEncoderLayer`].

    Args:
        config: BlipConfig
    """

343
344
345
346
347
348
349
    def __init__(
        self,
        config: BlipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        num_hidden_layers_override: Optional[int] = None,
        prefix: str = "",
    ) -> None:
350
351
352
353
354
355
356
357
358
359
        super().__init__()

        self.config = config

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

        self.layers = nn.ModuleList([
360
361
362
363
            BlipEncoderLayer(config=config,
                             quant_config=quant_config,
                             prefix=f"{prefix}.layers.{layer_idx}")
            for layer_idx in range(num_hidden_layers)
364
365
366
367
368
369
370
371
372
373
374
375
376
377
        ])

    def forward(self, inputs_embeds: torch.Tensor):
        hidden_states = inputs_embeds
        for encoder_layer in self.layers:
            hidden_states = encoder_layer(hidden_states)

        return hidden_states


class BlipVisionModel(nn.Module):
    config_class = BlipVisionConfig
    main_input_name = "pixel_values"

378
379
380
381
382
383
384
385
386
    def __init__(
        self,
        config: BlipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        *,
        num_hidden_layers_override: Optional[int] = None,
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
    ) -> None:
387
388
389
390
391
392
393
394
        super().__init__()
        self.config = config

        self.embeddings = BlipVisionEmbeddings(config)
        self.encoder = BlipEncoder(
            config=config,
            quant_config=quant_config,
            num_hidden_layers_override=num_hidden_layers_override,
395
            prefix=f"{prefix}.encoder",
396
        )
397

398
        num_hidden_layers = config.num_hidden_layers
399
400
        if len(self.encoder.layers) > config.num_hidden_layers:
            raise ValueError(
401
                f"The original encoder only has {num_hidden_layers} "
402
403
                f"layers, but you requested {len(self.encoder.layers)} layers."
            )
404
405
406
407
408
409

        # If possible, skip post_layernorm to conserve memory
        if require_post_norm is None:
            require_post_norm = len(self.encoder.layers) == num_hidden_layers

        if require_post_norm:
410
411
412
413
            self.post_layernorm = nn.LayerNorm(config.hidden_size,
                                               eps=config.layer_norm_eps)
        else:
            self.post_layernorm = None
414
415
416
417
418

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        hidden_states = self.embeddings(pixel_values)
        hidden_states = self.encoder(inputs_embeds=hidden_states)

419
420
421
        if self.post_layernorm is None:
            return hidden_states

422
        return self.post_layernorm(hidden_states)
423

424
425
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
426
427
428
429
430
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
431
        ]
432
        params_dict = dict(self.named_parameters())
433
        loaded_params: Set[str] = set()
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
        layer_count = len(self.encoder.layers)

        for name, loaded_weight in weights:
            # post_layernorm is not needed in BlipVisionModel
            if (name.startswith("post_layernorm")
                    and self.post_layernorm is None):
                continue

            # omit layers when num_hidden_layers_override is set
            if name.startswith("encoder.layers"):
                layer_idx = int(name.split(".")[2])
                if layer_idx >= layer_count:
                    continue

            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
451
452
                name = name.replace(weight_name, param_name)
                param = params_dict[name]
453
454
455
456
457
458
459
460
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
461
462
            loaded_params.add(name)
        return loaded_params