MathLLMs/MathCoder-VL-2B

1# --------------------------------------------------------2# InternVL3# Copyright (c) 2023 OpenGVLab4# Licensed under The MIT License [see LICENSE for details]5# --------------------------------------------------------6from typing import Optional, Tuple, Union7 8import torch9import torch.nn.functional as F10import torch.utils.checkpoint11from einops import rearrange12from timm.models.layers import DropPath13from torch import nn14from transformers.activations import ACT2FN15from transformers.modeling_outputs import (BaseModelOutput,16                                           BaseModelOutputWithPooling)17from transformers.modeling_utils import PreTrainedModel18from transformers.utils import logging19 20from .configuration_intern_vit import InternVisionConfig21 22try:23    try:  # v124        from flash_attn.flash_attn_interface import \25            flash_attn_unpadded_qkvpacked_func26    except:  # v227        from flash_attn.flash_attn_interface import \28            flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func29 30    from flash_attn.bert_padding import pad_input, unpad_input31 32    has_flash_attn = True33except:34    print('FlashAttention is not installed.')35    has_flash_attn = False36 37logger = logging.get_logger(__name__)38 39 40class FlashAttention(nn.Module):41    """Implement the scaled dot product attention with softmax.42    Arguments43    ---------44        softmax_scale: The temperature to use for the softmax attention.45                      (default: 1/sqrt(d_keys) where d_keys is computed at46                      runtime)47        attention_dropout: The dropout rate to apply to the attention48                           (default: 0.0)49    """50 51    def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):52        super().__init__()53        self.softmax_scale = softmax_scale54        self.dropout_p = attention_dropout55 56    def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None,57                max_s=None, need_weights=False):58        """Implements the multihead softmax attention.59        Arguments60        ---------61            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None62                if unpadded: (nnz, 3, h, d)63            key_padding_mask: a bool tensor of shape (B, S)64        """65        assert not need_weights66        assert qkv.dtype in [torch.float16, torch.bfloat16]67        assert qkv.is_cuda68 69        if cu_seqlens is None:70            batch_size = qkv.shape[0]71            seqlen = qkv.shape[1]72            if key_padding_mask is None:73                qkv = rearrange(qkv, 'b s ... -> (b s) ...')74                max_s = seqlen75                cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,76                                          device=qkv.device)77                output = flash_attn_unpadded_qkvpacked_func(78                    qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,79                    softmax_scale=self.softmax_scale, causal=causal80                )81                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)82            else:83                nheads = qkv.shape[-2]84                x = rearrange(qkv, 'b s three h d -> b s (three h d)')85                x_unpad, indices, cu_seqlens, max_s = unpad_input(x, key_padding_mask)86                x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)87                output_unpad = flash_attn_unpadded_qkvpacked_func(88                    x_unpad, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,89                    softmax_scale=self.softmax_scale, causal=causal90                )91                output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),92                                             indices, batch_size, seqlen),93                                   'b s (h d) -> b s h d', h=nheads)94        else:95            assert max_s is not None96            output = flash_attn_unpadded_qkvpacked_func(97                qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,98                softmax_scale=self.softmax_scale, causal=causal99            )100 101        return output, None102 103 104class InternRMSNorm(nn.Module):105    def __init__(self, hidden_size, eps=1e-6):106        super().__init__()107        self.weight = nn.Parameter(torch.ones(hidden_size))108        self.variance_epsilon = eps109 110    def forward(self, hidden_states):111        input_dtype = hidden_states.dtype112        hidden_states = hidden_states.to(torch.float32)113        variance = hidden_states.pow(2).mean(-1, keepdim=True)114        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)115        return self.weight * hidden_states.to(input_dtype)116 117 118try:119    from apex.normalization import FusedRMSNorm120 121    InternRMSNorm = FusedRMSNorm  # noqa122 123    logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternRMSNorm')124except ImportError:125    # using the normal InternRMSNorm126    pass127except Exception:128    logger.warning('discovered apex but it failed to load, falling back to InternRMSNorm')129    pass130 131 132NORM2FN = {133    'rms_norm': InternRMSNorm,134    'layer_norm': nn.LayerNorm,135}136 137 138class InternVisionEmbeddings(nn.Module):139    def __init__(self, config: InternVisionConfig):140        super().__init__()141        self.config = config142        self.embed_dim = config.hidden_size143        self.image_size = config.image_size144        self.patch_size = config.patch_size145 146        self.class_embedding = nn.Parameter(147            torch.randn(1, 1, self.embed_dim),148        )149 150        self.patch_embedding = nn.Conv2d(151            in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size152        )153 154        self.num_patches = (self.image_size // self.patch_size) ** 2155        self.num_positions = self.num_patches + 1156 157        self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))158 159    def _get_pos_embed(self, pos_embed, H, W):160        target_dtype = pos_embed.dtype161        pos_embed = pos_embed.float().reshape(162            1, self.image_size // self.patch_size, self.image_size // self.patch_size, -1).permute(0, 3, 1, 2)163        pos_embed = F.interpolate(pos_embed, size=(H, W), mode='bicubic', align_corners=False). \164            reshape(1, -1, H * W).permute(0, 2, 1).to(target_dtype)165        return pos_embed166 167    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:168        target_dtype = self.patch_embedding.weight.dtype169        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, channel, width, height]170        batch_size, _, height, width = patch_embeds.shape171        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)172        class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)173        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)174        position_embedding = torch.cat([175            self.position_embedding[:, :1, :],176            self._get_pos_embed(self.position_embedding[:, 1:, :], height, width)177        ], dim=1)178        embeddings = embeddings + position_embedding.to(target_dtype)179        return embeddings180 181 182class InternAttention(nn.Module):183    """Multi-headed attention from 'Attention Is All You Need' paper"""184 185    def __init__(self, config: InternVisionConfig):186        super().__init__()187        self.config = config188        self.embed_dim = config.hidden_size189        self.num_heads = config.num_attention_heads190        self.use_flash_attn = config.use_flash_attn and has_flash_attn191        if config.use_flash_attn and not has_flash_attn:192            print('Warning: Flash Attention is not available, use_flash_attn is set to False.')193        self.head_dim = self.embed_dim // self.num_heads194        if self.head_dim * self.num_heads != self.embed_dim:195            raise ValueError(196                f'embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:'197                f' {self.num_heads}).'198            )199 200        self.scale = self.head_dim ** -0.5201        self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)202        self.attn_drop = nn.Dropout(config.attention_dropout)203        self.proj_drop = nn.Dropout(config.dropout)204 205        self.qk_normalization = config.qk_normalization206 207        if self.qk_normalization:208            self.q_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)209            self.k_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)210 211        if self.use_flash_attn:212            self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)213        self.proj = nn.Linear(self.embed_dim, self.embed_dim)214 215    def _naive_attn(self, x):216        B, N, C = x.shape217        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)218        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)219 220        if self.qk_normalization:221            B_, H_, N_, D_ = q.shape222            q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)223            k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)224 225        attn = ((q * self.scale) @ k.transpose(-2, -1))226        attn = attn.softmax(dim=-1)227        attn = self.attn_drop(attn)228 229        x = (attn @ v).transpose(1, 2).reshape(B, N, C)230        x = self.proj(x)231        x = self.proj_drop(x)232        return x233 234    def wk_naive_attn(self, x):235        B, N, C = x.shape236        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)237        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)238 239        if self.qk_normalization:240            B_, H_, N_, D_ = q.shape241            q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)242            k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)243 244        attn = ((q * self.scale) @ k.transpose(-2, -1))245        attn = attn.softmax(dim=-1)246        attn = self.attn_drop(attn)247        print(attn.shape)248        return attn249 250    def _flash_attn(self, x, key_padding_mask=None, need_weights=False):251        qkv = self.qkv(x)252        qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)253 254        if self.qk_normalization:255            q, k, v = qkv.unbind(2)256            q = self.q_norm(q.flatten(-2, -1)).view(q.shape)257            k = self.k_norm(k.flatten(-2, -1)).view(k.shape)258            qkv = torch.stack([q, k, v], dim=2)259 260        context, _ = self.inner_attn(261            qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=False262        )263        outs = self.proj(rearrange(context, 'b s h d -> b s (h d)'))264        outs = self.proj_drop(outs)265        return outs266 267    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:268        x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states)269        return x270 271 272class InternMLP(nn.Module):273    def __init__(self, config: InternVisionConfig):274        super().__init__()275        self.config = config276        self.act = ACT2FN[config.hidden_act]277        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)278        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)279 280    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:281        hidden_states = self.fc1(hidden_states)282        hidden_states = self.act(hidden_states)283        hidden_states = self.fc2(hidden_states)284        return hidden_states285 286 287class InternVisionEncoderLayer(nn.Module):288    def __init__(self, config: InternVisionConfig, drop_path_rate: float):289        super().__init__()290        self.embed_dim = config.hidden_size291        self.intermediate_size = config.intermediate_size292        self.norm_type = config.norm_type293 294        self.attn = InternAttention(config)295        self.mlp = InternMLP(config)296        self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)297        self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)298 299        self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))300        self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))301        self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()302        self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()303 304    def forward(305            self,306            hidden_states: torch.Tensor,307    ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:308        """309        Args:310            hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`311        """312        hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states)) * self.ls1)313 314        hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states)) * self.ls2)315 316        return hidden_states317 318 319class InternVisionEncoder(nn.Module):320    """321    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a322    [`InternEncoderLayer`].323 324    Args:325        config (`InternConfig`):326            The corresponding vision configuration for the `InternEncoder`.327    """328 329    def __init__(self, config: InternVisionConfig):330        super().__init__()331        self.config = config332        # stochastic depth decay rule333        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]334        self.layers = nn.ModuleList([335            InternVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])336        self.gradient_checkpointing = True337 338    def forward(339            self,340            inputs_embeds,341            output_hidden_states: Optional[bool] = None,342            return_dict: Optional[bool] = None,343    ) -> Union[Tuple, BaseModelOutput]:344        r"""345        Args:346            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):347                Embedded representation of the inputs. Should be float, not int tokens.348            output_hidden_states (`bool`, *optional*):349                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors350                for more detail.351            return_dict (`bool`, *optional*):352                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.353        """354        output_hidden_states = (355            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states356        )357        return_dict = return_dict if return_dict is not None else self.config.use_return_dict358 359        encoder_states = () if output_hidden_states else None360        hidden_states = inputs_embeds361 362        for idx, encoder_layer in enumerate(self.layers):363            if output_hidden_states:364                encoder_states = encoder_states + (hidden_states,)365            if self.gradient_checkpointing and self.training:366                layer_outputs = torch.utils.checkpoint.checkpoint(367                    encoder_layer,368                    hidden_states)369            else:370                layer_outputs = encoder_layer(371                    hidden_states,372                )373            hidden_states = layer_outputs374 375        if output_hidden_states:376            encoder_states = encoder_states + (hidden_states,)377 378        if not return_dict:379            return tuple(v for v in [hidden_states, encoder_states] if v is not None)380        return BaseModelOutput(381            last_hidden_state=hidden_states, hidden_states=encoder_states382        )383 384 385class InternVisionModel(PreTrainedModel):386    main_input_name = 'pixel_values'387    config_class = InternVisionConfig388    _no_split_modules = ['InternVisionEncoderLayer']389 390    def __init__(self, config: InternVisionConfig):391        super().__init__(config)392        self.config = config393 394        self.embeddings = InternVisionEmbeddings(config)395        self.encoder = InternVisionEncoder(config)396 397    def resize_pos_embeddings(self, old_size, new_size, patch_size):398        pos_emb = self.embeddings.position_embedding399        _, num_positions, embed_dim = pos_emb.shape400        cls_emb = pos_emb[:, :1, :]401        pos_emb = pos_emb[:, 1:, :].reshape(1, old_size // patch_size, old_size // patch_size, -1).permute(0, 3, 1, 2)402        pos_emb = F.interpolate(pos_emb.float(), size=new_size // patch_size, mode='bicubic', align_corners=False)403        pos_emb = pos_emb.to(cls_emb.dtype).reshape(1, embed_dim, -1).permute(0, 2, 1)404        pos_emb = torch.cat([cls_emb, pos_emb], dim=1)405        self.embeddings.position_embedding = nn.Parameter(pos_emb)406        self.embeddings.image_size = new_size407        logger.info('Resized position embeddings from {} to {}'.format(old_size, new_size))408 409    def get_input_embeddings(self):410        return self.embeddings411 412    def forward(413            self,414            pixel_values: Optional[torch.FloatTensor] = None,415            output_hidden_states: Optional[bool] = None,416            return_dict: Optional[bool] = None,417            pixel_embeds: Optional[torch.FloatTensor] = None,418    ) -> Union[Tuple, BaseModelOutputWithPooling]:419        output_hidden_states = (420            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states421        )422        return_dict = return_dict if return_dict is not None else self.config.use_return_dict423 424        if pixel_values is None and pixel_embeds is None:425            raise ValueError('You have to specify pixel_values or pixel_embeds')426 427        if pixel_embeds is not None:428            hidden_states = pixel_embeds429        else:430            if len(pixel_values.shape) == 4:431                hidden_states = self.embeddings(pixel_values)432            else:433                raise ValueError(f'wrong pixel_values size: {pixel_values.shape}')434        encoder_outputs = self.encoder(435            inputs_embeds=hidden_states,436            output_hidden_states=output_hidden_states,437            return_dict=return_dict,438        )439        last_hidden_state = encoder_outputs.last_hidden_state440        pooled_output = last_hidden_state[:, 0, :]441 442        if not return_dict:443            return (last_hidden_state, pooled_output) + encoder_outputs[1:]444 445        return BaseModelOutputWithPooling(446            last_hidden_state=last_hidden_state,447            pooler_output=pooled_output,448            hidden_states=encoder_outputs.hidden_states,449            attentions=encoder_outputs.attentions,450        )451