FIX ATTENTION

modeling_alinlight.py

@@ -18,16 +18,17 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from typing import Optional, Tuple, List, Union
-from transformers import PreTrainedModel
+from transformers import PreTrainedModel, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
-from transformers import GenerationMixin
-from configuration_alinlight import AlinlightConfig # Импортируем конфиг из соседнего файла
+
+from configuration_alinlight import AlinlightConfig
 
 class AlinlightRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.eps = eps
+
     def forward(self, x):
         input_dtype = x.dtype
         x = x.to(torch.float32)
@@ -42,8 +43,12 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.base = base
         self.max_position_embeddings = max_position_embeddings
         self.scaling_factor = scaling_factor
+        
+        # We calculate frequencies immediately upon initialization
         inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32).to(device) / self.dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
+        
+        # Initialize the cache
         self._set_cos_sin_cache(seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype())
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
@@ -55,9 +60,14 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
     def forward(self, x, seq_len=None):
+        # If the length is greater than the cache, we recalculate (a rare case, but needed for reliability)
         if seq_len > self.cos_cached.shape[0]:
              self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-        return self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype)
+        
+        return (
+            self.cos_cached[:seq_len].to(dtype=x.dtype, device=x.device),
+            self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device)
+        )
 
 def rotate_half(x):
     x1 = x[..., : x.shape[-1] // 2]
@@ -95,18 +105,28 @@ class AlinlightAttention(nn.Module):
         self.num_key_value_heads = config.num_key_value_heads
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
         self.sliding_window = config.sliding_window
-        self.attention_dropout = config.attention_dropout
-
+        
         self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
         self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
         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)
 
-    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, output_attentions=False, use_cache=False, cos_sin=None):
+    def forward(
+        self, 
+        hidden_states, 
+        attention_mask=None, 
+        position_ids=None, 
+        past_key_value=None, 
+        output_attentions=False, 
+        use_cache=False, 
+        cos_sin=None
+    ):
         bsz, q_len, _ = hidden_states.size()
+        
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
+        
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
@@ -119,19 +139,33 @@ 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)
 
-        # Truncation logic for sliding window
+        # === TPU OPTIMIZATION: Physical cropping instead of mask ===
         if self.sliding_window is not None and key_states.shape[2] > self.sliding_window:
             key_states = key_states[:, :, -self.sliding_window:, :]
             value_states = value_states[:, :, -self.sliding_window:, :]
 
         past_key_value = (key_states, value_states) if use_cache else None
         
+        # GQA / MQA processing
         if self.num_key_value_groups > 1:
-            key_states = key_states[:, :, None, :, :].expand(bsz, self.num_key_value_heads, self.num_key_value_groups, key_states.shape[-2], self.head_dim).reshape(bsz, self.num_heads, key_states.shape[-2], self.head_dim)
-            value_states = value_states[:, :, None, :, :].expand(bsz, self.num_key_value_heads, self.num_key_value_groups, value_states.shape[-2], self.head_dim).reshape(bsz, self.num_heads, value_states.shape[-2], self.head_dim)
+            key_states = key_states[:, :, None, :, :].expand(
+                bsz, self.num_key_value_heads, self.num_key_value_groups, key_states.shape[-2], self.head_dim
+            ).reshape(bsz, self.num_heads, key_states.shape[-2], self.head_dim)
+            
+            value_states = value_states[:, :, None, :, :].expand(
+                bsz, self.num_key_value_heads, self.num_key_value_groups, value_states.shape[-2], self.head_dim
+            ).reshape(bsz, self.num_heads, value_states.shape[-2], self.head_dim)
 
-        # Use Scaled Dot Product Attention
-        attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=None, dropout_p=0.0, is_causal=True)
+        # SDPA (Flash Attention backend compatible)
+        attn_output = F.scaled_dot_product_attention(
+            query_states, 
+            key_states, 
+            value_states, 
+            attn_mask=None, 
+            dropout_p=0.0, 
+            is_causal=True 
+        )
+        
         attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
         return self.o_proj(attn_output), None, past_key_value
 
@@ -146,12 +180,16 @@ class AlinlightDecoderLayer(nn.Module):
     def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, output_attentions=False, use_cache=False, cos_sin=None):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states, _, present_key_value = self.self_attn(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cos_sin)
+        hidden_states, _, present_key_value = self.self_attn(
+            hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cos_sin
+        )
         hidden_states = residual + hidden_states
+        
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
+        
         return hidden_states, None, present_key_value
 
 class AlinlightModel(PreTrainedModel):
@@ -166,7 +204,16 @@ class AlinlightModel(PreTrainedModel):
         if config.rope_scaling and config.rope_scaling.get("type") == "linear":
              scaling_factor = config.rope_scaling.get("factor", 1.0)
              
-        self.rotary_emb = AlinlightRotaryEmbedding(config.hidden_size // config.num_attention_heads, max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, scaling_factor=scaling_factor)
+        self.rotary_emb = AlinlightRotaryEmbedding(
+            config.hidden_size // config.num_attention_heads, 
+            max_position_embeddings=config.max_position_embeddings, 
+            base=config.rope_theta, 
+            scaling_factor=scaling_factor
+        )
+        self.post_init()
+
+    def get_input_embeddings(self): return self.embed_tokens
+    def set_input_embeddings(self, value): self.embed_tokens = value
 
     def forward(self, input_ids=None, past_key_values=None, use_cache=None, **kwargs):
         if input_ids is not None:
@@ -174,6 +221,7 @@ class AlinlightModel(PreTrainedModel):
         else:
              inputs_embeds = kwargs.get("inputs_embeds")
 
+        # Length calculation for RoPE
         seq_len = inputs_embeds.shape[1]
         if past_key_values is not None:
              seq_len += past_key_values[0][0].shape[2]
@@ -190,7 +238,13 @@ class AlinlightModel(PreTrainedModel):
 
         for idx, layer in enumerate(self.layers):
             past_key_value = past_key_values[idx] if past_key_values is not None else None
-            layer_outputs = layer(hidden_states, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache, cos_sin=(cos, sin))
+            layer_outputs = layer(
+                hidden_states, 
+                position_ids=position_ids, 
+                past_key_value=past_key_value, 
+                use_cache=use_cache, 
+                cos_sin=(cos, sin)
+            )
             hidden_states = layer_outputs[0]
             if use_cache:
                 next_decoder_cache += (layer_outputs[2],)
@@ -210,9 +264,16 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
         super().__init__(config)
         self.model = AlinlightModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.lm_head.weight = self.model.embed_tokens.weight
+        self.lm_head.weight = self.model.embed_tokens.weight # Связываем веса (Weight Tying)
+        self.post_init()
+
+    def get_input_embeddings(self): return self.model.embed_tokens
+    def set_input_embeddings(self, value): self.model.embed_tokens = value
+    def get_output_embeddings(self): return self.lm_head
+    def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        # Optimization for generation: if there is a cache, we serve only the last token
         if past_key_values:
             input_ids = input_ids[:, -1:]