Delete visual_modeling_baichuan.py

visual_modeling_baichuan.py

@@ -1,166 +0,0 @@
-
-from typing import List, Optional, Tuple, Union
-import torch, math
-import torch.utils.checkpoint
-from torch import nn
-import transformers
-from flash_attn import flash_attn_varlen_func
-from transformers.activations import ACT2FN
-from PIL import Image
-import io, fire
-from torch.nn import functional as F
-
-
-class BaichuanVisualAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
-
-    def __init__(self, config):
-        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(
-                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
-                f" {self.num_heads})."
-            )
-        self.scale = self.head_dim**-0.5  # flash attention不需要使用
-        self.dropout = config.attention_dropout
-
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        # print("*******BaichuanVisualAttention monkey patched!!*******")
-
-        # initializer transformer from_pretrain下不生效
-        factor = self.config.initializer_factor
-        in_proj_std = (self.embed_dim**-0.5) * ((2 * self.config.num_hidden_layers) ** -0.5) * factor
-        out_proj_std = (self.embed_dim**-0.5) * factor
-        nn.init.normal_(self.q_proj.weight, std=in_proj_std)
-        nn.init.normal_(self.k_proj.weight, std=in_proj_std)
-        nn.init.normal_(self.v_proj.weight, std=in_proj_std)
-        nn.init.normal_(self.out_proj.weight, std=out_proj_std)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        causal_attention_mask: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = False,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        """Input shape: Batch x Time x Channel"""
-
-        bsz, tgt_len, embed_dim = hidden_states.size()
-        src_len = tgt_len
-
-        query_states = self.q_proj(hidden_states).view(bsz * tgt_len, self.num_heads, self.head_dim)
-        key_states = self.k_proj(hidden_states).view(bsz * tgt_len, self.num_heads, self.head_dim)
-        value_states = self.v_proj(hidden_states).view(bsz * tgt_len, self.num_heads, self.head_dim)
-
-        # 暂时不考虑变长patch nums 固定长度为256/1024
-        cu_len = torch.arange(0, (bsz + 1) * tgt_len, step=tgt_len, dtype=torch.int32, device=query_states.device)
-        # print(self.config.s2a, self.config.rope_scaling, cu_len, torch.sum(cu_len), q_len, kv_seq_len)
-        # 如果不是f16 bf16不用flash attn
-        if query_states.dtype in [torch.float16, torch.bfloat16]:
-            attn_output = flash_attn_varlen_func(query_states, key_states, value_states, cu_len, cu_len, tgt_len, tgt_len, causal=False)  # (bsz * qlen, nheads, headdim)
-            attn_output = attn_output.view(bsz, tgt_len, self.num_heads, self.head_dim)
-        else:
-            with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
-                attn_output = torch.nn.functional.scaled_dot_product_attention(query_states, key_states, value_states, attention_mask, 0.0)
-                attn_output = attn_output.transpose(1, 2)    
-        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
-        attn_output = self.out_proj(attn_output)
-
-        return attn_output, None
-
-# monkey patch for flash attention
-# transformers.models.siglip.modeling_siglip.SiglipAttention = BaichuanVisualAttention
-from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VisionTransformerPretrainedModel
-from qwen_vl_utils import process_vision_info
-
-class BaichuanVisualEncoder(transformers.models.qwen2_vl.modeling_qwen2_vl.Qwen2VisionTransformerPretrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.gradient_checkpointing = True  # 强制开启
-        self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint
-        del self.merger
-        
-    def forward(
-        self,
-        pixel_values: torch.Tensor, 
-        grid_thw: torch.Tensor,
-    ):
-        hidden_states = pixel_values.to(self.get_dtype())
-        grid_thw = grid_thw.to(pixel_values.device)
-        
-        hidden_states = self.patch_embed(hidden_states)
-        rotary_pos_emb = self.rot_pos_emb(grid_thw)
-
-        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), value=0)
-
-        for blk in self.blocks:
-            if self.gradient_checkpointing and self.training:
-                hidden_states = self._gradient_checkpointing_func(blk.__call__, hidden_states, cu_seqlens, rotary_pos_emb)
-            else:
-                hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb)
-
-        return hidden_states
-    
-    @torch.no_grad()
-    def fake_input(self, device, merge_size=2):
-        merge_size = max(merge_size, self.config.spatial_merge_size)
-        fake_image = [torch.zeros([
-            1,
-            self.config.temporal_patch_size,
-            3,
-            merge_size // self.config.spatial_merge_size,
-            self.config.spatial_merge_size,
-            self.config.patch_size,
-            merge_size // self.config.spatial_merge_size,
-            self.config.spatial_merge_size,
-            self.config.patch_size,
-        ], dtype=torch.float32, device=device)]
-        return fake_image
-
-
-class BaichuanVisualBridge(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.merge_size = self.config.merge_size if hasattr(self.config, 'merge_size') else 2
-        self.hidden_size = config.embed_dim * (self.merge_size**2)
-        self.ln_q = nn.LayerNorm(config.embed_dim, eps=1e-6)
-        self.mlp = nn.Sequential(
-            nn.Linear(self.hidden_size, self.hidden_size),
-            nn.GELU(),
-            nn.Linear(self.hidden_size, config.hidden_size),
-        )
-        
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
-        return x
-    
-
-def test_vision():
-    from transformers.models.clip.modeling_clip import CLIPPreTrainedModel
-    from transformers import AutoConfig
-    config = AutoConfig.from_pretrained("./", trust_remote_code=True)
-
-    ae = BaichuanVisualEncoder(config.visual_config).cuda().to(torch.bfloat16)
-    bg = BaichuanVisualBridge(config).cuda().to(torch.bfloat16)
-    print(ae)
-    pixel_input = torch.rand([4, 3, config.visual_config.image_size, config.visual_config.image_size], dtype=torch.float32).cuda()
-
-    visual_embedding = ae(pixel_input)[0][:, 1:]  # 删除class token
-    visual_proj = bg(visual_embedding)
-    print(visual_proj.shape)
-    print(ae.fake_input(visual_proj.device))
-
-if __name__ == '__main__':
-    fire.Fire()
-