Update modeling_alinlight.py

modeling_alinlight.py

@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# -*- coding: utf-8 -*-
+# Copyright 2026 EngineerGL Research.
+
 import math
 import torch
 import torch.nn as nn
@@ -41,7 +44,6 @@ class AlinlightPreTrainedModel(PreTrainedModel):
     def _init_weights(self, module):
         std = self.config.initializer_range
         if isinstance(module, nn.Linear):
-            # Scale down residual projections to improve training stability at depth
             if getattr(module, '_is_residual_projection', False):
                 module.weight.data.normal_(mean=0.0, std=std / math.sqrt(2 * self.config.num_hidden_layers))
             else:
@@ -72,23 +74,6 @@ 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__()
@@ -152,22 +137,14 @@ class AlinlightMLP(nn.Module):
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
         self.act_fn = nn.SiLU()
         
-        # 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
 
     def forward(self, x):
-        intermediate = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
-        intermediate = self.pre_down_norm(intermediate)
-        return self.down_proj(intermediate)
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
 
 
 # ==========================================
-# 3. ATTENTION
+# 3. ATTENTION (Стабильная QK-Norm без гейтов)
 # ==========================================
 
 class AlinlightAttention(nn.Module):
@@ -192,9 +169,8 @@ class AlinlightAttention(nn.Module):
 
         self.use_qk_norm = getattr(config, "use_qk_norm", True)
         if self.use_qk_norm:
-            # 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.q_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
         
         self.attn_logit_softcapping = getattr(config, 'attn_logit_softcapping', None)
 
@@ -233,8 +209,7 @@ class AlinlightAttention(nn.Module):
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
             value_states = torch.cat([past_key_value[1], value_states], dim=2)
 
-        # 3. Sliding Window (Slicing)
-        kv_seq_len = key_states.shape[2] # NOTE: This is the length BEFORE slicing
+        kv_seq_len = key_states.shape[2]
         
         if self.sliding_window is not None and kv_seq_len > self.sliding_window:
             slicing_tokens = kv_seq_len - self.sliding_window
@@ -246,12 +221,12 @@ class AlinlightAttention(nn.Module):
 
         past_key_value = (key_states, value_states) if use_cache else None
 
-        # 4. GQA Repeat
+        # 3. GQA Repeat
         if self.num_key_value_groups > 1:
             key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
             value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
 
-        # 5. Attention Mechanism
+        # 4. Attention
         attn_weights = None
         
         if output_attentions or self.attn_logit_softcapping is not None:
@@ -264,9 +239,7 @@ class AlinlightAttention(nn.Module):
                 attn_weights = attn_weights + attention_mask
             
             attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-            
             attn_weights_for_output = attn_weights if output_attentions else None
-
             attn_weights_dropped = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
             attn_output = torch.matmul(attn_weights_dropped, value_states)
         else:
@@ -339,7 +312,6 @@ class AlinlightModel(AlinlightPreTrainedModel):
         self.vocab_size = config.vocab_size
 
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        
         self.embed_scale = math.sqrt(config.hidden_size) if getattr(config, 'embed_scale', False) else 1.0
         
         embed_pdrop = getattr(config, 'embed_pdrop', 0.0)
@@ -414,7 +386,6 @@ class AlinlightModel(AlinlightPreTrainedModel):
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        # --- SAFETY CHECK FOR GRADIENT CHECKPOINTING ---
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 logger.warning_once(
@@ -459,7 +430,6 @@ class AlinlightModel(AlinlightPreTrainedModel):
             if self.gradient_checkpointing and self.training:
                 def create_custom_forward(module):
                     def custom_forward(*inputs):
-                        # Force use_cache=False inside checkpoint to be safe
                         return module(*inputs, output_attentions=output_attentions, use_cache=False, rotary_pos_emb=(cos, sin))
                     return custom_forward
                 
@@ -520,8 +490,6 @@ class AlinlightForCausalLM(AlinlightPreTrainedModel, GenerationMixin):
         if config.tie_word_embeddings:
             self.lm_head.weight = self.model.embed_tokens.weight
 
-        # Note: self.post_init() is called here, and inside AlinlightModel.
-        # This re-initialization is consistent with standard HF models (e.g. Llama).
         self.post_init()
 
     def get_input_embeddings(self): return self.model.embed_tokens