Initial FanFormer checkpoint with architecture and README

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:607ebeab78df2738e7039a379d2e0b022cdf42f7ff9b675f38be06d91f72a160
-size 331514552
+oid sha256:39fb924d47a28f1b2456828b984b26bc5c05382f02b7fc878e067b729edfd88a
+size 331511816

model_architecture.py

@@ -353,6 +353,8 @@ class FANformerMultiheadAttention(nn.Module):
     Implementación de la atención multi-cabeza con FANformer.
     Aplica normalización a Q, K, V individualmente y utiliza unpadding para mejorar el rendimiento.
     Incorpora modelado de periodicidad a través de proyecciones CoLA_FAN.
+    [MODIFICADO] Se eliminó el escalado ssmax_scale y seq_scale de Q.
+    [MODIFICADO] Se aplica conversión explícita a bfloat16 *después* de las operaciones de normalización.
     """
     def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.12, use_rope: bool = True,
                  layer_index: int = 1, max_seq_len: int = 512, p: float = 0.15,
@@ -366,21 +368,26 @@ class FANformerMultiheadAttention(nn.Module):
         self.layer_index = layer_index
         self.use_pre_norm = use_pre_norm
         self.p = p  # Proporción para periodicidad
-        
+
         if embed_dim % num_heads != 0:
             raise ValueError("embed_dim debe ser divisible por num_heads")
-        
+
         self.head_dim = embed_dim // num_heads
         self.use_rope = use_rope
-        
+
         if num_gqa_groups is None:
             num_gqa_groups = num_heads
+        # Añadido chequeo de divisibilidad para GQA
+        elif num_heads % num_gqa_groups != 0:
+            raise ValueError("num_heads debe ser divisible por num_gqa_groups")
+
 
         try:
             from flash_attn import flash_attn_func, flash_attn_varlen_func
             self.flash_attn_func = flash_attn_func
             self.flash_attn_varlen_func = flash_attn_varlen_func
         except ImportError as e:
+            # Mantener el comportamiento original de lanzar error si no se encuentra
             raise ImportError(f"Error al inicializar FlashAttention: {e}")
 
         # Para el unpadding
@@ -389,20 +396,21 @@ class FANformerMultiheadAttention(nn.Module):
             self.unpad_input = unpad_input
             self.pad_input = pad_input
         except ImportError as e:
+            # Mantener el comportamiento original de lanzar error si no se encuentra
             raise ImportError(f"Error al importar funciones de padding: {e}")
 
-        # Inicialización de parámetros de escala
-        self.ssmax_scale = nn.Parameter(torch.ones(num_heads, dtype=torch.bfloat16) * 0.168)
-        nn.init.uniform_(self.ssmax_scale, a=0.166, b=0.170)
-        self.register_buffer('seq_scale', torch.log(torch.tensor(max_seq_len, dtype=torch.bfloat16)))
-        
+        # Eliminada la inicialización de parámetros de escala ssmax_scale y seq_scale
+        # self.ssmax_scale = nn.Parameter(torch.ones(num_heads, dtype=torch.bfloat16) * 0.168)
+        # nn.init.uniform_(self.ssmax_scale, a=0.166, b=0.170)
+        # self.register_buffer('seq_scale', torch.log(torch.tensor(max_seq_len, dtype=torch.bfloat16)))
+
         # Capas de normalización para la entrada (Pre-Norm en primer bloque o QKV-Norm para los demás)
         self.norm = nn.RMSNorm(embed_dim, eps=1e-5)
-        
+
         # Capas de dropout (simplificadas)
-        self.attention_dropout = progressive_dropout(dropout, depth=1)
+        self.attention_dropout = progressive_dropout(dropout, depth=layer_index) # Usar layer_index
         # Eliminado: self.projection_dropout = progressive_dropout(dropout * 1.1, depth=1)
-        self.output_dropout = progressive_dropout(dropout, depth=1)
+        self.output_dropout = progressive_dropout(dropout, depth=layer_index) # Usar layer_index
 
         # Proyecciones para Q, K, V usando GQAFANLinear (implementación FANformer)
         self.Wq = GQAFANLinear(embed_dim, embed_dim, num_heads, num_gqa_groups, p=p)
@@ -413,165 +421,172 @@ class FANformerMultiheadAttention(nn.Module):
         self.out_proj = CoLA_Linear(embed_dim, embed_dim, rank=embed_dim // 4)
 
     def scaled_dot_product_attention_flash_unpadded(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
-                                                   attention_mask: Optional[torch.Tensor] = None,
+                                                   attention_mask: Optional[torch.Tensor] = None, # Revertido a Optional
                                                    is_causal: bool = False) -> torch.Tensor:
         B, H, S, D = q.shape  # batch, heads, sequence length, head dimension
-        
+
+        # Mantener la lógica original de manejo de máscara opcional
         if attention_mask is None:
             # Si no hay máscara de atención, usamos la versión regular
             return self.scaled_dot_product_attention_flash(q, k, v, mask=None, is_causal=is_causal)
-        
+
         # Convertir las tensiones a [B, S, H, D] para unpad_input
         q_unpad = q.permute(0, 2, 1, 3)  # [B, S, H, D]
         k_unpad = k.permute(0, 2, 1, 3)  # [B, S, H, D]
         v_unpad = v.permute(0, 2, 1, 3)  # [B, S, H, D]
-        
+
         # Preparar máscara: convertir a bool si es necesario
+        # Mantener la lógica original
         if attention_mask.dtype != torch.bool:
             attention_mask = attention_mask.bool()
-        
+
         # Hacer unpadding de los tensores
+        # Se mantienen las salidas originales, incluyendo el quinto elemento descartado
         q_unpadded, indices_q, cu_seqlens_q, max_seqlen_q, _ = self.unpad_input(q_unpad, attention_mask)
         k_unpadded, indices_k, cu_seqlens_k, max_seqlen_k, _ = self.unpad_input(k_unpad, attention_mask)
         v_unpadded, _, _, _, _ = self.unpad_input(v_unpad, attention_mask)
-        
+
         # Reacomodar para flash_attn_varlen_func: [Total, H, D]
         q_unpadded = q_unpadded.reshape(-1, H, D)
         k_unpadded = k_unpadded.reshape(-1, H, D)
         v_unpadded = v_unpadded.reshape(-1, H, D)
-        
+
         # Normalizar vectores Q y K para mejorar estabilidad numérica
-        q_norm = F.normalize(q_unpadded, p=2, dim=-1).to(torch.bfloat16)
-        k_norm = F.normalize(k_unpadded, p=2, dim=-1).to(torch.bfloat16)
-        
-        # Ajustar q con factor de escala
-        s = self.ssmax_scale.view(1, H, 1)
-        q_adjusted = q_norm * (self.seq_scale * s)
-        
-        # Factor de escala para softmax
-        softmax_scale = 1.0 / math.sqrt(D)
-        
+        q_norm = F.normalize(q_unpadded, p=2, dim=-1)
+        k_norm = F.normalize(k_unpadded, p=2, dim=-1)
+
+        # Eliminado el ajuste de q con factor de escala ssmax_scale y seq_scale
+        # s = self.ssmax_scale.view(1, H, 1)
+        # q_adjusted = q_norm * (self.seq_scale * s)
+
+        # Factor de escala estándar para softmax
+
         try:
-            # Usar flash attention sin padding
+            # Usar flash attention sin padding, pasando q_norm
             output_unpadded = self.flash_attn_varlen_func(
-                q_adjusted, k_norm, v_unpadded,
+                q_norm, k_norm, v_unpadded, # Usar q_norm directamente
                 cu_seqlens_q, cu_seqlens_k,
                 max_seqlen_q, max_seqlen_k,
                 dropout_p=self.attention_dropout.p,  # Aplicamos dropout aquí
-                softmax_scale=softmax_scale,
+                softmax_scale=None,         # Escala estándar
                 causal=is_causal
             )
-            
+
             # Volver a aplicar padding
             output_padded = self.pad_input(output_unpadded, indices_q, B, S)
-            
+
             # Reorganizar a [B, H, S, D]
             output = output_padded.reshape(B, S, H, D).permute(0, 2, 1, 3)
-            
+
             return output
-            
+
         except Exception as e:
+            # Mantener el manejo de errores original
             raise RuntimeError(f"Error en flash_attn_varlen_func: {e}")
 
     def scaled_dot_product_attention_flash(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
-                                           mask: Optional[torch.Tensor] = None,
+                                           mask: Optional[torch.Tensor] = None, # Mantener mask opcional
                                            is_causal: bool = False) -> torch.Tensor:
         # Normalizar vectores Q y K para mejorar estabilidad numérica
-        q_norm = F.normalize(q, p=2, dim=-1).to(torch.bfloat16)
-        k_norm = F.normalize(k, p=2, dim=-1).to(torch.bfloat16)
-        
-        # Ajustar q con factor de escala
-        s = self.ssmax_scale.view(-1, 1, 1)
-        q_adjusted = q_norm * (self.seq_scale * s)
-        
+        q_norm = F.normalize(q, p=2, dim=-1)
+        k_norm = F.normalize(k, p=2, dim=-1)
+
+        # Eliminado el ajuste de q con factor de escala ssmax_scale y seq_scale
+        # s = self.ssmax_scale.view(-1, 1, 1)
+        # q_adjusted = q_norm * (self.seq_scale * s)
+
         # Preparar tensores para Flash Attention (requiere shape [B, S, H, D])
-        q_trans = q_adjusted.permute(0, 2, 1, 3)
+        q_trans = q_norm.permute(0, 2, 1, 3) # Usar q_norm directamente
         k_trans = k_norm.permute(0, 2, 1, 3)
         v_trans = v.permute(0, 2, 1, 3)
-        
-        # Verificar dimensiones
+
+        # Mantener la verificación de dimensiones original
         if q_trans.size(-1) != k_trans.size(-1):
             raise ValueError(f"Las dimensiones de head no coinciden: q={q_trans.size(-1)}, k={k_trans.size(-1)}")
-        
-        # Factor de escala para softmax
-        softmax_scale = 1.0 / math.sqrt(q_trans.size(-1))
-        
+
+        # Factor de escala estándar para softmax
         try:
-            # Aplicar Flash Attention
+            # Aplicar Flash Attention, pasando q_trans
             output = self.flash_attn_func(
                 q_trans, k_trans, v_trans,
                 dropout_p=self.attention_dropout.p,  # Aplicamos dropout aquí
-                softmax_scale=softmax_scale,
+                softmax_scale=None,         # Escala estándar
                 causal=is_causal
+                # mask no se usa aquí
             )
-            
+
+            # Mantener la verificación de salida None original
             if output is None:
                 raise ValueError("flash_attn_func devolvió None. Verifica las dimensiones y tipos de los tensores de entrada.")
-            
+
             # Volver a la forma original
             output = output.permute(0, 2, 1, 3)
             return output
-            
+
         except Exception as e:
+            # Mantener el manejo de errores original
             raise RuntimeError(f"Error en flash_attn_func: {e}")
 
     def forward(self, X: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, causal: bool = True) -> torch.Tensor:
         B, T, _ = X.shape
-        
+        norm_func = self.norm # Referencia a la capa de normalización
+
         # Implementación de HybridNorm*
         if self.use_pre_norm:
             # Primer bloque: Pre-Norm en atención
-            X_norm = self.norm(X)
+            # Aplicar norm y luego convertir explícitamente a bfloat16
+            X_norm = norm_func(X).to(torch.bfloat16)
             # Proyecciones para Q, K, V con FANformer
             Q = self.Wq(X_norm)  # [B, T, num_heads, head_dim]
             K = self.Wk(X_norm)  # [B, T, num_heads, head_dim]
             V = self.Wv(X_norm)  # [B, T, num_heads, head_dim]
         else:
             # Otros bloques: QKV-Norm
-            Q = self.Wq(self.norm(X))  # [B, T, num_heads, head_dim]
-            K = self.Wk(self.norm(X))  # [B, T, num_heads, head_dim]
-            V = self.Wv(self.norm(X))  # [B, T, num_heads, head_dim]
-        
+            # Aplicar norm y convertir explícitamente a bfloat16 antes de cada proyección
+            Q = self.Wq(norm_func(X).to(torch.bfloat16))
+            K = self.Wk(norm_func(X).to(torch.bfloat16))
+            V = self.Wv(norm_func(X).to(torch.bfloat16))
+
         # Permutar a formato [B, num_heads, T, head_dim]
         Q = Q.permute(0, 2, 1, 3)
         K = K.permute(0, 2, 1, 3)
         V = V.permute(0, 2, 1, 3)
-        
+
         # Aplicar RoPE si está activado
         if self.use_rope:
             Q = apply_rope_vectorized(Q)
             K = apply_rope_vectorized(K)
-        
-        # Convertir a bfloat16 para flash attention
+
+        # Convertir a bfloat16 para flash attention (mantener esta conversión explícita)
         Q = Q.to(torch.bfloat16)
         K = K.to(torch.bfloat16)
         V = V.to(torch.bfloat16)
-        
+
         # Procesar la secuencia utilizando unpadding si hay máscara de atención
+        # Mantener la lógica original para decidir la ruta
         if attention_mask is not None:
             attn_output = self.scaled_dot_product_attention_flash_unpadded(
-                Q, K, V, 
-                attention_mask=attention_mask, 
+                Q, K, V,
+                attention_mask=attention_mask,
                 is_causal=causal
             )
         else:
             # Si no hay máscara, usar la versión regular
             attn_output = self.scaled_dot_product_attention_flash(
-                Q, K, V, 
-                mask=None, 
+                Q, K, V,
+                mask=None,
                 is_causal=causal
             )
-        
-        # Eliminada la aplicación redundante de dropout:
+
+        # Eliminada la aplicación redundante de dropout (ya estaba eliminada)
         # attn_output = self.attention_dropout(attn_output)
-        
+
         # Reorganizar la salida y aplicar proyección final
         out = attn_output.permute(0, 2, 1, 3).contiguous()
         out = out.reshape(B, T, self.embed_dim)
         out = self.output_dropout(self.out_proj(out))
-        
-        return out
 
+        return out
 ############################################
 # NUEVO MÓDULO: SWIGLU CON COLA (MLP)
 ############################################