Initial FanFormer checkpoint with architecture and README

model.safetensors

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

model_architecture.py

@@ -452,8 +452,8 @@ class FANformerMultiheadAttention(nn.Module):
         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)
-        k_norm = F.normalize(k_unpadded, p=2, dim=-1)
+        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)
 
         # Eliminado el ajuste de q con factor de escala ssmax_scale y seq_scale
         # s = self.ssmax_scale.view(1, H, 1)
@@ -488,8 +488,8 @@ class FANformerMultiheadAttention(nn.Module):
                                            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)
-        k_norm = F.normalize(k, p=2, dim=-1)
+        q_norm = F.normalize(q, p=2, dim=-1).to(torch.bfloat16)
+        k_norm = F.normalize(k, p=2, dim=-1).to(torch.bfloat16)
 
         # Eliminado el ajuste de q con factor de escala ssmax_scale y seq_scale
         # s = self.ssmax_scale.view(-1, 1, 1)