Update modular_Sewy2.py

modular_Sewy2.py

@@ -2,7 +2,7 @@ from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
 """ PyTorch Sewy model."""
-"""Used deepseek-V3 as a starting point."""
+"""Used deepseekv3 as starting point"""
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
@@ -214,6 +214,8 @@ class SewyV2Config(PretrainedConfig):
         unit_norm_eps = 1e-6,
         resformer_lambda = 2.0,
         neutreno_lambda=0.4,
+        final_logit_softcapping=30.0,
+        attn_logit_softcapping=50.0,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -260,6 +262,8 @@ class SewyV2Config(PretrainedConfig):
         self.unit_norm_eps = unit_norm_eps
         self.resformer_lambda = resformer_lambda
         self.neutreno_lambda = neutreno_lambda
+        self.final_logit_softcapping = final_logit_softcapping
+        self.attn_logit_softcapping = attn_logit_softcapping
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,
@@ -907,6 +911,9 @@ class SewyV2Attention(nn.Module):
 
         self.neutreno_lambda = nn.Parameter(torch.tensor(float(config.neutreno_lambda)))
 
+        self.attn_logit_softcapping = self.config.attn_logit_softcapping
+
+
     def _get_unit_norm(self, x,eps=1e-6):
         """
         Normalize a tensor to unit norm
@@ -1080,6 +1087,10 @@ class SewyV2Attention(nn.Module):
                 )
             attn_weights = attn_weights + attention_mask
 
+        ## tanh softcapping 
+
+        attn_weights = self.attn_logit_softcapping * torch.tanh(attn_weights/self.attn_logit_softcapping)
+
         # upcast attention to fp32
         attn_weights = nn.functional.softmax(
             attn_weights, dim=-1, dtype=torch.float32
@@ -1279,6 +1290,7 @@ class SewyV2FlashAttention2(SewyV2Attention):
             q_len,
             dropout=dropout_rate,
             softmax_scale=self.softmax_scale,
+            softcap=self.attn_logit_softcapping,
         )
         if self.q_head_dim != self.v_head_dim:
             attn_output = attn_output[:, :, :, : self.v_head_dim]
@@ -1865,6 +1877,10 @@ class SewyV2ForCausalLM(SewyV2PreTrainedModel):
         self.s_z = nn.Parameter(torch.ones(self.vocab_size) * (1/config.hidden_size ** 0.5))
         self.s_z_init = 1
         self.s_z_scale = 1/config.hidden_size ** 0.5
+
+        # tanh softcapping
+
+        self.tanh_softcapping = config.final_logit_softcapping
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -2017,6 +2033,11 @@ class SewyV2ForCausalLM(SewyV2PreTrainedModel):
         logits = self.lm_head(hidden_states)
         logits = logits.float()
 
+        ## tanh softcapping
+
+        logits = self.tanh_softcapping * torch.tanh(logits/self.tanh_softcapping)
+
+        
         ## nGPT
 
         s_z = self.s_z * (self.s_z_init/self.s_z_scale)