Ankit2802/phi3_vision_128k

1# coding=utf-82# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15import warnings16 17import torch18from torch import nn19from transformers import CLIPVisionConfig, CLIPVisionModel, PretrainedConfig20from transformers.models.clip.modeling_clip import CLIPAttention21from transformers.utils import logging22 23try:24    from flash_attn import flash_attn_func25except ImportError:26    pass27 28logger = logging.get_logger(__name__)29 30 31MAX_INPUT_ID = int(1e9)32 33CLIP_VIT_LARGE_PATCH14_336_CONFIG = CLIPVisionConfig(34  attention_dropout=0.0,35  dropout=0.0,36  hidden_act="quick_gelu",37  hidden_size=1024,38  image_size=336,39  initializer_factor=1.0,40  initializer_range=0.02,41  intermediate_size=4096,42  layer_norm_eps=1e-05,43  num_attention_heads=16,44  num_channels=3,45  num_hidden_layers=24,46  patch_size=14,47  projection_dim=76848)49 50class CLIPAttentionFA2(CLIPAttention):51    """Add flash attention 2 to CLIPAttention. (This is only used in the vision encoder)"""52 53    def forward(self,54        hidden_states,55        attention_mask=None,56        causal_attention_mask=None,57        output_attentions=False,58    ):59        """Input shape: Batch x Time x Channel"""60 61        assert attention_mask is None, "CLIPAttentionFA2 does not support attention_mask"62        assert causal_attention_mask is None, "CLIPAttentionFA2 does not support causal_attention_mask"63        assert output_attentions is False, "CLIPAttentionFA2 does not support output_attentions"64 65        bsz, tgt_len, embed_dim = hidden_states.size()66        query_states = self.q_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)67        key_states = self.k_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)68        value_states = self.v_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)69 70        attn_output = flash_attn_func(71            query_states,72            key_states,73            value_states,74            dropout_p=self.dropout if self.training else 0.0,75            softmax_scale=self.scale,76            causal=False,77        ).reshape(bsz, tgt_len, embed_dim)78 79        attn_output = self.out_proj(attn_output)80        return attn_output, None81 82 83class Phi3ImageEmbedding(nn.Module):84    """Phi3 Image embedding."""85 86    def __init__(self, config: PretrainedConfig, wte=None, **kwargs) -> None:87        super().__init__()88 89        # n_embed or hidden_size90        hidden_size = config.n_embd if hasattr(config, 'n_embd') else config.hidden_size91        if hasattr(config, 'embd_pdrop') or hasattr(config, 'embed_pdrop'):92            embd_drop = config.embd_pdrop if hasattr(config, 'embd_pdrop') else config.embed_pdrop93            self.drop = nn.Dropout(embd_drop)94        else:95            self.drop = None96 97        self.wte = wte98 99        if isinstance(config.img_processor, dict) and config.img_processor.get('name', None) == 'clip_vision_model':100            assert 'model_name' in config.img_processor, 'model_name must be provided for CLIPVisionModel'101            assert 'image_dim_out' in config.img_processor, 'image_dim_out must be provided for CLIPVisionModel'102            assert 'num_img_tokens' in config.img_processor, 'num_img_tokens must be provided for CLIPVisionModel'103            assert config.img_processor['model_name'] == 'openai/clip-vit-large-patch14-336'104            clip_config = CLIP_VIT_LARGE_PATCH14_336_CONFIG105            self.img_processor = CLIPVisionModel(clip_config)106            image_dim_out = config.img_processor['image_dim_out']107            self.num_img_tokens = config.img_processor['num_img_tokens']108 109            # FA2 in CLIP110            if config._attn_implementation == 'flash_attention_2':111                for layer in self.img_processor.vision_model.encoder.layers:112                    clip_fa2 = CLIPAttentionFA2(clip_config)113                    del layer.self_attn114                    layer.self_attn = clip_fa2115        else:116            raise NotImplementedError(f'img_processor = {config.img_processor}, not implemented')117 118        self.image_dim_out = image_dim_out119        self.img_sizes = None120 121        # global_gn and sub_gn for hd transform, serves as line separator122        self.use_hd_transform = kwargs.get('use_hd_transform', False)123        self.with_learnable_separator = kwargs.get('with_learnable_separator', False)124        self.hd_transform_order = kwargs.get('hd_transform_order', 'glb_sub')125        # with_hd_transform and with_learnable_separator should have same value126        assert self.use_hd_transform == self.with_learnable_separator, 'use_hd_transform and with_learnable_separator should have same value'127        if self.with_learnable_separator:128            assert self.use_hd_transform, 'learnable separator is only for hd transform'129            # 1024 * 4, merge spatial to channel dimension130            self.glb_GN = nn.Parameter(torch.zeros([1, 1, self.image_dim_out * 4]))131            self.sub_GN = nn.Parameter(torch.zeros([1, 1, 1, self.image_dim_out * 4]))132            logger.info(f'learnable separator enabled for hd transform, hd_transform_order = {self.hd_transform_order}')133 134        projection_cls = kwargs.get('projection_cls', 'linear')135        if projection_cls == 'linear':136            self.img_projection = nn.Linear(image_dim_out, hidden_size)137        elif projection_cls == 'mlp' and self.use_hd_transform:138            dim_projection = hidden_size139            depth = 2140            layers = [nn.Linear(image_dim_out * 4, dim_projection)]141            for _ in range(1, depth):142                layers.extend([nn.GELU(),143                                nn.Linear(dim_projection, dim_projection)])144            self.img_projection = nn.Sequential(*layers)145        elif projection_cls == 'mlp':146            dim_projection = hidden_size147            depth = 2148            layers = [nn.Linear(image_dim_out, dim_projection)]149            for _ in range(1, depth):150                layers.extend([nn.GELU(),151                                nn.Linear(dim_projection, dim_projection)])152            self.img_projection = nn.Sequential(*layers)153        else:154            raise NotImplementedError(f'projection_cls = {projection_cls}, not implemented')155 156        self.vocab_size = config.vocab_size157        self.img_features = None158 159        if isinstance(config.img_processor, dict):160            self.layer_idx = config.img_processor.get('layer_idx', -2)161            self.type_feature = config.img_processor.get('type_feature', 'patch')162        else:163            self.layer_idx = -2164            self.type_feature = 'patch'165 166 167    def set_img_features(self, img_features: torch.FloatTensor) -> None:168        self.img_features = img_features169 170    def set_img_sizes(self, img_sizes: torch.LongTensor) -> None:171        self.img_sizes = img_sizes172 173    def get_img_features(self, img_embeds: torch.FloatTensor) -> torch.FloatTensor:174        LAYER_IDX = self.layer_idx175        TYPE_FEATURE = self.type_feature176 177        img_processor_output = self.img_processor(img_embeds, output_hidden_states=True)178        img_feature = img_processor_output.hidden_states[LAYER_IDX]179 180        if TYPE_FEATURE == "patch":181            patch_feature = img_feature[:, 1:]182            return patch_feature183 184        raise NotImplementedError185 186    def forward(187        self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes=None188    ) -> torch.FloatTensor:189        input_shape = input_ids.size()190        input_ids = input_ids.view(-1, input_shape[-1])191 192        # positions for image tokens193        positions = torch.nonzero((input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=True)194        has_image = len(positions[0].tolist()) > 0195        # input_ids = input_ids.clamp_min(0).clamp_max(self.vocab_size).detach()196        input_ids.clamp_min_(0).clamp_max_(self.vocab_size)197        warnings.warn(198            "Phi-3-V modifies `input_ids` in-place and the tokens indicating images will be "199            "removed after model forward. If your workflow requires multiple forward passes on "200            "the same `input_ids`, please make a copy of `input_ids` before passing it to the "201            "model."202        )203 204        hidden_states = self.wte(input_ids)205 206        if has_image:207            assert self.use_hd_transform208            num_images, num_crops, c, h, w = pixel_values.shape209            assert c == 3 and h == w == 336210            img_features = self.get_img_features(pixel_values.flatten(0, 1)).reshape(211                num_images, num_crops, -1, self.image_dim_out212            )213            image_features_proj = self.hd_feature_transform(img_features, image_sizes)214            hidden_states = hidden_states.index_put(215                positions, image_features_proj, accumulate=False216            )217 218        if self.drop is not None:219            hidden_states = self.drop(hidden_states)220 221        return hidden_states222 223    def hd_feature_transform(self, image_features, image_sizes):224        """225        image_features: (num_images, num_crops+1, 24*24, 1024)226        """227        assert (228            self.hd_transform_order == 'sub_glb'229        ), f'hd_transform_order `{self.hd_transform_order}` not implemented'230        if isinstance(self.img_projection, nn.Sequential):231            target_device = self.img_projection[0].bias.device232            target_dtype = self.img_projection[0].bias.dtype233        else:  # It's a single nn.Linear layer234            target_device = self.img_projection.bias.device235            target_dtype = self.img_projection.bias.dtype236 237        global_image_features = image_features[:, 0]  # (num_images, 24*24, 1024)238        # global feature can be viewed as a special HD case with num_crops 1x1239        global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)240        global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)241 242        all_image_embeddings = []243        # need a for loop to process each image because of different image sizes244        # (patch arrangement is different for each image)245        for i, img_size in enumerate(image_sizes):246            h, w = img_size247            h_crop = h // 336248            w_crop = w // 336249            num_crops = h_crop * w_crop250 251            # NOTE: real num_crops is padded252            # (num_crops, 24*24, 1024)253            sub_image_features = image_features[i, 1 : 1 + num_crops]254            sub_image_features_hd = self.reshape_hd_patches_2x2merge(255                sub_image_features, h_crop, w_crop256            )257            sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)258 259            # [sub features, separator, global features]260            all_image_embeddings.extend(261                [262                    sub_image_features_hd_newline.squeeze(0),  # (h_crop*12*(w_crop*12+1), 4096)263                    self.glb_GN.squeeze(0),264                    global_image_features_hd_newline[i],265                ]266            )267 268        image_features_proj = self.img_projection(269            torch.cat(all_image_embeddings, dim=0).to(target_device).to(target_dtype)270        )271 272        return image_features_proj273 274    def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):275        """276        image_features: (num_images*num_crops, 24*24, 1024)277        output: (num_images, h_crop*12, w_crop*12, 4096), h_crop*w_crop == num_crops278        """279        N, L, C = image_features.shape280        assert L == 24 * 24 and C == 1024 and N % (h_crop * w_crop) == 0281        num_images = N // (h_crop * w_crop)282        H = int(L**0.5)283        image_features_hd = (284            image_features.reshape(N, H, H, C)  # N, 24, 24, 1024285            .reshape(N, H // 2, 2, H // 2, 2, C)  # N, 12, 2, 12, 2, 1024286            .permute(0, 1, 3, 2, 4, 5)  # N, 12, 12, 2, 2, 1024287            .reshape(N, -1, 4 * C)  # N, 144, 4096288            .reshape(289                num_images, h_crop, w_crop, H // 2, H // 2, -1290            )  # n_img, h_crop, w_crop, 12, 12, 4096291            .permute(0, 1, 3, 2, 4, 5)  # n_img, h_crop, 12, w_crop, 12, 4096292            .reshape(293                num_images, h_crop * H // 2, w_crop * H // 2, 4 * C294            )  # n_img, h_crop*12, w_crop*12, 4096295        )296 297        # alternative implementation using einops298        # from einops import rearrange299        # image_features_nhwc = rearrange(300        #     image_features,301        #     'N (H W) c -> N H W c',302        #     H=H,303        #     W=H,304        # )305        # image_features_2x2merge = rearrange(306        #     image_features_nhwc,307        #     'N (h h_pool) (w w_pool) c -> N h w (h_pool w_pool c)',308        #     h_pool=2,309        #     w_pool=2,310        # )311        # image_features_hd = rearrange(312        #     image_features_2x2merge,313        #     '(n_img h_crop w_crop) h w C -> n_img (h_crop h) (w_crop w) C',314        #     h_crop=h_crop,315        #     w_crop=w_crop,316        # )317 318        return image_features_hd319 320    def add_image_newline(self, image_features_hd):321        """322        image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)323        output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)324        """325        num_images, h, w, hid_dim = image_features_hd.shape326        # add the newline token to the HD image feature patches327        newline_embeddings = self.sub_GN.expand(num_images, h, -1, -1)  # (n_img, h, 1, hid_dim)328        image_features_hd_newline = torch.cat(329            [image_features_hd, newline_embeddings], dim=2330        ).reshape(num_images, -1, hid_dim)331        return image_features_hd_newline332