Update modeling_alinlight.py

modeling_alinlight.py

@@ -72,6 +72,23 @@ class AlinlightRMSNorm(nn.Module):
         return self.weight * x.to(input_dtype)
 
 
+class GatedNorm(nn.Module):
+    """
+    Gated Normalization wrapper.
+    Allows the model to learn to skip normalization via a learnable gate.
+    """
+    def __init__(self, original_norm, initial_gate_value=-1.0):
+        super().__init__()
+        self.norm = original_norm
+        # Initialize gate to -1.0 (sigmoid(-1) ≈ 0.27) to start conservatively
+        self.gate = nn.Parameter(torch.tensor(initial_gate_value))
+
+    def forward(self, x, *args, **kwargs):
+        normed = self.norm(x, *args, **kwargs)
+        g = torch.sigmoid(self.gate)
+        return (1.0 - g) * x + g * normed
+
+
 class AlinlightRotaryEmbedding(nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
         super().__init__()
@@ -134,7 +151,11 @@ class AlinlightMLP(nn.Module):
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
         self.act_fn = nn.SiLU()
-        self.pre_down_norm = AlinlightRMSNorm(self.intermediate_size, eps=config.rms_norm_eps)
+        
+        # Use GatedNorm for the inner normalization
+        self.pre_down_norm = GatedNorm(
+            AlinlightRMSNorm(self.intermediate_size, eps=config.rms_norm_eps)
+        )
         
         # Tag for specialized initialization
         self.down_proj._is_residual_projection = True
@@ -171,8 +192,9 @@ class AlinlightAttention(nn.Module):
 
         self.use_qk_norm = getattr(config, "use_qk_norm", True)
         if self.use_qk_norm:
-            self.q_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
-            self.k_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            # Use GatedNorm for QK Normalization
+            self.q_norm = GatedNorm(AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps))
+            self.k_norm = GatedNorm(AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps))
         
         self.attn_logit_softcapping = getattr(config, 'attn_logit_softcapping', None)