siglip.py 20.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
3
4
5
"""Implementation of SiglipVisionModel intended to be only used
within a vision language model."""

import math
6
from typing import Iterable, Optional, Set, Tuple, Union
7
8
9
10

import torch
from PIL import Image
from torch import nn
11
from transformers import SiglipVisionConfig
12

13
from vllm.attention.layer import MultiHeadAttention
14
from vllm.distributed import divide, get_tensor_model_parallel_world_size
15
16
17
18
19
20
21
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.vocab_parallel_embedding import (
    VocabParallelEmbedding)
22
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23
from vllm.multimodal.utils import consecutive_placeholder_ranges
24
from vllm.sequence import SequenceData
zhuwenwen's avatar
zhuwenwen committed
25
import vllm.envs as envs
26

27
from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs
28

29
30

def get_siglip_patch_grid_length(*, image_size: int, patch_size: int) -> int:
31
32
    # Since interpolation is applied, the image size need not be divisible
    # assert image_size % patch_size == 0
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
    return image_size // patch_size


def get_siglip_num_patches(*, image_size: int, patch_size: int) -> int:
    grid_length = get_siglip_patch_grid_length(image_size=image_size,
                                               patch_size=patch_size)
    return grid_length * grid_length


def get_siglip_image_feature_size(hf_config: SiglipVisionConfig) -> int:
    return get_siglip_num_patches(image_size=hf_config.image_size,
                                  patch_size=hf_config.patch_size)


def get_max_siglip_image_tokens(hf_config: SiglipVisionConfig) -> int:
    return get_siglip_image_feature_size(hf_config)


def dummy_seq_data_for_siglip(
    hf_config: SiglipVisionConfig,
    seq_len: int,
54
    num_images: int,
55
56
57
    *,
    image_token_id: int,
    image_feature_size_override: Optional[int] = None,
58
    mm_key: str = "image",
59
60
61
62
63
64
):
    if image_feature_size_override is None:
        image_feature_size = get_siglip_image_feature_size(hf_config)
    else:
        image_feature_size = image_feature_size_override

65
    return SequenceData.from_prompt_token_counts(
66
67
        (image_token_id, image_feature_size * num_images),
        (0, seq_len - image_feature_size * num_images),
68
69
70
71
72
    ), {
        mm_key:
        consecutive_placeholder_ranges(num_items=num_images,
                                       item_size=image_feature_size)
    }
73
74
75
76


def dummy_image_for_siglip(
    hf_config: SiglipVisionConfig,
77
    num_images: int,
78
79
80
81
82
83
84
85
86
87
88
    *,
    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)
89
    return {"image": image if num_images == 1 else [image] * num_images}
90
91


92
93
94
95
96
97
98
99
100
101
102
103
104
class SiglipEncoderInfo(VisionEncoderInfo[SiglipVisionConfig]):

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
        return get_siglip_image_feature_size(self.vision_config)

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

105
106
107
108
109
110
111
    def get_image_size(self) -> int:
        return self.vision_config.image_size

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

    def get_patch_grid_length(self) -> int:
112
113
114
115
116
117
        return get_siglip_patch_grid_length(
            image_size=self.vision_config.image_size,
            patch_size=self.vision_config.patch_size,
        )


118
119
120
121
122
123
124
125
126
127
128
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
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
201
202
203
204
205
206
207
208
209
210
# Adapted from https://github.com/huggingface/transformers/blob/v4.43.3/src/transformers/models/siglip/modeling_siglip.py#L249 # noqa
class SiglipVisionEmbeddings(nn.Module):

    def __init__(self, config: SiglipVisionConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.patch_embedding = nn.Conv2d(
            in_channels=config.num_channels,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
            padding="valid",
        )

        self.num_patches = (self.image_size // self.patch_size)**2
        self.num_positions = self.num_patches
        self.position_embedding = VocabParallelEmbedding(
            self.num_positions, self.embed_dim)
        self.register_buffer(
            "position_ids",
            torch.arange(self.num_positions, dtype=torch.int64).expand(
                (1, -1)),
            persistent=False,
        )

    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int,
                                 width: int) -> torch.Tensor:
        """
        This method is an adapted method for SigLIP (due to SigLIP not having
        class embedding unlike other ViTs) that allows the model to interpolate
        the pre-trained position encodings such that it can be usable on higher
        resolution images.

        Source:
        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
        """
        position_embeddings = self.position_embedding.weight.unsqueeze(0)
        num_patches = embeddings.shape[1]
        num_positions = position_embeddings.shape[1]
        if num_patches == num_positions and height == width:
            return position_embeddings

        dim = embeddings.shape[-1]
        height = height // self.patch_size
        width = width // self.patch_size
        # we add a small number to avoid floating point error
        # in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        height, width = height + 0.1, width + 0.1

        patch_pos_embed = position_embeddings.reshape(
            1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)),
            dim)
        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed,
            scale_factor=(
                height / math.sqrt(num_positions),
                width / math.sqrt(num_positions),
            ),
            mode="bicubic",
            align_corners=False,
        )
        if (int(height) != patch_pos_embed.shape[-2]
                or int(width) != patch_pos_embed.shape[-1]):
            raise ValueError("Width or height does not match with "
                             "the interpolated position embeddings")

        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return patch_pos_embed

    def forward(self,
                pixel_values: torch.Tensor,
                interpolate_pos_encoding: bool = False) -> torch.Tensor:
        _, _, height, width = pixel_values.shape
        target_dtype = self.patch_embedding.weight.dtype
        patch_embeds = self.patch_embedding(pixel_values.to(
            dtype=target_dtype))  # shape = [*, width, grid, grid]
        embeddings = patch_embeds.flatten(2).transpose(1, 2)

        if interpolate_pos_encoding:
            embeddings = embeddings + self.interpolate_pos_encoding(
                embeddings, height, width)
        else:
            embeddings = embeddings + self.position_embedding(
                self.position_ids)
        return embeddings


211
class SiglipAttention(nn.Module):
212
213
214

    def __init__(
        self,
215
        config: SiglipVisionConfig,
216
        quant_config: Optional[QuantizationConfig] = None,
217
218
        prefix: str = "",
    ) -> None:
219
        super().__init__()
220

221
222
        self.config = config
        self.embed_dim = config.hidden_size
223
224
225
        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:
226
227
228
            raise ValueError(f"embed_dim must be divisible by num_heads (got "
                             "`embed_dim`: {self.embed_dim} and `num_heads`:"
                             f" {self.num_heads}).")
229

230
231
232
233
234
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_dropout
        self.qkv_proj = QKVParallelLinear(
            hidden_size=self.embed_dim,
            head_size=self.head_dim,
235
            total_num_heads=self.num_heads,
236
            quant_config=quant_config,
237
            prefix=f"{prefix}.qkv_proj",
238
        )
239

240
241
242
243
        self.out_proj = RowParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            quant_config=quant_config,
244
            prefix=f"{prefix}.out_proj",
245
246
        )

247
248
        self.tp_size = get_tensor_model_parallel_world_size()
        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
249

250
251
        self.attn = MultiHeadAttention(self.num_heads_per_partition,
                                       self.head_dim, self.scale)
252

253
254
255
256
257
258
    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        """Input shape: Batch x Time x Channel"""
        qkv_states, _ = self.qkv_proj(hidden_states)
259
260
        query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)

261
        out = self.attn(query_states, key_states, value_states)
262
        attn_output, _ = self.out_proj(out)
263

264
        return attn_output, None
265
266
267
268
269
270


class SiglipMLP(nn.Module):

    def __init__(
        self,
271
        config: SiglipVisionConfig,
272
        quant_config: Optional[QuantizationConfig] = None,
273
274
        prefix: str = "",
    ) -> None:
275
        super().__init__()
276

277
278
        self.config = config
        self.activation_fn = get_act_fn(config.hidden_act)
279
280
281
282
        # Special handling for BNB quantization
        if quant_config and quant_config.get_name() == "bitsandbytes":
            quantizable = True
        else:
283
            # For other quantization, we require the hidden size to be a
284
            # multiple of 64
285
286
            quantizable = (config.hidden_size % 64 == 0
                           and config.intermediate_size % 64 == 0)
287
288
289
290
        self.fc1 = ColumnParallelLinear(
            config.hidden_size,
            config.intermediate_size,
            quant_config=quant_config if quantizable else None,
291
            prefix=f"{prefix}.fc1",
292
293
294
295
296
        )
        self.fc2 = RowParallelLinear(
            config.intermediate_size,
            config.hidden_size,
            quant_config=quant_config if quantizable else None,
297
            prefix=f"{prefix}.fc2",
298
299
300
301
302
303
304
305
306
307
308
309
310
        )

    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 SiglipEncoderLayer(nn.Module):

    def __init__(
        self,
311
        config: SiglipVisionConfig,
312
        quant_config: Optional[QuantizationConfig] = None,
313
314
        prefix: str = "",
    ) -> None:
315
316
317
        super().__init__()

        self.embed_dim = config.hidden_size
318

319
320
321
322
323
        self.self_attn = SiglipAttention(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
324
325
326
327
328
        self.layer_norm1 = nn.LayerNorm(self.embed_dim,
                                        eps=config.layer_norm_eps)
        self.mlp = SiglipMLP(
            config,
            quant_config=quant_config,
329
            prefix=f"{prefix}.mlp",
330
331
332
333
334
335
336
        )
        self.layer_norm2 = nn.LayerNorm(self.embed_dim,
                                        eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
337
    ) -> Tuple[torch.Tensor, None]:
338
339
340
        residual = hidden_states

        hidden_states = self.layer_norm1(hidden_states)
341
        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
342
343
344
345
346
347
348
349
350
351
352
353
354
355
        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, None


class SiglipEncoder(nn.Module):

    def __init__(
        self,
356
        config: SiglipVisionConfig,
357
        quant_config: Optional[QuantizationConfig] = None,
358
        num_hidden_layers_override: Optional[int] = None,
359
360
        prefix: str = "",
    ) -> None:
361
        super().__init__()
362

363
        self.config = config
364
365
366
367
368
369

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

370
        self.layers = nn.ModuleList([
371
372
373
374
            SiglipEncoderLayer(config,
                               quant_config=quant_config,
                               prefix=f"{prefix}.layers.{layer_idx}")
            for layer_idx in range(num_hidden_layers)
375
376
377
378
379
        ])

    def forward(
        self,
        inputs_embeds: torch.Tensor,
380
381
        return_all_hidden_states: bool,
    ) -> Union[torch.Tensor, list[torch.Tensor]]:
382
        hidden_states_pool = [inputs_embeds]
383
        hidden_states = inputs_embeds
384

385
386
        for encoder_layer in self.layers:
            hidden_states, _ = encoder_layer(hidden_states)
387
388
389
390
391
392
            if return_all_hidden_states:
                hidden_states_pool.append(hidden_states)
        # If we have multiple feature sample layers, we return all hidden
        # states in order and grab the ones we need by index.
        if return_all_hidden_states:
            return hidden_states_pool
393
394
395
396
397
398
399
400
401
402
        return hidden_states


class SiglipMultiheadAttentionPoolingHead(nn.Module):
    """Multihead Attention Pooling."""

    def __init__(
        self,
        config: SiglipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
403
404
        prefix: str = "",
    ) -> None:
405
406
407
408
409
410
411
412
        super().__init__()

        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
        # TODO(ChristopherCho): Implement vLLM version of MultiheadAttention
        self.attention = torch.nn.MultiheadAttention(
            config.hidden_size, config.num_attention_heads, batch_first=True)
        self.layernorm = nn.LayerNorm(config.hidden_size,
                                      eps=config.layer_norm_eps)
413
414
415
        self.mlp = SiglipMLP(config=config,
                             quant_config=quant_config,
                             prefix=f"{prefix}.mlp")
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        batch_size = hidden_state.shape[0]
        probe = self.probe.repeat(batch_size, 1, 1)

        hidden_state = self.attention(probe, hidden_state, hidden_state)[0]

        residual = hidden_state
        hidden_state = self.layernorm(hidden_state)
        hidden_state = residual + self.mlp(hidden_state)

        return hidden_state[:, 0]


class SiglipVisionTransformer(nn.Module):

    def __init__(
        self,
        config: SiglipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
436
        *,
437
        num_hidden_layers_override: Optional[int] = None,
438
439
440
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
    ) -> None:
441
        super().__init__()
442

443
444
445
446
        self.config = config
        embed_dim = config.hidden_size

        self.embeddings = SiglipVisionEmbeddings(config)
447

448
449
450
        self.encoder = SiglipEncoder(
            config,
            quant_config=quant_config,
451
            num_hidden_layers_override=num_hidden_layers_override,
452
            prefix=f"{prefix}.encoder",
453
        )
454

455
        num_hidden_layers = config.num_hidden_layers
456
457
        if len(self.encoder.layers) > config.num_hidden_layers:
            raise ValueError(
458
                f"The original encoder only has {num_hidden_layers} "
459
460
                f"layers, but you requested {len(self.encoder.layers)} layers."
            )
461
462
463
464
465
466

        # 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:
467
468
469
            self.post_layernorm = nn.LayerNorm(embed_dim,
                                               eps=config.layer_norm_eps)
        else:
470
471
            self.post_layernorm = None

472
473
474
475
        self.use_head = (True if not hasattr(config, "vision_use_head") else
                         config.vision_use_head)
        if self.use_head:
            self.head = SiglipMultiheadAttentionPoolingHead(
476
477
478
479
                config=config,
                quant_config=quant_config,
                prefix=f"{prefix}.head",
            )
480
481
482
483
484

    def forward(
        self,
        pixel_values: torch.Tensor,
        interpolate_pos_encoding: bool = True,
485
        feature_sample_layers: Optional[list[int]] = None,
486
    ) -> torch.Tensor:
487

488
489
490
491
492
        hidden_states = self.embeddings(
            pixel_values,
            interpolate_pos_encoding=interpolate_pos_encoding,
        )

493
        return_all_hidden_states = feature_sample_layers is not None
494

495
496
497
498
499
500
        # Produces either the last layer output or all of the hidden states,
        # depending on if we have feature_sample_layers or not
        encoder_outputs = self.encoder(
            inputs_embeds=hidden_states,
            return_all_hidden_states=return_all_hidden_states,
        )
501

502
503
504
505
        # Handle post-norm (if applicable) and stacks feature layers if needed
        encoder_outputs = resolve_visual_encoder_outputs(
            encoder_outputs, feature_sample_layers, self.post_layernorm,
            self.config.num_hidden_layers)
506

507
        # TODO: add this back when pooled_output is used in inference.
508
        # if self.use_head:
509
        # pooled_output = self.head(encoder_outputs)
510

511
        return encoder_outputs
512
513
514
515
516
517
518
519
520
521


class SiglipVisionModel(nn.Module):
    config_class = SiglipVisionConfig
    main_input_name = "pixel_values"

    def __init__(
        self,
        config: SiglipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
522
        *,
523
        num_hidden_layers_override: Optional[int] = None,
524
525
526
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
    ) -> None:
527
        super().__init__()
528

529
530
531
        self.vision_model = SiglipVisionTransformer(
            config,
            quant_config,
532
            num_hidden_layers_override=num_hidden_layers_override,
533
534
            require_post_norm=require_post_norm,
            prefix=f"{prefix}.vision_model",
535
536
537
538
539
540
541
542
543
        )

    def get_input_embeddings(self) -> nn.Module:
        return self.vision_model.embeddings.patch_embedding

    def forward(
        self,
        pixel_values: torch.Tensor,
        interpolate_pos_encoding: bool = False,
544
        feature_sample_layers: Optional[list[int]] = None,
545
546
547
548
    ) -> torch.Tensor:
        return self.vision_model(
            pixel_values=pixel_values,
            interpolate_pos_encoding=interpolate_pos_encoding,
549
            feature_sample_layers=feature_sample_layers,
550
        )
551

552
553
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
554
555
556
557
558
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
559
        ]
560
        params_dict = dict(self.named_parameters())
561
        loaded_params: Set[str] = set()
562
563
564
        layer_count = len(self.vision_model.encoder.layers)

        for name, loaded_weight in weights:
565
            # post_layernorm is optional in SiglipVisionModel
566
567
            if (name.startswith("vision_model.post_layernorm")
                    and self.vision_model.post_layernorm is None):
568
569
                continue

570
            # omit layers when num_hidden_layers_override is set
571
            if name.startswith("vision_model.encoder.layers"):
572
573
574
575
                layer_idx = int(name.split(".")[3])
                if layer_idx >= layer_count:
                    continue

576
577
578
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
579
                name = name.replace(weight_name, param_name)
580

581
                param = params_dict[name]
582
583
584
585
586
587
588
589
                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)
590
591
            loaded_params.add(name)
        return loaded_params