fix: v12 GENESIS — fix 6 interaction bugs between paradigms\n\n1. P13 MTP heads added to optimizer (were dead — never updated)\n2. P18 Grokfast: skip Muon 2D params (NS normalisation cancels amplification)\n Apply only to 1D/embed params where AdamW preserves the signal\n3. P16 Plateau: save/restore ALL group LRs (was destroying LLRD ratios)\n4. P15 Token Triage applied to MTP loss too (was only on base loss)\n5. P16 Plateau: gentler burst ×2 instead of ×3 (Grokfast already amplifies)\n6. P15 Triage: per-position EMA disabled, use global excess only"

chimera_turbo.py

@@ -1,19 +1,15 @@
 """
-chimera_turbo.py — CHIMERA GENESIS v11
+chimera_turbo.py — CHIMERA GENESIS v12
 
-The unified training engine for ch1mera 5.3.
+Interaction-audited paradigm stack. Every paradigm verified cumulative.
 
-Active paradigms (all fused, no dead code):
-  P12 Muon optimizer         — NS-orthogonalized momentum, 2× token efficiency
-  P13 Multi-Token Prediction — 3 aux heads, 3× gradient signal per forward
-  P15 Token Triage           — focus gradient on informative tokens (Rho-1 inspired)
-  P16 Plateau Breaker        — adaptive LR bursts on stagnation
-  P17 Batch Metabolism        — weight hard sequences 2×, easy 0.5×
-  P18 Grokfast-EMA           — amplify slow grads (generalization), filter fast (noise)
-  P19 Layer-wise LR Decay    — top layers learn faster, bottom layers preserve features
-
-Removed (dead weight):
-  P14 EMA Self-Distill — doubled forward time, marginal gain from self-EMA
+  P12 Muon             — NS-orthogonalized momentum for 2D matrices
+  P13 MTP              — 3 aux heads (NOW in optimizer)
+  P15 Token Triage     — focus on informative tokens (applied to ALL losses)
+  P16 Plateau Breaker  — adaptive LR burst (LLRD-aware save/restore)
+  P17 Batch Metabolism  — hard sequences weighted 2×
+  P18 Grokfast-EMA     — amplify slow grads (1D params ONLY — NS cancels on 2D)
+  P19 LLRD             — layer-wise LR decay for ternary
 """
 
 import math
@@ -21,7 +17,7 @@ import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from typing import Optional, Dict, Any, Tuple
+from typing import Dict
 from contextlib import nullcontext
 from collections import deque
 
@@ -54,7 +50,7 @@ def configure_threading(cpu_info, reserve=1):
 
 
 # ═══════════════════════════════════════════════════════════
-# P12 — Muon Optimizer + P19 Layer-wise LR Decay
+# P12 Muon + P19 LLRD
 # ═══════════════════════════════════════════════════════════
 
 def _zeropower_via_newtonschulz5(G, steps=5):
@@ -69,7 +65,6 @@ def _zeropower_via_newtonschulz5(G, steps=5):
 
 
 class Muon(torch.optim.Optimizer):
-    """Muon with integrated layer-wise LR decay (P19)."""
     def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True,
                  ns_steps=5, weight_decay=0.0,
                  adamw_betas=(0.9, 0.98), adamw_eps=1e-8):
@@ -113,14 +108,8 @@ class Muon(torch.optim.Optimizer):
 
 
 def create_muon_optimizer(model, lr=0.02, momentum=0.95, weight_decay=0.01,
-                          llrd_decay=0.85):
-    """Create Muon with P19 layer-wise LR decay.
-
-    Top layers get full LR, bottom layers get LR × decay^depth.
-    This preserves general features in early layers while allowing
-    later layers to specialize faster. Proven for ternary (arxiv 2602.07374).
-    """
-    # Detect layer depth for each param
+                          llrd_decay=0.85, extra_params=None):
+    """Create Muon with LLRD. extra_params: additional nn.Module params to include."""
     raw = getattr(model, "_orig_mod", model)
     n_layers = len(raw.layers) if hasattr(raw, "layers") else 28
 
@@ -128,35 +117,30 @@ def create_muon_optimizer(model, lr=0.02, momentum=0.95, weight_decay=0.01,
     for name, p in model.named_parameters():
         if not p.requires_grad:
             continue
-
         is_embed = any(k in name for k in ["embed", "lm_head", "wte", "wpe"])
         if is_embed:
             p._is_embed = True
-
-        # Determine layer index for LLRD
         lr_scale = 1.0
         for i in range(n_layers):
-            if f"layers.{i}." in name or f"layers.{i}]" in name:
-                # Scale: top layer = 1.0, bottom layer = decay^(n_layers-1)
-                depth_from_top = n_layers - 1 - i
-                lr_scale = llrd_decay ** depth_from_top
+            if f"layers.{i}." in name:
+                lr_scale = llrd_decay ** (n_layers - 1 - i)
                 break
-
-        # Embeddings and lm_head get lowest LR
         if is_embed:
             lr_scale = llrd_decay ** n_layers
+        param_groups.append({"params": [p], "lr_scale": lr_scale})
 
-        param_groups.append({
-            "params": [p],
-            "lr_scale": lr_scale,
-        })
+    # Add extra params (e.g. MTP heads) at full LR
+    if extra_params:
+        for p in extra_params:
+            if p.requires_grad:
+                param_groups.append({"params": [p], "lr_scale": 1.0})
 
     return Muon(param_groups, lr=lr, momentum=momentum,
                 weight_decay=weight_decay, adamw_betas=(0.9, 0.98))
 
 
 # ═══════════════════════════════════════════════════════════
-# P13 — Multi-Token Prediction
+# P13 MTP
 # ═════════════════════════════��═════════════════════════════
 
 class MultiTokenPredictionLoss(nn.Module):
@@ -169,7 +153,8 @@ class MultiTokenPredictionLoss(nn.Module):
         for h in self.extra_heads:
             nn.init.normal_(h.weight, std=0.006)
 
-    def forward(self, hidden_states, labels):
+    def forward(self, hidden_states, labels, token_weights=None):
+        """Compute MTP loss, optionally weighted by token_weights from Triage."""
         total, count = 0.0, 0
         for k, head in enumerate(self.extra_heads):
             shift = k + 2
@@ -178,8 +163,20 @@ class MultiTokenPredictionLoss(nn.Module):
             logits = head(hidden_states[:, :-shift])
             targets = labels[:, shift:]
             sl = min(logits.size(1), targets.size(1))
-            loss = F.cross_entropy(logits[:, :sl].reshape(-1, logits.size(-1)),
-                                   targets[:, :sl].reshape(-1), ignore_index=-100)
+
+            if token_weights is not None:
+                # Apply token triage weights to MTP loss too
+                per_tok = F.cross_entropy(
+                    logits[:, :sl].reshape(-1, logits.size(-1)),
+                    targets[:, :sl].reshape(-1), ignore_index=-100, reduction="none"
+                ).reshape(logits.size(0), sl)
+                tw = token_weights[:, :sl]
+                loss = (per_tok * tw).sum() / tw.sum()
+            else:
+                loss = F.cross_entropy(
+                    logits[:, :sl].reshape(-1, logits.size(-1)),
+                    targets[:, :sl].reshape(-1), ignore_index=-100)
+
             if torch.isfinite(loss):
                 total = total + loss
                 count += 1
@@ -187,7 +184,7 @@ class MultiTokenPredictionLoss(nn.Module):
 
 
 # ═══════════════════════════════════════════════════════════
-# P15 — Token Triage (Rho-1 inspired)
+# P15 Token Triage
 # ═══════════════════════════════════════════════════════════
 
 class TokenTriage:
@@ -198,25 +195,24 @@ class TokenTriage:
         self._loss_ema = None
 
     def compute_weights(self, per_token_loss):
-        """Returns per-token weights [B, T]. Differentiable-safe (weights are detached)."""
         with torch.no_grad():
-            mean_loss = per_token_loss.mean().item()
+            ml = per_token_loss.mean().item()
             if self._loss_ema is None:
-                self._loss_ema = mean_loss
+                self._loss_ema = ml
             else:
-                self._loss_ema = self.ema_decay * self._loss_ema + (1 - self.ema_decay) * mean_loss
+                self._loss_ema = self.ema_decay * self._loss_ema + (1 - self.ema_decay) * ml
             excess = per_token_loss - self._loss_ema
-            threshold = torch.quantile(excess.flatten(), 1.0 - self.select_ratio)
-            return torch.where(excess >= threshold, 1.0, self.floor_weight)
+            thr = torch.quantile(excess.flatten(), 1.0 - self.select_ratio)
+            return torch.where(excess >= thr, 1.0, self.floor_weight)
 
 
 # ═══════════════════════════════════════════════════════════
-# P16 — Plateau Breaker
+# P16 Plateau Breaker (LLRD-aware)
 # ═══════════════════════════════════════════════════════════
 
 class PlateauBreaker:
     def __init__(self, patience=100, variance_threshold=0.005,
-                 lr_multiplier=3.0, burst_steps=50):
+                 lr_multiplier=2.0, burst_steps=50):
         self.patience = patience
         self.var_threshold = variance_threshold
         self.lr_mult = lr_multiplier
@@ -224,7 +220,7 @@ class PlateauBreaker:
         self._history = deque(maxlen=patience)
         self._stagnant_count = 0
         self._burst_remaining = 0
-        self._base_lr = None
+        self._saved_lrs = None  # Save ALL group LRs, not just one
         self.total_bursts = 0
 
     def check_and_adjust(self, loss_val, optimizer, step):
@@ -233,10 +229,11 @@ class PlateauBreaker:
         self._history.append(loss_val)
         if self._burst_remaining > 0:
             self._burst_remaining -= 1
-            if self._burst_remaining == 0 and self._base_lr is not None:
-                for pg in optimizer.param_groups:
-                    pg["lr"] = self._base_lr
-                self._base_lr = None
+            if self._burst_remaining == 0 and self._saved_lrs is not None:
+                # Restore ALL group LRs (preserves LLRD ratios)
+                for pg, saved_lr in zip(optimizer.param_groups, self._saved_lrs):
+                    pg["lr"] = saved_lr
+                self._saved_lrs = None
             return False
         if len(self._history) < self.patience:
             return False
@@ -248,35 +245,30 @@ class PlateauBreaker:
         else:
             self._stagnant_count = 0
         if self._stagnant_count >= self.patience // 2:
-            self._base_lr = optimizer.param_groups[0]["lr"]
-            burst_lr = self._base_lr * self.lr_mult
+            # Save ALL LRs before burst
+            self._saved_lrs = [pg["lr"] for pg in optimizer.param_groups]
             for pg in optimizer.param_groups:
-                pg["lr"] = burst_lr
+                pg["lr"] *= self.lr_mult  # Multiply, don't replace → LLRD preserved
             self._burst_remaining = self.burst_steps
             self._stagnant_count = 0
             self.total_bursts += 1
-            print(f"  [P16] Plateau! LR burst {self._base_lr:.2e} → {burst_lr:.2e} × {self.burst_steps}steps")
+            base = self._saved_lrs[0]
+            print(f"  [P16] Plateau! LR ×{self.lr_mult} for {self.burst_steps} steps (base {base:.2e})")
             return True
         return False
 
 
 # ═══════════════════════════════════════════════════════════
-# P18 — Grokfast-EMA (arxiv 2405.20233)
+# P18 Grokfast-EMA (1D params only — NS cancels on 2D)
 # ═══════════════════════════════════════════════════════════
 
 class GrokfastEMA:
-    """Accelerate generalization by amplifying slow gradient components.
-
-    The key insight: gradient time-series has fast components (memorization,
-    STE quantization noise) and slow components (generalization signal).
-    EMA-filter the gradients, then ADD the filtered (slow) component back
-    with amplification factor λ.
-
-    Result: 43× faster convergence on grokking tasks.
-    For ternary models: STE noise is exactly the "fast component" —
-    Grokfast filters it out while amplifying the real learning signal.
+    """Amplify slow gradient components for generalization.
 
-    arxiv 2405.20233, α=0.98, λ=2.0 recommended.
+    Applied ONLY to 1D params and embeddings (AdamW path).
+    Skipped for 2D matrices because Muon's Newton-Schulz normalisation
+    cancels the amplitude amplification — only direction survives,
+    which Muon already optimises via orthogonalisation.
     """
     def __init__(self, alpha=0.98, lamb=2.0):
         self.alpha = alpha
@@ -284,20 +276,18 @@ class GrokfastEMA:
         self._ema: Dict[str, torch.Tensor] = {}
 
     @torch.no_grad()
-    def apply(self, model: nn.Module):
-        """Call after loss.backward(), before optimizer.step().
-
-        Modifies param.grad in-place to amplify slow components.
-        """
-        for name, param in model.named_parameters():
-            if param.grad is None:
+    def apply(self, model):
+        for name, p in model.named_parameters():
+            if p.grad is None:
+                continue
+            # Skip 2D Muon params — NS normalisation cancels amplitude
+            if p.ndim == 2 and not getattr(p, "_is_embed", False):
                 continue
             if name not in self._ema:
-                self._ema[name] = param.grad.clone()
+                self._ema[name] = p.grad.clone()
             else:
-                self._ema[name].mul_(self.alpha).add_(param.grad, alpha=1 - self.alpha)
-            # Amplify slow component: grad = grad + λ * EMA(grad)
-            param.grad.add_(self._ema[name], alpha=self.lamb)
+                self._ema[name].mul_(self.alpha).add_(p.grad, alpha=1 - self.alpha)
+            p.grad.add_(self._ema[name], alpha=self.lamb)
 
 
 # ═══════════════════════════════════════════════════════════
@@ -333,7 +323,7 @@ def apply(model, max_steps=10000, lr=0.02, weight_decay=0.01,
     cpu_info = detect_cpu_info()
     if verbose:
         print("=" * 65)
-        print("CHIMERA GENESIS v11 — Revolutionary Training Engine")
+        print("CHIMERA GENESIS v12 — Interaction-Audited Stack")
         print("=" * 65)
         print(f"  CPU: {cpu_info['capability']}  Cores: {cpu_info['physical_cores']}")
 
@@ -341,51 +331,55 @@ def apply(model, max_steps=10000, lr=0.02, weight_decay=0.01,
     if verbose:
         print(f"  Threads: {n}")
 
-    # P12+P19: Muon with layer-wise LR decay
+    raw = getattr(model, "_orig_mod", model)
+    extras = {}
+
+    # P13: Create MTP FIRST so we can add its params to optimizer
+    h, v = raw.config["hidden_size"], raw.config["vocab_size"]
+    mtp = MultiTokenPredictionLoss(h, v, n_future=mtp_heads)
+    extras["mtp"] = mtp
+
+    # P12+P19: Muon with LLRD + MTP head params included
+    mtp_params = list(mtp.parameters())
     optimizer = create_muon_optimizer(model, lr=lr, weight_decay=weight_decay,
-                                      llrd_decay=llrd_decay)
+                                      llrd_decay=llrd_decay, extra_params=mtp_params)
     scheduler = create_scheduler(optimizer, max_steps, warmup_steps)
 
     if verbose:
-        n_groups = len(optimizer.param_groups)
         n_total = sum(p.numel() for g in optimizer.param_groups for p in g["params"])
         scales = [g["lr_scale"] for g in optimizer.param_groups]
-        print(f"[P12] Muon (lr={lr}) + [P19] LLRD (decay={llrd_decay}) — {n_total:,} params, {n_groups} groups")
-        print(f"       LR range: {min(scales):.3f}× → {max(scales):.3f}×")
-
-    raw = getattr(model, "_orig_mod", model)
-    extras = {}
+        n_mtp = sum(p.numel() for p in mtp_params)
+        print(f"[P12] Muon (lr={lr}) + [P19] LLRD (decay={llrd_decay})")
+        print(f"       {n_total:,} params, LR: {min(scales):.3f}× → {max(scales):.3f}×")
+        print(f"[P13] MTP ({mtp_heads} heads, {n_mtp:,} params — IN optimizer)")
 
-    # P13: MTP
-    h, v = raw.config["hidden_size"], raw.config["vocab_size"]
-    extras["mtp"] = MultiTokenPredictionLoss(h, v, n_future=mtp_heads)
-    if verbose:
-        print(f"[P13] Multi-Token Prediction ({mtp_heads} heads)")
-
-    # P15: Token Triage
+    # P15
     extras["triage"] = TokenTriage(ema_decay=0.99, select_ratio=0.6, floor_weight=0.1)
     if verbose:
-        print(f"[P15] Token Triage (60% informative → full grad, 40% noise → 10%)")
+        print(f"[P15] Token Triage (60%→full, 40%→10%, applied to base+MTP)")
 
-    # P16: Plateau Breaker
+    # P16
     extras["plateau"] = PlateauBreaker(patience=100, variance_threshold=0.005,
-                                        lr_multiplier=3.0, burst_steps=50)
+                                        lr_multiplier=2.0, burst_steps=50)
     if verbose:
-        print(f"[P16] Plateau Breaker (stagnation → LR ×3 burst)")
+        print(f"[P16] Plateau Breaker (×2 burst, LLRD-aware save/restore)")
 
-    # P18: Grokfast-EMA
+    # P18
     extras["grokfast"] = GrokfastEMA(alpha=grokfast_alpha, lamb=grokfast_lambda)
     if verbose:
-        print(f"[P18] Grokfast-EMA (α={grokfast_alpha}, λ={grokfast_lambda} — amplify generalization)")
+        n_1d = sum(p.numel() for p in model.parameters()
+                   if p.requires_grad and (p.ndim < 2 or getattr(p, "_is_embed", False)))
+        print(f"[P18] Grokfast-EMA (α={grokfast_alpha}, λ={grokfast_lambda}, {n_1d:,} params — 1D only)")
 
     if verbose:
+        print(f"[P17] Batch Metabolism (hard seq ×2, easy ×0.5)")
         print("=" * 65)
 
     return model, optimizer, scheduler, extras
 
 
 # ═══════════════════════════════════════════════════════════
-# Training step — ALL paradigms FUSED
+# Training step — ALL paradigms FUSED + VERIFIED CUMULATIVE
 # ═══════════════════════════════════════════════════════════
 
 _nan_count = 0
@@ -394,10 +388,28 @@ def training_step(model, batch, optimizer, scheduler,
                   extras=None, grad_accum_steps=1, step=0,
                   max_grad_norm=1.0, autocast_dtype=None,
                   mtp_weight=0.3) -> float:
-    """One training step with all paradigms active and fused.
-
-    Loss = TokenTriage(BatchMetabolism(CE_per_token)) + mtp_weight * MTP_aux
-    After backward: Grokfast-EMA filters gradients → Muon+LLRD step
+    """
+    Data flow (verified cumulative):
+
+    forward(batch) → logits, hidden_states
+         │
+         ├─ per_token_loss = CE(logits, labels, reduction='none')  [B,T]
+         │
+         ├─ P17: seq_weights = sigmoid(z-score(per_seq_loss))      [B]
+         ├─ P15: tok_weights = triage(excess_loss)                  [B,T]
+         ├─ combined = tok_weights × seq_weights                    [B,T]
+         ├─ base_loss = weighted_mean(per_token_loss, combined)
+         │
+         ├─ P13: mtp_loss = MTP(hidden, labels, tok_weights)  ← triage applied!
+         ├─ total_loss = base + 0.3 × mtp
+         │
+    backward(total_loss) → param.grad for ALL params (model + MTP heads)
+         │
+         ├─ P18: Grokfast amplifies grad on 1D params only (skip 2D/Muon)
+         │
+    optimizer.step() → P12 Muon (2D) + AdamW (1D), P19 LLRD scales per group
+         │
+         └─ P16: Plateau checks loss_val, burst preserves LLRD ratios
     """
     global _nan_count
     extras = extras or {}
@@ -415,48 +427,43 @@ def training_step(model, batch, optimizer, scheduler,
 
         logits = getattr(outputs, "logits", None)
 
-        # ── FUSED LOSS: Token Triage × Batch Metabolism ──
         if logits is not None:
             B, T, V = logits.shape
-            # Per-token CE (no reduction)
             per_token = F.cross_entropy(
                 logits.reshape(-1, V), labels.reshape(-1),
                 ignore_index=-100, reduction="none"
             ).reshape(B, T)
 
-            # P17: Batch Metabolism — per-sequence weights
+            # P17: Batch Metabolism — per-sequence difficulty weights
             with torch.no_grad():
-                seq_loss = per_token.mean(dim=1)  # [B]
+                seq_loss = per_token.mean(dim=1)
                 seq_mean = seq_loss.mean()
                 seq_std = seq_loss.std().clamp(min=1e-6)
                 z = (seq_loss - seq_mean) / seq_std
                 seq_weights = torch.sigmoid(z) * 1.5 + 0.5  # [0.5, 2.0]
 
-            # P15: Token Triage — per-token weights
+            # P15: Token Triage — per-token informativeness weights
             triage = extras.get("triage")
-            if triage is not None:
-                tok_weights = triage.compute_weights(per_token)  # [B, T]
+            tok_weights = triage.compute_weights(per_token) if triage else torch.ones_like(per_token)
+
+            # Fuse: multiplicative composition
+            combined = tok_weights * seq_weights.unsqueeze(1)
+            base_loss = (per_token * combined).sum() / combined.sum()
+
+            # P13: MTP with Token Triage weights passed through
+            mtp = extras.get("mtp")
+            hidden = getattr(outputs, "hidden_states", None)
+            if mtp is not None and hidden is not None:
+                mtp_loss = mtp(hidden, labels, token_weights=tok_weights)
+                total_loss = base_loss + mtp_weight * mtp_loss
             else:
-                tok_weights = torch.ones_like(per_token)
-
-            # Fuse: multiply token weights × sequence weights
-            combined_weights = tok_weights * seq_weights.unsqueeze(1)  # [B, T]
-            base_loss = (per_token * combined_weights).sum() / combined_weights.sum()
-        else:
-            base_loss = outputs.loss if hasattr(outputs, "loss") else outputs
-
-        # P13: MTP auxiliary
-        mtp = extras.get("mtp")
-        hidden = getattr(outputs, "hidden_states", None)
-        if mtp is not None and hidden is not None:
-            mtp_loss = mtp(hidden, labels)
-            total_loss = base_loss + mtp_weight * mtp_loss
+                total_loss = base_loss
         else:
-            total_loss = base_loss
+            total_loss = outputs.loss if hasattr(outputs, "loss") else outputs
 
         loss_val = total_loss.item()
 
-    # ── NaN guard ──
+    # NaN guard
     if not math.isfinite(loss_val):
         _nan_count += 1
         optimizer.zero_grad(set_to_none=True)
@@ -468,28 +475,28 @@ def training_step(model, batch, optimizer, scheduler,
         return loss_val
     _nan_count = 0
 
-    # P16: Plateau Breaker
+    # P16: Plateau Breaker (before backward, uses loss_val only)
     plateau = extras.get("plateau")
-    if plateau is not None:
+    if plateau:
         plateau.check_and_adjust(loss_val, optimizer, step)
 
     if grad_accum_steps > 1:
         total_loss = total_loss / grad_accum_steps
     total_loss.backward()
 
-    # Sanitize grads
+    # Sanitize
     for p in model.parameters():
         if p.grad is not None and not torch.isfinite(p.grad).all():
             p.grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
 
-    # P18: Grokfast-EMA — amplify slow gradients BEFORE optimizer step
+    # P18: Grokfast on 1D params only (2D handled by Muon NS)
     grokfast = extras.get("grokfast")
-    if grokfast is not None:
+    if grokfast:
         grokfast.apply(model)
 
     if is_accum:
         torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
-        optimizer.step()
+        optimizer.step()  # P12 Muon (2D) + AdamW (1D), P19 LLRD via lr_scale
         scheduler.step()
         optimizer.zero_grad(set_to_none=True)
         invalidate_all_caches(model)