Upload folder using huggingface_hub

configuration_monkeyocrv2vit.py

@@ -1,5 +1,8 @@
 from typing import Any, Optional
 from transformers.configuration_utils import PretrainedConfig
+# Remove Qwen3Config import as it causes error on older transformers/python versions
+# from transformers.models.qwen3 import Qwen3Config 
+#from transformers import Qwen2_5_VLProcessor, AutoProcessor
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING
 
 
@@ -8,8 +11,8 @@ class MonkeyOCRv2VisionConfig(PretrainedConfig):
 
     def __init__(
         self,
-        embed_dim: int = 1536,
-        hidden_size: int = 1536,
+        embed_dim: int = 1536,  # vision encoder embed size
+        hidden_size: int = 1536,  # after merger hidden size
         intermediate_size: int = 4224,
         num_hidden_layers: int = 42,
         num_attention_heads: int = 12,
@@ -22,7 +25,7 @@ class MonkeyOCRv2VisionConfig(PretrainedConfig):
         vision_attn_implementation="flash_attention_2",  # "eager","sdpa","flash_attention_2"
         initializer_range=0.02,
         init_merger_std=0.02,
-        is_causal=False,
+        is_causal=False,  # ve causal forward
         post_norm=True,
         gradient_checkpointing=False,
         **kwargs: Any,
@@ -46,4 +49,13 @@ class MonkeyOCRv2VisionConfig(PretrainedConfig):
         self.post_norm = post_norm
         self.gradient_checkpointing = gradient_checkpointing
 
+
+# Commented out Processor definition to avoid dependencies on Qwen2_5_VLProcessor
+# class MonkeyOCRv2Processor(Qwen2_5_VLProcessor):
+#     attributes = ["image_processor", "tokenizer"]
+#     def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
+#         super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+
+# AutoProcessor.register("MonkeyOCRv2VisionTransformer", MonkeyOCRv2Processor)
 CONFIG_MAPPING.register("MonkeyOCRv2VisionTransformer", MonkeyOCRv2VisionConfig)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c61ccbd8db2e0f655565d8502256f9fb75f2f24491e8b0d2f2f9a28ca4c354a7
+oid sha256:a9f059fc5a7afb6e9d2ad64b5db02dfc8e846d4b713460f4b3c0b2468a4d52e1
 size 454882592

modeling_monkeyocrv2_vision.py

@@ -16,6 +16,7 @@ except ImportError:
 from torch.nn import LayerNorm
 from transformers.modeling_utils import PreTrainedModel
 from .configuration_monkeyocrv2vit import MonkeyOCRv2VisionConfig
+# from configuration_monkeyocrv2vit import MonkeyOCRv2VisionConfig
 
 
 try:
@@ -28,7 +29,7 @@ except ImportError:
 def rotate_half(x):
     """Rotates half the hidden dims of the input."""
     x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2:]
+    x2 = x[..., x.shape[-1] // 2:] #这里是q0和q(d/2)一组，而不是q0和q1
     return torch.cat((-x2, x1), dim=-1)  
 
 
@@ -194,6 +195,8 @@ class VisionAttentionV2(nn.Module):
         q_list = torch.split(q, seqlens, 0)
         k_list = torch.split(k, seqlens, 0)
         v_list = torch.split(v, seqlens, 0)
+        # eager attention  空间复杂度为 O(n^2) , n 为  b*s（batch_size * seq_len）,  序列太长容易OOM， 这个实现 更具batch 切分 seq
+        # 减少内存需求， 计算相对 continus  batching 较慢。
         outputs = []
         for q_i, k_i, v_i in zip(q_list, k_list, v_list):
             q_i = q_i.transpose(0, 1)
@@ -274,18 +277,21 @@ class VisionSdpaAttention(nn.Module):
         for i in range(1, len(cu_seqlens)):
             attention_mask[..., cu_seqlens[i - 1]: cu_seqlens[i], cu_seqlens[i - 1]: cu_seqlens[i]] = True
 
-        q = q.transpose(0, 1).unsqueeze(0)
+        # Convert q, k, v to 4D to enable : (1, num_heads, seq_length, head_dim)
+        q = q.transpose(0, 1).unsqueeze(0)   # (1, num_heads, seq_length, head_dim)
         k = k.transpose(0, 1).unsqueeze(0)
         v = v.transpose(0, 1).unsqueeze(0)
 
+        # See: https://github.com/pytorch/pytorch/issues/127523
         if attention_mask.stride(-1) != 1:
             attention_mask = torch.empty_like(attention_mask, memory_format=torch.contiguous_format).copy_(attention_mask)
 
+        # use memory efficient backend
         from torch.nn.attention import SDPBackend, sdpa_kernel
         with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
             attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
 
-        attn_output = attn_output.squeeze(0).transpose(0, 1)
+        attn_output = attn_output.squeeze(0).transpose(0, 1)  # (seq_length, num_heads, head_dim)
         attn_output = attn_output.reshape(seq_length, -1)
 
         attn_output = self.proj(attn_output)
@@ -294,10 +300,10 @@ class VisionSdpaAttention(nn.Module):
 
 VISION_ATTENTION_CLASSES = {
     "eager": VisionAttention,
-    "eager_v2": VisionAttentionV2,
+    "eager_v2": VisionAttentionV2,  # 内存更少
     "flash_attention_2": VisionFlashAttention2,
     "sdpa": VisionSdpaAttention,
-    "ascend_fa": VisionAscendAttention,
+    "ascend_fa": VisionAscendAttention,  # ascend， 长序列精度下降严重。
 }
 
 
@@ -404,6 +410,132 @@ class VisionBlock(nn.Module):
         return hidden_states
 
 
+class VisionTransformerDecoder(PreTrainedModel):
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _no_split_modules = ["VisionBlock"]
+    def __init__(self, config: MonkeyOCRv2VisionConfig) -> None:
+        super().__init__(config)
+        self.num_classes = 3 * config.patch_size ** 2
+
+        self.spatial_merge_size = config.spatial_merge_size
+        self.patch_size = config.patch_size
+
+        head_dim = config.embed_dim // config.num_attention_heads
+
+        self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
+
+        _num_hidden_layers = config.num_hidden_layers
+
+        self.blocks = nn.ModuleList(
+            [VisionBlock(config, config.vision_attn_implementation) for _ in range(_num_hidden_layers)]
+        )
+
+        if self.config.post_norm:
+            self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
+        self.head = nn.Linear(config.embed_dim, self.num_classes) if self.num_classes > 0 else nn.Identity()
+
+        self.gradient_checkpointing = False
+        self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear, nn.Conv3d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.blocks[0].mlp.fc2.weight.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.blocks[0].mlp.fc2.weight.device
+
+    def get_pos_ids_by_grid(self, grid_thw):
+        pos_ids = []
+        for t, h, w in grid_thw:
+            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+            hpos_ids = hpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+            hpos_ids = hpos_ids.flatten()
+
+            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+            wpos_ids = wpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+            wpos_ids = wpos_ids.flatten()
+            
+            pos_ids.append(
+                torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)
+            )
+
+
+
+        return pos_ids
+
+    def rot_pos_emb(self, grid_thw):
+        pos_ids = self.get_pos_ids_by_grid(grid_thw) #得到在旋转编码表中的hw坐标
+        pos_ids = torch.cat(pos_ids, dim=0)
+        max_grid_size = grid_thw[:, 1:].max()
+        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) # max_size x dim/2
+
+        # # mrope
+        # rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) 
+
+
+        # mrope-i
+        emb = rotary_pos_emb_full[pos_ids]
+        rotary_pos_emb = torch.stack([emb[:,0], emb[:,1]], dim=2).reshape(emb.shape[0], -1)
+        # mrope-i
+
+        return rotary_pos_emb
+    
+    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, bf16=True) -> torch.Tensor:
+        # if bf16:
+        #     hidden_states = hidden_states.bfloat16()
+
+        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=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+        )#得到每个图像的token起始
+        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) #左侧补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)
+
+        if self.config.post_norm:
+            hidden_states = self.post_trunk_norm(hidden_states)
+
+        hidden_states = self.head(hidden_states)
+
+        return hidden_states
+
+
 class MonkeyOCRv2VisionTransformer(PreTrainedModel):
     config_class = MonkeyOCRv2VisionConfig
     _supports_flash_attn = True
@@ -485,13 +617,20 @@ class MonkeyOCRv2VisionTransformer(PreTrainedModel):
         return pos_ids
 
     def rot_pos_emb(self, grid_thw):
-        pos_ids = self.get_pos_ids_by_grid(grid_thw)
+        pos_ids = self.get_pos_ids_by_grid(grid_thw) #得到在旋转编码表中的hw坐标
         pos_ids = torch.cat(pos_ids, dim=0)
         max_grid_size = grid_thw[:, 1:].max()
-        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) # max_size x dim/2
+
+        # # mrope
+        # rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) 
 
+
+        # mrope-i
         emb = rotary_pos_emb_full[pos_ids]
         rotary_pos_emb = torch.stack([emb[:,0], emb[:,1]], dim=2).reshape(emb.shape[0], -1)
+        # mrope-i
+
         return rotary_pos_emb
 
     def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, bf16=True) -> torch.Tensor:
@@ -505,8 +644,8 @@ class MonkeyOCRv2VisionTransformer(PreTrainedModel):
         cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
             dim=0,
             dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
-        )
-        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+        )#得到每个图像的token起始
+        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) #左侧补1个值，右侧补0个值，这个值的value是0
 
         for blk in self.blocks:
             if self.gradient_checkpointing and self.training: