Update modeling_alinlight.py

modeling_alinlight.py

@@ -16,6 +16,9 @@
 # -*- coding: utf-8 -*-
 # Copyright 2026 EngineerGL Research.
 
+# -*- coding: utf-8 -*-
+# Copyright 2026 EngineerGL Research.
+
 import math
 import torch
 import torch.nn as nn
@@ -46,11 +49,15 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.max_position_embeddings = max_position_embeddings
         self.scaling_factor = scaling_factor
         
+        # Вычисляем частоты сразу
         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)
+        
+        # Инициализируем кэш
         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):
+        # RoPE кэш считаем во float32 для точности, потом кастим
         t = torch.arange(seq_len, device=device, dtype=torch.int64).type_as(self.inv_freq)
         t = t / self.scaling_factor
         freqs = torch.outer(t, self.inv_freq)
@@ -59,8 +66,10 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
     def forward(self, x, seq_len=None):
+        # Если вдруг длина последовательности больше кэша (редкий случай)
         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, device=x.device),
             self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device)
@@ -136,13 +145,15 @@ 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 (TPU Optimization)
+        # === TPU/Efficiency Optimization: Physical Truncation ===
+        # Вместо сложной маски просто обрезаем старые токены, вышедшие за пределы окна.
         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 Expansion
         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
@@ -152,9 +163,10 @@ class AlinlightAttention(nn.Module):
                 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)
 
-        # ✅ FIX: Динамический флаг is_causal
-        # Если q_len > 1 (обучение или prefill) -> True (маскируем будущее)
-        # Если q_len == 1 (генерация) -> False (видим всё прошлое, что есть в key_states)
+        # === Fix for "Blind" Generation ===
+        # is_causal=True нужно ТОЛЬКО при обучении (или prefill), когда q_len > 1.
+        # Когда мы генерируем по 1 токену (q_len == 1), мы смотрим на весь кэш (прошлое).
+        # Если поставить True при q_len=1, SDPA может замаскировать всё.
         is_causal = q_len > 1
 
         attn_output = F.scaled_dot_product_attention(
@@ -220,14 +232,17 @@ class AlinlightModel(PreTrainedModel):
         else:
              inputs_embeds = kwargs.get("inputs_embeds")
 
+        # Определяем длину последовательности для RoPE
         seq_len = inputs_embeds.shape[1]
         if past_key_values is not None:
              seq_len += past_key_values[0][0].shape[2]
              
+        # Получаем cos/sin
         cos, sin = self.rotary_emb(inputs_embeds, seq_len=seq_len)
         
         position_ids = kwargs.get("position_ids")
         if position_ids is None:
+             # Генерируем position_ids сами
              position_ids = torch.arange(seq_len - inputs_embeds.shape[1], seq_len, dtype=torch.long, device=inputs_embeds.device)
              position_ids = position_ids.unsqueeze(0).expand(inputs_embeds.shape[0], -1)
 
@@ -262,7 +277,14 @@ 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
+        
+        # === ВАЖНЕЙШЕЕ ИСПРАВЛЕНИЕ ===
+        # Связываем веса ТОЛЬКО если это разрешено в конфиге.
+        # Раньше это было принудительно, что ломало генерацию для моделей типа Mistral/Llama3,
+        # где tie_word_embeddings=False.
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+            
         self.post_init()
 
     def get_input_embeddings(self): return self.model.embed_tokens
@@ -271,15 +293,18 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
     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):
+        # Если есть кэш, нам нужен только последний токен из input_ids
         if past_key_values:
             input_ids = input_ids[:, -1:]
         
         position_ids = kwargs.get("position_ids", None)
         if position_ids is None:
             if past_key_values:
+                 # Если есть кэш, позиция = длина кэша
                  past_length = past_key_values[0][0].shape[2]
                  position_ids = torch.tensor([[past_length]], dtype=torch.long, device=input_ids.device)
             else:
+                 # Если кэша нет (первый токен), позиции = [0, 1, 2...]
                  position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device).unsqueeze(0)
 
         return {