fix: make model loadable via AutoModelForCausalLM (add auto_map, flatten module files to repo root, inherit GenerationMixin)

differential.py

@@ -0,0 +1,135 @@
+"""
+Differential Transformer (PRODUCTION - causal-safe)
+====================================================
+
+Microsoft 2410.05258 (ICLR 2025)
+Cancels noise via dual-softmax attention map subtraction.
+
+Core formula:
+  DiffAttn(X) = (softmax(Q1 K1^T / sqrt(d)) - λ × softmax(Q2 K2^T / sqrt(d))) V
+  λ = exp(λ_q1 · λ_k1) - exp(λ_q2 · λ_k2) + λ_init
+
+Implementation notes:
+  - GroupNorm (sequence-spanning) → LayerNorm(head_dim) (per-token, causal-safe)
+  - Causal triu mask explicit (was relying on SDPA, but split + softmax direct now)
+  - forward signature: layer_idx kwarg + (Tensor, past_kv) tuple return for modeling compat
+
+Hyperparameter (configuration_aether_v2_7way.py):
+  - diff_lambda_init = 0.8
+"""
+from __future__ import annotations
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DifferentialAttention(nn.Module):
+    """λ-gated dual-softmax differential attention with causal mask + LayerNorm."""
+
+    def __init__(self, config, layer_idx: int = 0):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.h = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        # Differential splits head_dim in half
+        self.diff_head_dim = self.head_dim // 2
+
+        # Layer-depth lambda init (paper Eq. 3): λ_init = 0.8 - 0.6 * exp(-0.3 * (depth - 1))
+        # Use config base + depth decay if provided.
+        base = getattr(config, "diff_lambda_init", 0.8)
+        decay = getattr(config, "diff_lambda_layer_decay", 0.6)
+        # Simple per-layer init
+        self.lambda_init = base - decay * math.exp(-0.3 * max(layer_idx - 1, 0))
+
+        # Q/K/V/O projections
+        self.q_proj = nn.Linear(self.h, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.h, bias=False)
+
+        # Lambda parameters per-head
+        self.lambda_q1 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))
+        self.lambda_k1 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))
+        self.lambda_q2 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))
+        self.lambda_k2 = nn.Parameter(torch.zeros(self.num_heads, self.diff_head_dim).normal_(0, 0.1))
+
+        # ★ FIX 5/7: LayerNorm(head_dim) — per-token norm, causal-safe.
+        # Was nn.GroupNorm(num_groups=num_heads, num_channels=num_heads*head_dim) which
+        # normalized over sequence dim and leaked future info.
+        self.subln = nn.LayerNorm(self.head_dim, eps=getattr(config, "rms_norm_eps", 1e-6))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value=None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[object]]:
+        B, S, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
+        k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)
+        v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)
+
+        # AetherCache: cache full k,v (pre-split) in HF layout [B, kv_h, S, D]
+        if past_key_value is not None:
+            kc, vc = past_key_value.update(k.transpose(1, 2), v.transpose(1, 2), getattr(self, "cache_idx", self.layer_idx))
+            k = kc.transpose(1, 2)
+            v = vc.transpose(1, 2)
+        S_kv = k.shape[1]
+
+        # Split Q, K halves
+        q1, q2 = q.split(self.diff_head_dim, dim=-1)  # [B, S, H, d/2]
+        k1, k2 = k.split(self.diff_head_dim, dim=-1)  # [B, S_kv, kv_h, d/2]
+
+        # GQA repeat
+        repeat = self.num_heads // self.num_kv_heads
+        k1 = k1.repeat_interleave(repeat, dim=2)
+        k2 = k2.repeat_interleave(repeat, dim=2)
+        v = v.repeat_interleave(repeat, dim=2)
+
+        # Transpose to [B, H, S, d]
+        q1 = q1.transpose(1, 2)
+        q2 = q2.transpose(1, 2)
+        k1 = k1.transpose(1, 2)
+        k2 = k2.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        scale = 1.0 / math.sqrt(self.diff_head_dim)
+        scores1 = torch.matmul(q1, k1.transpose(-2, -1)) * scale
+        scores2 = torch.matmul(q2, k2.transpose(-2, -1)) * scale
+
+        # ★ FIX 5/7: causal mask (triu(1) blocks future)
+        # AetherCache: query i is at absolute pos (S_kv - S + i) -> shift the diagonal
+        causal = torch.ones(S, S_kv, device=q1.device, dtype=torch.bool).triu(1 + S_kv - S)
+        scores1 = scores1.masked_fill(causal, float("-inf"))
+        scores2 = scores2.masked_fill(causal, float("-inf"))
+
+        attn1 = F.softmax(scores1, dim=-1)
+        attn2 = F.softmax(scores2, dim=-1)
+
+        # Lambda per-head (paper Eq. 3)
+        lam_q1k1 = (self.lambda_q1 * self.lambda_k1).sum(dim=-1)  # [H]
+        lam_q2k2 = (self.lambda_q2 * self.lambda_k2).sum(dim=-1)
+        lam = torch.exp(lam_q1k1) - torch.exp(lam_q2k2) + self.lambda_init  # [H]
+        lam = lam.view(1, -1, 1, 1)
+
+        diff_attn = attn1 - lam * attn2
+
+        out = torch.matmul(diff_attn, v)  # [B, H, S, head_dim]
+        # ★ FIX 5/7: per-token LayerNorm (was sequence-spanning GroupNorm = leak)
+        out = self.subln(out)
+        out = out * (1 - self.lambda_init)
+
+        out = out.transpose(1, 2).reshape(B, S, -1)
+        return self.o_proj(out), past_key_value
+
+
+__all__ = ["DifferentialAttention"]