modeling_unified.py

# ====================================================================
# modeling_unified.py  
# ====================================================================

"""
Unified Language Model with GPAS + LNS Integration + xIELU Activation + CoLA (Linear Only) + LaX + Weight Tying + Canon Layers (A+C Only)
MIGRATED TO HUGGINGFACE TRANSFORMERS - FINAL VERSION WITH ALL FIXES
UPDATED: Standard Transformer with advanced variance control, parameter efficiency, and Canon horizontal information flow
Combines advanced Transformer architecture with CORRECTED variance control mechanisms,
advanced variance control via GPAS and LNS, xIELU activation function, LaX integration, and Canon Layers (A+C only)
Based on LLaMA 3 architecture with 30M parameters

MIGRATION TO HUGGINGFACE - FINAL FIXED VERSION:
==============================================

1. **HUGGINGFACE INTEGRATION**: Migrado de PyTorch Lightning a Transformers v4.53.3
2. **UPDATED API**: processing_class en lugar de tokenizer (deprecated)
3. **UPDATED COMPUTE_LOSS**: Método actualizado con num_items_in_batch parameter
4. **FIXED LOGGING**: Corregido self.log() syntax según documentación oficial HF
5. **RESTORED PAD HANDLING**: pad_token_id → -100 conversion for CrossEntropyLoss (from original code)
6. **NATIVE TORCH COMPILE**: Moved to TrainingArguments (torch_compile=True)
7. **FIXED WEIGHT TYING**: Corrected _tied_weights_keys as class attribute (HF standard)
8. **VALIDATION DIAGNOSTIC**: Added simple method to diagnose validation loss issues
9. **CUSTOM CONFIGURATION**: PretrainedConfig personalizada con todos los parámetros
10. **PRETRAINED MODEL**: Hereda de PreTrainedModel para compatibilidad completa
11. **MAINTAINED OPTIMIZERS**: Muon + AdamW híbrido preservado
12. **MAINTAINED PRECISION**: bf16-true preservado
13. **MAINTAINED TRAINING**: Custom Trainer con todas las métricas y logging
14. **MAINTAINED ARCHITECTURE**: Toda la arquitectura personalizada preservada
15. **AUTO TOKENIZER**: Integración completa con AutoTokenizer dinámico
16. **AUTOCLASS SUPPORT**: Registro completo para AutoConfig y AutoModel
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from transformers import (
    AutoTokenizer, 
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    PreTrainedModel,
)
import math
import os
from typing import Optional, Tuple, Dict, Any, cast, List
from flash_attn import flash_attn_func
import numpy as np

# ✅ ABSOLUTE IMPORT - No relative imports for Hub compatibility
from configuration_unified import UnifiedModelConfig

# Fix tokenizer parallelism warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
torch.set_float32_matmul_precision('high')

def init_cola_components(A: nn.Linear, B: nn.Linear):
    nn.init.kaiming_normal_(A.weight, mode='fan_in', nonlinearity='relu')
    nn.init.xavier_normal_(B.weight, gain=0.8)
    if B.bias is not None:
        nn.init.zeros_(B.bias)

def init_embedding(embedding: nn.Embedding):
    nn.init.normal_(embedding.weight, mean=0.0, std=0.02)

class CanonLayer(nn.Module):
    def __init__(self, hidden_dim: int, kernel_size: int = 4):
        """
        Canon layer using a 1D causal convolution with residual connection.
        """
        super().__init__()
        self.hidden_dim = hidden_dim
        self.kernel_size = kernel_size
        
        # Use causal convolution with explicit initialization
        self.causal_conv1d = nn.Conv1d(
            in_channels=hidden_dim,
            out_channels=hidden_dim,
            kernel_size=kernel_size,
            groups=hidden_dim,  # Depthwise convolution
            padding=0,  # No automatic padding
            bias=True
        )
        
        # Initialize weights more conservatively (as per paper)
        nn.init.zeros_(self.causal_conv1d.weight)
        nn.init.zeros_(self.causal_conv1d.bias)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        """
        Applies the Canon layer transformation with causal masking.
        """
        # Conv1d expects input shape (batch_size, channels, sequence_length)
        h_permuted = h.permute(0, 2, 1)  # (batch, hidden_dim, seq_len)
        
        # Add padding of (kernel_size - 1) only to the left side
        padding = self.kernel_size - 1
        h_padded = F.pad(h_permuted, (padding, 0))
        
        # Apply causal convolution
        conv_out = self.causal_conv1d(h_padded)
        
        # Permute back to the original shape
        conv_out_permuted = conv_out.permute(0, 2, 1)
        
        # Add the residual connection
        output = h + conv_out_permuted
        
        return output

class CoLA_Linear(nn.Module):
    def __init__(self, in_features: int, out_features: int, rank: Optional[int] = None, activation=F.gelu, bias: bool = True):
        super().__init__()
        if rank is None:
            rank = in_features // 4
        self.rank = rank
        self.activation = activation
        
        self.A = nn.Linear(in_features, rank, bias=False)
        self.B = nn.Linear(rank, out_features, bias=bias)
        
        # LaX Gate components (Linear only)
        self.lax_gate = nn.Parameter(torch.zeros(1))
        
        # Storage for previous layer's latent representation
        self.prev_latent = None
        
        init_cola_components(self.A, self.B)
    
    def apply_lax_gate(self, prev_latent: torch.Tensor) -> torch.Tensor:
        """Apply linear gate to previous latent representation."""
        return F.sigmoid(self.lax_gate) * prev_latent
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Standard CoLA forward
        latent = self.A(x)
        latent_activated = self.activation(latent)
        
        # Apply LaX if previous latent exists
        if self.prev_latent is not None and self.prev_latent.shape == latent_activated.shape:
            gated_prev = self.apply_lax_gate(self.prev_latent)
            latent_activated = latent_activated + gated_prev
        
        output = self.B(latent_activated)
        
        # Store current latent for next layer (detached to avoid gradient issues)
        self.prev_latent = latent_activated.detach()
        
        return output
    
    def reset_lax_state(self):
        self.prev_latent = None

class LayerNormScaling(nn.Module):
    def __init__(self, layer_depth: int):
        super().__init__()
        
        if layer_depth < 1:
            raise ValueError(f"layer_depth debe ser ≥ 1, got {layer_depth}")
        
        self.layer_depth = layer_depth
        self.scaling_factor = 1.0 / math.sqrt(float(layer_depth))
    
    def forward(self, normalized_input: torch.Tensor) -> torch.Tensor:
        return normalized_input * self.scaling_factor

class GPAS(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        
        self.d_model = d_model
        self.alpha = nn.Parameter(torch.zeros(1))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_detached = x.detach()
        scaled_component = F.silu(self.alpha) * x_detached
        x_scaled = x - scaled_component
        
        return x_scaled

class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_position_embeddings: int = 2048, base: float = 10000):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, x, seq_len=None):
        if seq_len is None:
            seq_len = x.shape[-2]
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos().to(x.dtype), emb.sin().to(x.dtype)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None):
    def rotate_half(x):
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2 :]
        return torch.cat((-x2, x1), dim=-1)
    
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

class XIELU(nn.Module):
    def __init__(self, alpha_p_init: float = 0.8, alpha_n_init: float = 0.8, beta: float = 0.5):
        super().__init__()
        
        self.beta = beta
        
        self.alpha_p = nn.Parameter(torch.log(torch.exp(torch.tensor(alpha_p_init)) - 1))
        self.alpha_n = nn.Parameter(torch.log(torch.exp(torch.tensor(alpha_n_init - self.beta)) - 1))
        
        self.register_buffer('eps', torch.tensor(-1e-6))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        alpha_p = F.softplus(self.alpha_p)
        alpha_n = self.beta + F.softplus(self.alpha_n)
        
        return torch.where(
            x > 0,
            alpha_p * x * x + self.beta * x,
            alpha_n * torch.expm1(torch.clamp(x, min=self.eps)) - alpha_n * x + self.beta * x
        )

class StandardMLP(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int, dropout: float = 0.0, config=None):
        super().__init__()
        
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.config = config
        
        self.up_proj = CoLA_Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = CoLA_Linear(intermediate_size, hidden_size, bias=False)
        
        if config is not None:
            self.activation = XIELU(
                alpha_p_init=config.xielu_alpha_p_init,
                alpha_n_init=config.xielu_alpha_n_init,
                beta=config.xielu_beta
            )
        else:
            self.activation = XIELU(alpha_p_init=0.8, alpha_n_init=0.8, beta=0.5)
        
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        
        # Canon-D is permanently disabled

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        intermediate = self.up_proj(x)
        
        # No Canon-D applied (eliminated)
        
        activated = self.activation(intermediate)
        activated = self.dropout(activated)
        output = self.down_proj(activated)
        
        return output
    
    def reset_lax_state(self):
        self.up_proj.reset_lax_state()
        self.down_proj.reset_lax_state()

class GroupedQueryAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        
        # FANFormer components
        self.fanformer_p = getattr(config, 'fanformer_p', 0.15)
        
        self.d_periodic = int(self.hidden_size * self.fanformer_p)
        self.d_standard = self.hidden_size - 2 * self.d_periodic
        
        assert self.d_standard > 0, \
            f"fanformer_p={self.fanformer_p} is too high. d_standard={self.d_standard} must be > 0"
        
        self.fan_w_p = CoLA_Linear(self.hidden_size, self.d_periodic, bias=False)
        self.fan_w_p_bar = CoLA_Linear(self.hidden_size, self.d_standard, bias=False)
        
        self.q_proj = CoLA_Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = CoLA_Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = CoLA_Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = CoLA_Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
        
        self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.v_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        
        self.rotary_emb = RotaryEmbedding(
            self.head_dim, 
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta
        )
        
        # Canon-B is permanently disabled (no more canon_b)
    
    def _fan_layer_prime(self, x: torch.Tensor) -> torch.Tensor:
        periodic_proj = self.fan_w_p(x)
        standard_proj = self.fan_w_p_bar(x)
        
        cos_component = torch.cos(periodic_proj)
        sin_component = torch.sin(periodic_proj)
        
        x_f = torch.cat([cos_component, sin_component, standard_proj], dim=-1)
        
        return x_f

    def _compute_flash_attention(
        self, 
        query_states: torch.Tensor, 
        key_states: torch.Tensor, 
        value_states: torch.Tensor,
        seq_len: int,
        position_ids: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        batch_size = query_states.shape[0]
        
        q_rope = query_states.transpose(1, 2)
        k_rope = key_states.transpose(1, 2)
        
        cos, sin = self.rotary_emb(value_states, seq_len=seq_len)
        q_rope, k_rope = apply_rotary_pos_emb(q_rope, k_rope, cos, sin, position_ids)
        
        query_states = q_rope.transpose(1, 2)
        key_states = k_rope.transpose(1, 2)

        from flash_attn import flash_attn_func
        
        attn_output = flash_attn_func(
            query_states,
            key_states,
            value_states,
            dropout_p=self.config.attention_dropout if self.training else 0.0,
            causal=True,
        )
        
        return attn_output

    def forward(self, hidden_states, position_ids=None, attention_mask=None):
        batch_size, seq_len, _ = hidden_states.shape
        
        enhanced_input = self._fan_layer_prime(hidden_states)
        
        query_states = self.q_proj(enhanced_input)
        key_states = self.k_proj(enhanced_input)
        value_states = self.v_proj(enhanced_input)
        
        # No Canon-B applied (eliminated)

        query_states = query_states.view(batch_size, seq_len, self.num_heads, self.head_dim)
        key_states = key_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
        value_states = value_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)

        q_flat = query_states.reshape(-1, self.head_dim)
        k_flat = key_states.reshape(-1, self.head_dim)
        v_flat = value_states.reshape(-1, self.head_dim)
        
        q_normalized = self.q_norm(q_flat)
        k_normalized = self.k_norm(k_flat)
        v_normalized = self.v_norm(v_flat)
        
        query_states = q_normalized.view(batch_size, seq_len, self.num_heads, self.head_dim)
        key_states = k_normalized.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
        value_states = v_normalized.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim)

        attn_output = self._compute_flash_attention(
            query_states=query_states,
            key_states=key_states,
            value_states=value_states,
            seq_len=seq_len,
            position_ids=position_ids
        )

        attn_output = attn_output.reshape(batch_size, seq_len, self.hidden_size)
        return self.o_proj(attn_output)
    
    def reset_lax_state(self):
        self.fan_w_p.reset_lax_state()
        self.fan_w_p_bar.reset_lax_state()
        self.q_proj.reset_lax_state()
        self.k_proj.reset_lax_state()
        self.v_proj.reset_lax_state()
        self.o_proj.reset_lax_state()

class DecoderLayer(nn.Module):
    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        
        if layer_idx < 0:
            raise ValueError(f"layer_idx debe ser >= 0, got {layer_idx}")
        
        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = GroupedQueryAttention(config)
        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        self.mlp = StandardMLP(
            config.hidden_size, 
            config.intermediate_size, 
            config.mlp_dropout,
            config
        )
        
        self.dropout_output = nn.Dropout(0.01)
        
        self.lns_attention = LayerNormScaling(layer_depth=layer_idx + 1)
        self.lns_mlp = LayerNormScaling(layer_depth=layer_idx + 1)
        
        self.gpas_attention = GPAS(config.hidden_size)
        self.gpas_mlp = GPAS(config.hidden_size)
        
        # Canon layers (A+C only)
        # Canon-A: Before attention block
        if config.canon_enabled and config.canon_a_enabled:
            self.canon_a = CanonLayer(config.hidden_size, config.canon_kernel_size)
        else:
            self.canon_a = None
            
        # Canon-C: Before MLP block
        if config.canon_enabled and config.canon_c_enabled:
            self.canon_c = CanonLayer(config.hidden_size, config.canon_kernel_size)
        else:
            self.canon_c = None

    def forward(self, hidden_states: torch.Tensor, position_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        residual = hidden_states
        
        # Apply Canon-A before attention
        if self.canon_a is not None:
            hidden_states = self.canon_a(hidden_states)
        
        attention_input = self.input_layernorm(hidden_states)
        attention_input = self.lns_attention(attention_input)
        attention_output = self.self_attn(attention_input, position_ids, attention_mask)
        hidden_states = residual + attention_output
        hidden_states = self.gpas_attention(hidden_states)
        hidden_states = self.dropout_output(hidden_states)
        
        residual = hidden_states
        
        # Apply Canon-C before MLP
        if self.canon_c is not None:
            hidden_states = self.canon_c(hidden_states)
        
        mlp_input = self.post_attention_layernorm(hidden_states)
        mlp_input = self.lns_mlp(mlp_input)
        mlp_output = self.mlp(mlp_input)
        hidden_states = residual + mlp_output
        hidden_states = self.gpas_mlp(hidden_states)
        hidden_states = self.dropout_output(hidden_states)
        
        return hidden_states
    
    def reset_lax_state(self):
        self.self_attn.reset_lax_state()
        self.mlp.reset_lax_state()

class UnifiedModel(PreTrainedModel):
    """
    UnifiedModel that inherits from PreTrainedModel for full HuggingFace compatibility.
    With AutoClass support for seamless Hub integration.
    """
    config_class = UnifiedModelConfig
    
    # ✅ FIXED: _tied_weights_keys as class attribute (HuggingFace standard)
    _tied_weights_keys = ["lm_head.weight"]
    
    def __init__(self, config: UnifiedModelConfig):
        super().__init__(config)
        self.config = config
        
        if config.vocab_size is None:
            raise ValueError("config.vocab_size must be set from tokenizer before model initialization")
        
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.embedding_dropout = nn.Dropout(config.embedding_dropout)
        self.output_dropout = nn.Dropout(0.05)

        # Create lm_head for output projections
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        self.layers = nn.ModuleList()
        for i in range(config.num_hidden_layers):
            self.layers.append(DecoderLayer(config, i))
        
        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        # Initialize weights
        self.post_init()
        
        self._print_configuration()

    def tie_weights(self):
        """
        ✅ FIXED: Simplified tie_weights method following HuggingFace standard.
        Tie the word embeddings and the output layer.
        This is called automatically if config.tie_word_embeddings is True.
        """
        if self.config.tie_word_embeddings:
            print("🔗 Applying weight tying: lm_head.weight = embed_tokens.weight")
            self.lm_head.weight = self.embed_tokens.weight
            print("✅ Weight tying successful: Parameters are properly shared")

    def _init_weights(self, module):
        """Initialize weights following the custom initialization scheme."""
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.trunc_normal_(module.weight, mean=0.0, std=0.02, a=-0.04, b=0.04)
        elif isinstance(module, CoLA_Linear):
            pass  # CoLA_Linear has its own initialization

    def _print_configuration(self):
        # Conteo ingenuo de todos los parámetros registrados
        total_params_naive = sum(p.numel() for p in self.parameters())
        
        # Conteo inteligente considerando weight tying
        total_params_actual = total_params_naive
        vocab_params = self.config.vocab_size * self.config.hidden_size
        tied_savings = 0
        
        # ✅ CORRECCIÓN: Detectar y ajustar por weight tying real
        if self.config.tie_word_embeddings:
            # Verificar si los tensors están realmente atados en memoria
            embed_weight = self.embed_tokens.weight
            lm_head_weight = self.lm_head.weight
            
            if embed_weight is lm_head_weight:
                # Los tensors son idénticos - restar la duplicación
                tied_savings = vocab_params
                total_params_actual = total_params_naive - tied_savings
            else:
                # Weight tying configurado pero no aplicado aún
                tied_savings = 0
        
        # Cálculos de optimización existentes
        total_linear_params = 0
        total_cola_params = 0
        canon_params = 0
        
        for name, module in self.named_modules():
            if isinstance(module, CoLA_Linear):
                in_features = module.A.in_features
                out_features = module.B.out_features
                rank = module.rank
                
                standard_params = in_features * out_features
                cola_params = (in_features * rank) + (rank * out_features)
                
                total_linear_params += standard_params
                total_cola_params += cola_params
            elif isinstance(module, CanonLayer):
                # Canon layer parameters: depthwise conv1d + bias
                canon_layer_params = module.hidden_dim * module.kernel_size + module.hidden_dim
                canon_params += canon_layer_params
        
        cola_reduction = ((total_linear_params - total_cola_params) / total_linear_params) * 100 if total_linear_params > 0 else 0
        canon_overhead = (canon_params / total_params_actual) * 100 if total_params_actual > 0 else 0
        
        print(f"\n📊 UNIFIED Model + GPAS + LNS + xIELU + CoLA (Linear Only) + LaX + Canon (A+C) + Weight Tying:")
        
        # ✅ MEJORADO: Mostrar conteo real vs ingenuo para transparencia
        if self.config.tie_word_embeddings and tied_savings > 0:
            print(f"🎯 Total Parameters: {total_params_actual/1e6:.2f}M (effective)")
            print(f"📊 Parameter Breakdown:")
            print(f"   • Naive count: {total_params_naive/1e6:.2f}M (all registered params)")
            print(f"   • Actual count: {total_params_actual/1e6:.2f}M (after weight tying)")
            print(f"   • Weight tying savings: {tied_savings/1e6:.2f}M ({tied_savings/total_params_naive*100:.1f}%)")
        else:
            print(f"🎯 Total Parameters: {total_params_actual/1e6:.2f}M")
        
        print(f"📚 DYNAMIC Vocabulary Size: {self.config.vocab_size} (from tokenizer)")
        print(f"🔗 ✅ PROPER Weight Tying: {'ENABLED' if self.config.tie_word_embeddings else 'DISABLED'}")
        
        # ✅ CORRECCIÓN: Mostrar estado real del weight tying
        if self.config.tie_word_embeddings:
            if tied_savings > 0:
                print(f"💾 Weight Tying Status: ✅ ACTIVE (tensors are shared in memory)")
            else:
                print(f"💾 Weight Tying Status: ⏳ CONFIGURED (will be applied during post_init)")
        
        print(f"🚀 ACTIVATION: xIELU (αp_init={self.config.xielu_alpha_p_init}, αn_init={self.config.xielu_alpha_n_init}, β={self.config.xielu_beta})")
        print(f"🔄 UPGRADE: SwiGLU → StandardMLP + xIELU (better efficiency & adaptability)")
        print(f"🗜️ CoLA Integration: {cola_reduction:.1f}% parameter reduction in internal projections")
        print(f"🔀 LaX Enabled: {'YES' if self.config.lax_enabled else 'NO'} (Gate: LINEAR ONLY)")
        print(f"🎼 Canon Layers Enabled: {'YES' if self.config.canon_enabled else 'NO'} (A+C ONLY)")
        if self.config.canon_enabled:
            print(f"   • Canon-A (Before Attention): {'✅' if self.config.canon_a_enabled else '❌'}")
            print(f"   • Canon-B (Inside Attention): ❌ PERMANENTLY DISABLED")
            print(f"   • Canon-C (Before MLP): {'✅' if self.config.canon_c_enabled else '❌'}")
            print(f"   • Canon-D (Inside MLP): ❌ PERMANENTLY DISABLED")
            print(f"   • Canon Kernel Size: {self.config.canon_kernel_size}")
            print(f"   • Canon Parameters Overhead: {canon_overhead:.3f}% ({canon_params/1e3:.1f}K params)")
        print(f"⚡ GPAS Enabled: ALWAYS (Dynamic variance control)")
        print(f"📏 LNS Enabled: ALWAYS (Static depth scaling)")
        print(f"🔧 Variance Control: Triple-level (LNS + GPAS + Canon A+C) ALWAYS")
        print(f"🔗 Residual Connections: STANDARD + HORIZONTAL (Canon A+C only)")
        print(f"🧹 CLEAN: Standard transformer architecture - CrossEntropyLoss manages PAD naturally")
        print(f"⚡ FlashAttention: Scaled Dot-Product Attention with GQA + automatic causal masking")
        print(f"🎯 TOKENIZER AGNOSTIC: Dynamic vocab_size and pad_token_id")
        print(f"🎯 SIMPLIFIED: CoLA Linear Only + Canon A+C Only = Better performance & less overhead")
        print(f"🔗 ✅ FIXED Weight Tying: _tied_weights_keys as class attribute (HF standard)")
        print(f"🎼 Canon A+C BENEFITS: Strategic horizontal information flow with minimal parameters")
        print(f"🤗 HUGGINGFACE COMPATIBLE: Full PreTrainedModel integration v4.53.3")
        print(f"⚡ ✅ NATIVE TORCH COMPILE: Will be handled by TrainingArguments")
        print(f"🚀 ✅ AUTOCLASS SUPPORT: Compatible with AutoConfig.from_pretrained() and AutoModel.from_pretrained()")
    def reset_lax_state(self):
        """Reset LaX state for all layers."""
        for layer in self.layers:
            layer.reset_lax_state()

    def forward(
        self, 
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        **kwargs
    ):
        batch_size, seq_len = input_ids.shape
        
        # Reset LaX state at the beginning of each forward pass
        self.reset_lax_state()
        
        hidden_states = self.embed_tokens(input_ids)
        hidden_states = hidden_states.detach()
        hidden_states = self.embedding_dropout(hidden_states)
        
        for layer in self.layers:
            hidden_states = layer(hidden_states, position_ids=position_ids, attention_mask=attention_mask)
        
        hidden_states = self.norm(hidden_states)
        hidden_states = self.output_dropout(hidden_states)
        
        logits = self.lm_head(hidden_states)
        
        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = nn.CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            
            # ✅ RESTORED: Change pad tokens to -100 so CrossEntropyLoss ignores them (from original code)
            if self.config.pad_token_id is not None:
                shift_labels[shift_labels == self.config.pad_token_id] = -100
            
            loss = loss_fct(shift_logits, shift_labels)
        
        # Return in HuggingFace format
        from transformers.modeling_outputs import CausalLMOutputWithPast
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 50,
        temperature: float = 1.0,
        top_p: float = 0.9,
        top_k: Optional[int] = None,
        do_sample: bool = True,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[int] = None,
        **kwargs
    ) -> torch.Tensor:
        """
        Generate sequences using the model.
        Compatible with AutoModelForCausalLM interface.
        """
        # Set default token IDs
        if pad_token_id is None:
            pad_token_id = self.config.pad_token_id
        if eos_token_id is None:
            eos_token_id = self.config.eos_token_id
        
        batch_size = input_ids.shape[0]
        device = input_ids.device
        
        # Reset LaX state for generation
        self.reset_lax_state()
        
        generated = input_ids.clone()
        
        for _ in range(max_new_tokens):
            # Forward pass
            outputs = self.forward(generated)
            logits = outputs.logits
            
            # Get the logits for the last token
            next_token_logits = logits[:, -1, :]
            
            if do_sample:
                # Apply temperature
                if temperature != 1.0:
                    next_token_logits = next_token_logits / temperature
                
                # Apply top-k filtering
                if top_k is not None:
                    values, indices = torch.topk(next_token_logits, top_k)
                    next_token_logits[next_token_logits < values[:, [-1]]] = -float('inf')
                
                # Apply top-p (nucleus) filtering
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                    
                    # Remove tokens with cumulative probability above the threshold
                    sorted_indices_to_remove = cumulative_probs > top_p
                    # Shift the indices to the right to keep also the first token above the threshold
                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                    sorted_indices_to_remove[..., 0] = 0
                    
                    # Scatter sorted tensors to original indexing
                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                    next_token_logits[indices_to_remove] = -float('inf')
                
                # Sample from the filtered distribution
                probs = F.softmax(next_token_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                # Greedy decoding
                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
            
            # Append the new token
            generated = torch.cat([generated, next_token], dim=1)
            
            # Check for EOS token
            if eos_token_id is not None and (next_token == eos_token_id).all():
                break
        
        return generated


# ✅ AUTOCLASS REGISTRATION - Required for Hub compatibility
# Register the configuration and model for AutoClass support
AutoConfig.register("unified_model", UnifiedModelConfig)
AutoModel.register(UnifiedModelConfig, UnifiedModel)
AutoModelForCausalLM.register(UnifiedModelConfig, UnifiedModel)

print("🚀 ✅ AUTOCLASS REGISTRATION COMPLETE:")
print("   • AutoConfig.register('unified_model', UnifiedModelConfig)")
print("   • AutoModel.register(UnifiedModelConfig, UnifiedModel)")
print("   • AutoModelForCausalLM.register(UnifiedModelConfig, UnifiedModel)")
print("   • Users can now load with: AutoModel.from_pretrained('your-repo', trust_remote_code=True)")