Upload processor

modeling_yangjian.py

@@ -26,7 +26,8 @@ class YangJianConfig(Qwen2_5_VLConfig):
         super().__init__(**kwargs)
         self.vision_config.compare_token_size = 100
         self.architectures = ["YangJianVLForConditionalGeneration"]
-
+        self.sequence_compare = False
+        
 class YangJianProcessor(Qwen2_5_VLProcessor):
     config_class = YangJianConfig
     def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs):
@@ -152,7 +153,6 @@ class OptimizedCrossAttention(nn.Module):
         self.dim = config.hidden_size
         self.num_heads = config.num_heads
         self.head_dim = self.dim // self.num_heads
-        self.num_key_value_groups = 1  # 对于 cross attention，通常设为 1
         self.scaling = self.head_dim**-0.5
         self.attention_dropout = 0.0
         self.is_causal = False  # cross attention 不需要因果掩码
@@ -173,103 +173,108 @@ class OptimizedCrossAttention(nn.Module):
         query_states: torch.Tensor,
         key_value_states: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,   # 只FA2用
+        kv_cu_seqlens: Optional[torch.Tensor] = None,# 只FA2用
         **kwargs,
     ) -> torch.Tensor:
-        """
-        Args:
-            query_states: [seq_len_q, hidden_size] 或 [batch_size, seq_len_q, hidden_size]
-            key_value_states: [seq_len_kv, hidden_size] 或 [batch_size, seq_len_kv, hidden_size]
-                             如果为 None，则执行 self attention
-        """
-        # 处理输入维度
+        # 允许 query_states [B,T,d] 或 [T,d]，自动扩展 batch 维
+        orig_2d = False
         if query_states.dim() == 2:
-            query_states = query_states.unsqueeze(0)  # [1, seq_len_q, hidden_size]
-            squeeze_output = True
-        else:
-            squeeze_output = False
-        
+            query_states = query_states.unsqueeze(0)
+            orig_2d = True
+
         batch_size, seq_len_q, _ = query_states.shape
-        
+
+        # Q/K/V投影
         if self.is_cross_attention and key_value_states is not None:
-            # Cross Attention
             if key_value_states.dim() == 2:
-                key_value_states = key_value_states.unsqueeze(0)  # [1, seq_len_kv, hidden_size]
-            
-            # 计算 Q
-            q = self.q_proj(query_states)  # [batch_size, seq_len_q, hidden_size]
-            
-            # 计算 K、V（融合计算）
-            kv = self.kv(key_value_states)  # [batch_size, seq_len_kv, hidden_size * 2]
+                key_value_states = key_value_states.unsqueeze(0)
+            q = self.q_proj(query_states)
+            kv = self.kv(key_value_states)
             seq_len_kv = kv.shape[1]
-            
-            # 分离 K、V
             k, v = kv.reshape(batch_size, seq_len_kv, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
-            # k, v: [batch_size, num_heads, seq_len_kv, head_dim]
-            
-            # 重塑 Q
             q = q.reshape(batch_size, seq_len_q, self.num_heads, self.head_dim).transpose(1, 2)
-            # q: [batch_size, num_heads, seq_len_q, head_dim]
-            
         else:
-            # Self Attention
             if key_value_states is None:
                 key_value_states = query_states
+            qkv = self.qkv(query_states)
+            q, k, v = qkv.reshape(batch_size, seq_len_q, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+
+        # 选用哪个 attention kernel
+        attn_impl = getattr(self.config, '_attn_implementation', 'sdpa')
+        attn_impl = 'sdpa'
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[attn_impl]
+
+        # ========= 支持 FA2 ==========
+        if attn_impl == "flash_attention_2":
+            # Qwen2_5 之所以能支持 FA2，是因为准备了 flatten+cu_seqlens
+            # 这里假设 query_states/key_value_states 按 batch 维是变长的
+
+            # 检查 cu_seqlens，有就用，否则尝试自动生成
+            if cu_seqlens is None:
+                # 默认把每个batch都视为长度=seq_len_q
+                cu_seqlens = torch.arange(0, (batch_size + 1) * seq_len_q, step=seq_len_q, dtype=torch.int32, device=q.device)
+            if kv_cu_seqlens is None:
+                cu_seqlens_k = torch.arange(0, (batch_size + 1) * k.shape[2], step=k.shape[2], dtype=torch.int32, device=k.device)
+            else:
+                cu_seqlens_k = kv_cu_seqlens
+
+            # flatten [B, nH, T, d] -> [total_T, nH, d]
+            # 注意！FlashAttn2是 (total, nH, d)，不是 (nH, total, d)，和普通实现不一样
+            # 更安全的 flatten 方式
+            # [B, nH, T, d] -> [B, T, nH, d] -> [total_T, nH, d]
+            q_ = q.transpose(1, 2).contiguous().view(-1, self.num_heads, self.head_dim)
+            k_ = k.transpose(1, 2).contiguous().view(-1, self.num_heads, self.head_dim)
+            v_ = v.transpose(1, 2).contiguous().view(-1, self.num_heads, self.head_dim)
             
-            # 融合计算 Q、K、V
-            qkv = self.qkv(query_states)  # [batch_size, seq_len, hidden_size * 3]
+            max_seqlen_q = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+            max_seqlen_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).max().item()
+
+            attn_output, _ = attention_interface(
+                self,
+                q_,
+                k_,
+                v_,
+                attention_mask=None,
+                scaling=self.scaling,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                cu_seq_lens_q=cu_seqlens,
+                cu_seq_lens_k=cu_seqlens_k,
+                max_length_q=max_seqlen_q,
+                max_length_k=max_seqlen_k,
+                is_causal=self.is_causal,
+                **kwargs,
+            )
             
-            # 分离 Q、K、V
-            q, k, v = qkv.reshape(batch_size, seq_len_q, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
-            # q, k, v: [batch_size, num_heads, seq_len, head_dim]
-        
-        # 选择 attention 实现
-        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS["sdpa"]
-        # if hasattr(self.config, '_attn_implementation') and self.config._attn_implementation != "eager":
-        #     attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-        
-        # 构造 cu_seqlens 参数（FlashAttention 必需）
-        cu_seqlens_q = torch.arange(0, (batch_size*self.num_heads + 1) * seq_len_q, step=seq_len_q, dtype=torch.int32, device=q.device)
-        if self.is_cross_attention and key_value_states is not None:
-            cu_seqlens_k = torch.arange(0, (batch_size*self.num_heads + 1) * seq_len_kv, step=seq_len_kv, dtype=torch.int32, device=k.device)
+            # 更简洁的输出重构
+            # [total_q, nH, d] -> [B, seq_len_q, nH, d]
+            attn_output = attn_output.view(batch_size, seq_len_q, self.num_heads, self.head_dim).contiguous()
         else:
-            cu_seqlens_k = cu_seqlens_q
-        
-        # 执行 attention 计算
-        attn_output, _ = attention_interface(
-            self,
-            q,
-            k, 
-            v,
-            attention_mask=attention_mask,
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            max_seqlen_q=seq_len_q,
-            max_seqlen_k=seq_len_kv if self.is_cross_attention and key_value_states is not None else seq_len_q,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            is_causal=False,
-            **kwargs,
-        )
-        
-        attn_output = attn_output.reshape(batch_size, self.num_heads, seq_len_q, self.head_dim)
-        
-        attn_output = attn_output.transpose(1, 2).contiguous()  # [batch_size, seq_len_q, num_heads, head_dim]
-        attn_output = attn_output.reshape(batch_size, seq_len_q, self.dim)  # [batch_size, seq_len_q, hidden_size]
-        
-        # 输出投影
+            # 普通实现，下游实现就是 [B, nH, T, d]
+            attn_output, _ = attention_interface(
+                self,
+                q, k, v,
+                attention_mask=attention_mask,
+                scaling=self.scaling,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                is_causal=self.is_causal,
+                **kwargs,
+            )
+            # attn_output: [B, nH, seq_q, d]
+            attn_output = attn_output.transpose(1, 2).contiguous()  # [B, seq_q, nH, d]
+
+        attn_output = attn_output.reshape(batch_size, seq_len_q, self.dim)  # [B, seq_q, D]
         attn_output = self.proj(attn_output)
-        
-        # 如果输入是 2D，则输出也应该是 2D
-        if squeeze_output:
-            attn_output = attn_output.squeeze(0)  # [seq_len_q, hidden_size]
-        
-        return attn_output
+        if orig_2d:
+            attn_output = attn_output.squeeze(0)
+        return attn_output.contiguous()
 
 
 class YangJianCompareVisualEncoder(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
+        self.sequence_compare = getattr(config, "sequence_compare", False)
         self.hidden_size = config.hidden_size
         # self.token_size = 100  * (config.spatial_merge_size**2) if "compare_token_size" not in config else config.compare_token_size  * (config.spatial_merge_size**2)
         self.token_size = 100 if "compare_token_size" not in config else config.compare_token_size
@@ -291,7 +296,6 @@ class YangJianCompareVisualEncoder(nn.Module):
         self.query_embeddings = nn.Parameter(
             torch.empty(self.token_size, self.hidden_size)
         )
-        
         # 只保留 Cross Attention for queries to attend to encoded features
         self.decoder_cross_attn = OptimizedCrossAttention(config, is_cross_attention=True)
         
@@ -301,33 +305,8 @@ class YangJianCompareVisualEncoder(nn.Module):
 
         self.compare_projector = nn.Linear(config.hidden_size, config.out_hidden_size)
 
-    def _ensure_device_dtype_consistency(self, target_tensor):
-        """
-        确保所有模块组件都在目标张量的设备上并使用相同的数据类型
-        """
-        device = target_tensor.device
-        dtype = target_tensor.dtype
-        
-        # 移动 attention 模块到正确设备
-        self.encoder_cross_attn1 = self.encoder_cross_attn1.to(device=device, dtype=dtype)
-        self.encoder_cross_attn2 = self.encoder_cross_attn2.to(device=device, dtype=dtype)
-        self.decoder_cross_attn = self.decoder_cross_attn.to(device=device, dtype=dtype)
-        
-        # 移动 norm 层到正确设备
-        self.encoder_norm1 = self.encoder_norm1.to(device=device, dtype=dtype)
-        self.encoder_norm2 = self.encoder_norm2.to(device=device, dtype=dtype)
-        self.encoder_norm3 = self.encoder_norm3.to(device=device, dtype=dtype)
-        self.encoder_norm4 = self.encoder_norm4.to(device=device, dtype=dtype)
-        self.decoder_norm1 = self.decoder_norm1.to(device=device, dtype=dtype)
-        self.decoder_norm2 = self.decoder_norm2.to(device=device, dtype=dtype)
-        
-        # 移动 MLP 到正确设备
-        self.encoder_mlp1 = self.encoder_mlp1.to(device=device, dtype=dtype)
-        self.encoder_mlp2 = self.encoder_mlp2.to(device=device, dtype=dtype)
-        self.decoder_mlp = self.decoder_mlp.to(device=device, dtype=dtype)
-
-    def _initialize_weights(self):
-        nn.init.normal_(self.query_embeddings.weight, mean=0.0, std=0.02)
+    def init_query_embeddings(self):
+        nn.init.normal_(self.query_embeddings, mean=0.0, std=0.02)
 
     def forward(self, images_hidden_states: list) -> torch.Tensor:
         """
@@ -340,13 +319,10 @@ class YangJianCompareVisualEncoder(nn.Module):
         if not images_hidden_states:
             return torch.empty(0, self.token_size, self.hidden_size)
         
-        # 确保所有组件的设备和数据类型一致
-        # self._ensure_device_dtype_consistency(images_hidden_states[0])
-        
         # 检查 query_embeddings 是否包含 NaN
         if torch.isnan(self.query_embeddings).any():
-            print("警告：query_embeddings 包含 NaN 值，重新初始化")
-            nn.init.normal_(self.query_embeddings, mean=0.0, std=0.02)
+            print("警告：query_embeddings 包含 NaN 值")
+            # nn.init.normal_(self.query_embeddings, mean=0.0, std=0.02)
         
         # 获取每个图像的序列长度
         seq_lengths = [state.size(0) for state in images_hidden_states]
@@ -380,9 +356,11 @@ class YangJianCompareVisualEncoder(nn.Module):
         # 创建循环移位的状态用于对比
         # 对于第一个图像，使用自身作为previous
         previous_states = torch.roll(batched_states, shifts=1, dims=0)
-        previous_states[0] = batched_states[0]
         previous_masks = torch.roll(attention_masks, shifts=1, dims=0)
-        previous_masks[0] = attention_masks[0]
+
+        if previous_states.size(0) > 1 and self.sequence_compare:
+            previous_states[0] = previous_states[1]
+            previous_masks[0] = previous_masks[1]
         
         # Encoder: 批量处理所有图像
         encoded_features = self._encoder_forward(
@@ -759,4 +737,5 @@ class YangJianVLForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
     
     def __init__(self, config):
         super().__init__(config)
-        self.model = YangJianVLModel(config)
+        self.model = YangJianVLModel(config)
+