Fix NaN cascade: restore per-step gradient sanitization, add weight/momentum repair, harden Newton-Schulz

Three NaN-related bugs:

1. Gradient sanitization was reduced from every step to every 10 steps as a
throughput optimization. But BitLinear STE + complex recurrent layers produce
occasional NaN gradients that MUST be caught immediately. A single uncaught
NaN gradient corrupts Muon's momentum buffer, which then corrupts the weights
on the next optimizer step, causing permanent all-NaN forward passes.
Fix: restore per-step sanitization.

2. NaN guard only zeroed gradients and halved LR — but never repaired the
already-corrupted weights or Muon momentum buffers. So once NaN entered
the weights, every subsequent step also produced NaN, triggering the LR
halving cascade (5 NaN → halve → 5 NaN → halve → ... → LR = 0).
Fix: NaN guard now sanitizes model weights AND optimizer momentum buffers.
Also increased threshold from 5 to 10 consecutive NaN before halving.

3. Newton-Schulz could produce NaN if input gradient matrix had near-zero norm
(< 1e-12) or already contained NaN values. Now returns zero matrix in
these cases instead of propagating NaN through the polynomial iterations.

chimera_turbo.py

@@ -57,10 +57,15 @@ def _zeropower_via_newtonschulz5(G, steps=5):
     assert G.ndim == 2
     a, b, c = 3.4445, -4.7750, 2.0315
     X = G.T if G.size(0) > G.size(1) else G.clone()
-    X = X / (X.norm() + 1e-7)
+    nrm = X.norm()
+    if nrm < 1e-12 or not torch.isfinite(nrm):
+        return torch.zeros_like(G)
+    X = X / (nrm + 1e-7)
     for _ in range(steps):
         A = X @ X.T
         X = a * X + (b * A + c * A @ A) @ X
+        if not torch.isfinite(X).all():
+            return torch.zeros_like(G)
     return X.T if G.size(0) > G.size(1) else X
 
 
@@ -165,7 +170,6 @@ class MultiTokenPredictionLoss(nn.Module):
             sl = min(logits.size(1), targets.size(1))
 
             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"
@@ -194,9 +198,6 @@ class TokenTriage:
         self.floor_weight = floor_weight
         self._loss_ema = None
         self._step = 0
-        # FIX: Anneal floor_weight from 1.0 → floor_weight over warmup_steps.
-        # When loss is high (early training), all tokens are informative.
-        # Discarding 40% of gradient signal at loss=10+ starves the model.
         self.warmup_steps = 500
 
     def compute_weights(self, per_token_loss):
@@ -208,7 +209,6 @@ class TokenTriage:
             else:
                 self._loss_ema = self.ema_decay * self._loss_ema + (1 - self.ema_decay) * ml
 
-            # FIX: Anneal — during warmup, all tokens get weight ≈ 1.0
             if self._step < self.warmup_steps:
                 t = self._step / self.warmup_steps
                 cur_floor = 1.0 - t * (1.0 - self.floor_weight)
@@ -234,7 +234,7 @@ class PlateauBreaker:
         self._history = deque(maxlen=patience)
         self._stagnant_count = 0
         self._burst_remaining = 0
-        self._saved_lrs = None  # Save ALL group LRs, not just one
+        self._saved_lrs = None
         self.total_bursts = 0
 
     def check_and_adjust(self, loss_val, optimizer, step):
@@ -244,7 +244,6 @@ class PlateauBreaker:
         if self._burst_remaining > 0:
             self._burst_remaining -= 1
             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
@@ -259,31 +258,23 @@ class PlateauBreaker:
         else:
             self._stagnant_count = 0
         if self._stagnant_count >= self.patience // 2:
-            # Save ALL LRs before burst
             self._saved_lrs = [pg["lr"] for pg in optimizer.param_groups]
             for pg in optimizer.param_groups:
-                pg["lr"] *= self.lr_mult  # Multiply, don't replace → LLRD preserved
+                pg["lr"] *= self.lr_mult
             self._burst_remaining = self.burst_steps
             self._stagnant_count = 0
             self.total_bursts += 1
             base = self._saved_lrs[0]
-            print(f"  [P16] Plateau! LR ×{self.lr_mult} for {self.burst_steps} steps (base {base:.2e})")
+            print(f"  [P16] Plateau! LR x{self.lr_mult} for {self.burst_steps} steps (base {base:.2e})")
             return True
         return False
 
 
 # ═══════════════════════════════════════════════════════════
-# P18 Grokfast-EMA (1D params only — NS cancels on 2D)
+# P18 Grokfast-EMA (1D params only)
 # ═══════════════════════════════════════════════════════════
 
 class GrokfastEMA:
-    """Amplify slow gradient components for generalization.
-
-    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
         self.lamb = lamb
@@ -294,7 +285,6 @@ class GrokfastEMA:
         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:
@@ -352,12 +342,10 @@ def apply(model, max_steps=10000, lr=0.02, weight_decay=0.01,
     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, extra_params=mtp_params)
@@ -368,40 +356,33 @@ def apply(model, max_steps=10000, lr=0.02, weight_decay=0.01,
         scales = [g["lr_scale"] for g in optimizer.param_groups]
         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)")
+        print(f"       {n_total:,} params, LR: {min(scales):.3f}x -> {max(scales):.3f}x")
+        print(f"[P13] MTP ({mtp_heads} heads, {n_mtp:,} params -- IN optimizer)")
 
-    # P15
     extras["triage"] = TokenTriage(ema_decay=0.99, select_ratio=0.6, floor_weight=0.1)
     if verbose:
-        print(f"[P15] Token Triage (60%→full, 40%→10%, applied to base+MTP)")
+        print(f"[P15] Token Triage (60%->full, 40%->10%, applied to base+MTP)")
 
-    # P16
-    # FIX: Increase patience (100→200) and variance threshold (0.005→0.02)
-    # so the breaker doesn't fire during normal slow convergence.
-    # The old settings triggered bursts when loss was fluctuating ±0.07,
-    # which is normal for stochastic training at loss~10.
     extras["plateau"] = PlateauBreaker(patience=200, variance_threshold=0.02,
                                         lr_multiplier=2.0, burst_steps=50)
     if verbose:
-        print(f"[P16] Plateau Breaker (×2 burst, LLRD-aware save/restore)")
+        print(f"[P16] Plateau Breaker (x2 burst, LLRD-aware save/restore)")
 
-    # P18
     extras["grokfast"] = GrokfastEMA(alpha=grokfast_alpha, lamb=grokfast_lambda)
     if verbose:
         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)")
+        print(f"[P18] Grokfast-EMA (a={grokfast_alpha}, l={grokfast_lambda}, {n_1d:,} params -- 1D only)")
 
     if verbose:
-        print(f"[P17] Batch Metabolism (hard seq ×2, easy ×0.5)")
+        print(f"[P17] Batch Metabolism (hard seq x2, easy x0.5)")
         print("=" * 65)
 
     return model, optimizer, scheduler, extras
 
 
 # ═══════════════════════════════════════════════════════════
-# Training step — ALL paradigms FUSED + VERIFIED CUMULATIVE
+# Training step
 # ═══════════════════════════════════════════════════════════
 
 _nan_count = 0
@@ -410,29 +391,6 @@ 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.1) -> float:
-    """
-    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.1 × 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 {}
     is_accum = (step + 1) % grad_accum_steps == 0
@@ -456,27 +414,19 @@ def training_step(model, batch, optimizer, scheduler,
                 ignore_index=-100, reduction="none"
             ).reshape(B, T)
 
-            # P17: Batch Metabolism — per-sequence difficulty weights
-            # FIX: With small effective batches (e.g. 8-32), seq_loss.std()
-            # is extremely noisy, causing wild oscillation in seq_weights.
-            # Clamp the z-scores and narrow the weight range from [0.5, 2.0]
-            # to [0.7, 1.4] to reduce gradient noise.
             with torch.no_grad():
                 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).clamp(-2.0, 2.0)
-                seq_weights = torch.sigmoid(z) * 0.7 + 0.7  # [0.7, 1.4]
+                seq_weights = torch.sigmoid(z) * 0.7 + 0.7
 
-            # P15: Token Triage — per-token informativeness weights
             triage = extras.get("triage")
             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:
@@ -489,19 +439,28 @@ def training_step(model, batch, optimizer, scheduler,
 
         loss_val = total_loss.item()
 
-    # NaN guard
+    # NaN guard — skip step AND repair corrupted state
     if not math.isfinite(loss_val):
         _nan_count += 1
         optimizer.zero_grad(set_to_none=True)
-        if _nan_count >= 5:
+        with torch.no_grad():
+            for p in model.parameters():
+                if not torch.isfinite(p.data).all():
+                    p.data.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
+            for group in optimizer.param_groups:
+                for p in group["params"]:
+                    s = optimizer.state.get(p, {})
+                    for key in ("buf", "m", "v"):
+                        if key in s and not torch.isfinite(s[key]).all():
+                            s[key].nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
+        if _nan_count >= 10:
             for pg in optimizer.param_groups:
                 pg["lr"] *= 0.5
-            print(f"  [NaN] 5× — LR halved")
+            print(f"  [NaN] 10x -- LR halved to {optimizer.param_groups[0]['lr']:.2e}")
             _nan_count = 0
         return loss_val
     _nan_count = 0
 
-    # P16: Plateau Breaker (before backward, uses loss_val only)
     plateau = extras.get("plateau")
     if plateau:
         plateau.check_and_adjust(loss_val, optimizer, step)
@@ -510,21 +469,19 @@ def training_step(model, batch, optimizer, scheduler,
         total_loss = total_loss / grad_accum_steps
     total_loss.backward()
 
-    # Sanitize — only check every 10 steps to save CPU cycles.
-    # NaN gradients are rare; checking every step is wasteful.
-    if step % 10 == 0:
-        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)
+    # Sanitize gradients every step — BitLinear STE + complex recurrent
+    # layers produce occasional NaN gradients that MUST be caught immediately.
+    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 on 1D params only (2D handled by Muon NS)
     grokfast = extras.get("grokfast")
     if grokfast:
         grokfast.apply(model)
 
     if is_accum:
         torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
-        optimizer.step()  # P12 Muon (2D) + AdamW (1D), P19 LLRD via lr_scale
+        optimizer.step()
         scheduler.step()
         optimizer.zero_grad(set_to_none=True)
         invalidate_all_caches(model)