Update modeling_alinlight.py

modeling_alinlight.py

@@ -17,13 +17,43 @@ import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import warnings
 from typing import Optional, Tuple, List, Union
 from torch.utils.checkpoint import checkpoint
 
 from transformers import PreTrainedModel, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
+from transformers.utils import logging
 from configuration_alinlight import AlinlightConfig
 
+logger = logging.get_logger(__name__)
+
+# ==========================================
+# 0. BASE PRETRAINED MODEL
+# ==========================================
+
+class AlinlightPreTrainedModel(PreTrainedModel):
+    config_class = AlinlightConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["AlinlightDecoderLayer"]
+    _supports_gradient_checkpointing = True
+
+    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:
+                module.weight.data.normal_(mean=0.0, std=std)
+            
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
 # ==========================================
 # 1. BASE COMPONENTS
 # ==========================================
@@ -137,7 +167,6 @@ class AlinlightAttention(nn.Module):
         self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
         
-        # Tag for specialized initialization
         self.o_proj._is_residual_projection = True
 
         self.use_qk_norm = getattr(config, "use_qk_norm", True)
@@ -155,7 +184,7 @@ class AlinlightAttention(nn.Module):
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
-        cos_sin: Optional[Tuple[torch.Tensor]] = None
+        rotary_pos_emb: Optional[Tuple[torch.Tensor]] = None
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 
         bsz, q_len, _ = hidden_states.size()
@@ -172,25 +201,25 @@ class AlinlightAttention(nn.Module):
             query_states = self.q_norm(query_states)
             key_states = self.k_norm(key_states)
 
-        # 1. RoPE (Applied before caching)
-        if cos_sin is not None:
-            cos, sin = cos_sin
+        # 1. RoPE
+        if rotary_pos_emb is not None:
+            cos, sin = rotary_pos_emb
             query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
-        # 2. KV Cache
+        # 2. KV Cache Update
         if past_key_value is not None:
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
             value_states = torch.cat([past_key_value[1], value_states], dim=2)
 
-        kv_seq_len = key_states.shape[2]
-
         # 3. Sliding Window (Slicing)
+        kv_seq_len = key_states.shape[2] # NOTE: This is the length BEFORE slicing
+        
         if self.sliding_window is not None and kv_seq_len > self.sliding_window:
             slicing_tokens = kv_seq_len - self.sliding_window
             key_states = key_states[:, :, slicing_tokens:, :]
             value_states = value_states[:, :, slicing_tokens:, :]
             
-            if attention_mask is not None:
+            if attention_mask is not None and attention_mask.shape[-1] == kv_seq_len:
                 attention_mask = attention_mask[:, :, :, slicing_tokens:]
 
         past_key_value = (key_states, value_states) if use_cache else None
@@ -203,9 +232,6 @@ class AlinlightAttention(nn.Module):
         # 5. Attention Mechanism
         attn_weights = None
         
-        # We must use manual implementation if:
-        # a) Output weights are requested
-        # b) Soft-capping is enabled (SDPA doesn't support intermediate logit transforms)
         if output_attentions or self.attn_logit_softcapping is not None:
             attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
             
@@ -217,16 +243,11 @@ class AlinlightAttention(nn.Module):
             
             attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
             
-            if not output_attentions:
-                # If we only calculated weights for soft-capping but user didn't ask for them, drop reference
-                attn_weights_for_output = None
-            else:
-                attn_weights_for_output = attn_weights
+            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:
-            # Fast Path (SDPA)
             attn_output = F.scaled_dot_product_attention(
                 query_states,
                 key_states,
@@ -264,7 +285,7 @@ class AlinlightDecoderLayer(nn.Module):
         past_key_value=None, 
         output_attentions=False, 
         use_cache=False, 
-        cos_sin=None
+        rotary_pos_emb=None
     ):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
@@ -276,7 +297,7 @@ class AlinlightDecoderLayer(nn.Module):
             past_key_value=past_key_value,
             output_attentions=output_attentions,
             use_cache=use_cache,
-            cos_sin=cos_sin
+            rotary_pos_emb=rotary_pos_emb
         )
         hidden_states = residual + self.resid_dropout(hidden_states)
 
@@ -288,9 +309,7 @@ class AlinlightDecoderLayer(nn.Module):
         return hidden_states, attn_weights, present_key_value
 
 
-class AlinlightModel(PreTrainedModel):
-    config_class = AlinlightConfig
-    
+class AlinlightModel(AlinlightPreTrainedModel):
     def __init__(self, config: AlinlightConfig):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
@@ -299,7 +318,9 @@ class AlinlightModel(PreTrainedModel):
         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
-        self.embed_dropout = nn.Dropout(config.embed_pdrop) if config.embed_pdrop > 0 else nn.Identity()
+        
+        embed_pdrop = getattr(config, 'embed_pdrop', 0.0)
+        self.embed_dropout = nn.Dropout(embed_pdrop) if embed_pdrop > 0 else nn.Identity()
 
         self.layers = nn.ModuleList([AlinlightDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
         self.norm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -369,6 +390,14 @@ class AlinlightModel(PreTrainedModel):
         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(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
         
@@ -406,10 +435,17 @@ class AlinlightModel(PreTrainedModel):
             if self.gradient_checkpointing and self.training:
                 def create_custom_forward(module):
                     def custom_forward(*inputs):
-                        return module(*inputs, output_attentions=output_attentions, use_cache=False, cos_sin=(cos, sin))
+                        # 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
+                
                 layer_outputs = checkpoint(
-                    create_custom_forward(layer), hidden_states, attention_mask, position_ids, past_key_value, use_reentrant=True
+                    create_custom_forward(layer), 
+                    hidden_states, 
+                    attention_mask, 
+                    position_ids, 
+                    past_key_value, 
+                    use_reentrant=False
                 )
             else:
                 layer_outputs = layer(
@@ -419,7 +455,7 @@ class AlinlightModel(PreTrainedModel):
                     past_key_value=past_key_value,
                     output_attentions=output_attentions,
                     use_cache=use_cache,
-                    cos_sin=(cos, sin)
+                    rotary_pos_emb=(cos, sin)
                 )
 
             hidden_states = layer_outputs[0]
@@ -448,11 +484,7 @@ class AlinlightModel(PreTrainedModel):
 # 5. CAUSAL LM HEAD
 # ==========================================
 
-class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
-    config_class = AlinlightConfig
-    _keys_to_ignore_on_load_missing = ["model.rotary_emb.inv_freq"]
-    _supports_gradient_checkpointing = True
-
+class AlinlightForCausalLM(AlinlightPreTrainedModel, GenerationMixin):
     def __init__(self, config):
         super().__init__(config)
         self.model = AlinlightModel(config)
@@ -464,6 +496,8 @@ class AlinlightForCausalLM(PreTrainedModel, 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
@@ -471,25 +505,11 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
     def get_output_embeddings(self): return self.lm_head
     def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings
 
-    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:
-                module.weight.data.normal_(mean=0.0, std=std)
-            
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-
     def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
         self.model.gradient_checkpointing = True
-        self.config.use_cache = False
+
+    def gradient_checkpointing_disable(self):
+        self.model.gradient_checkpointing = False
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
         if past_key_values is not None:
@@ -498,7 +518,11 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
         position_ids = kwargs.get("position_ids", None)
         if position_ids is None:
             if past_key_values:
-                position_ids = (attention_mask.long().sum(dim=-1) - 1).unsqueeze(-1)
+                if attention_mask is not None:
+                    position_ids = (attention_mask.long().sum(dim=-1) - 1).unsqueeze(-1)
+                else:
+                    past_length = past_key_values[0][0].shape[2]
+                    position_ids = torch.tensor([[past_length]], device=input_ids.device)
             else:
                 position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device).unsqueeze(0)
 
@@ -540,7 +564,6 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
 
-        # Final Logit Soft-Capping
         if self.final_logit_softcapping is not None:
             logits = self.final_logit_softcapping * torch.tanh(logits / self.final_logit_softcapping)
 
@@ -552,9 +575,7 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
             loss_fct = nn.CrossEntropyLoss()
             ce_loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
             
-            # Z-Loss Regularization
             if self.z_loss_weight > 0 and self.training:
-                # log(sum(exp(x)))^2
                 z_loss = torch.logsumexp(shift_logits, dim=-1).pow(2).mean()
                 loss = ce_loss + self.z_loss_weight * z_loss
             else: