Fix loss plateau + throughput collapse: 7 bugs resolved

1. LR was hardcoded at 0.02 (13x too high), now uses args.lr (1.5e-3)
2. GrowLength spent 30% of training at seq=16 (useless fragments), now 10% at seq/2 + 90% at full
3. Token Triage discarded 40% of gradient signal at loss=10+ (no easy tokens exist yet), now anneals floor from 1.0→0.1 over 500 steps
4. LLRD decay=0.85 created 80x LR gap between layers, softened to 0.92 (10x gap)
5. Evolution engine fired 7x per forward pass (expensive HDC/Hamming ops), now 1x
6. MTP heads 3→2 (saves 51M params of gradient overhead), weight 0.3→0.1
7. Batch Metabolism z-scores unclamped with [0.5,2.0] range on B=32 batches, now clamped ±2σ with [0.7,1.4] range

Performance fixes:
- Muon NS steps 5→3 (40% fewer matmuls per optimizer step)
- BitLinear cache lookup amortized (was walking all modules every step)
- Gradient sanitization every 10 steps instead of every step
- Loop classifier bypassed during training (was calling .item() every forward)
- Plateau breaker patience 100→200, variance threshold 0.005→0.02
- Progressive loops delayed: 1→2→3 at 50%/80%/100% (was 20%/60%/100%)

chimera_turbo.py

@@ -66,7 +66,7 @@ def _zeropower_via_newtonschulz5(G, steps=5):
 
 class Muon(torch.optim.Optimizer):
     def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True,
-                 ns_steps=5, weight_decay=0.0,
+                 ns_steps=3, weight_decay=0.0,
                  adamw_betas=(0.9, 0.98), adamw_eps=1e-8):
         defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov,
                         ns_steps=ns_steps, weight_decay=weight_decay,
@@ -193,17 +193,31 @@ class TokenTriage:
         self.select_ratio = select_ratio
         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):
         with torch.no_grad():
+            self._step += 1
             ml = per_token_loss.mean().item()
             if self._loss_ema is None:
                 self._loss_ema = ml
             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)
+            else:
+                cur_floor = self.floor_weight
+
             excess = per_token_loss - self._loss_ema
             thr = torch.quantile(excess.flatten(), 1.0 - self.select_ratio)
-            return torch.where(excess >= thr, 1.0, self.floor_weight)
+            return torch.where(excess >= thr, 1.0, cur_floor)
 
 
 # ═══════════════════════════════════════════════════════════
@@ -294,12 +308,16 @@ class GrokfastEMA:
 # Utilities
 # ═══════════════════════════════════════════════════════════
 
+_bitlinear_cache = []
+
 def invalidate_all_caches(model):
     from chimera.quantization import BitLinear
-    raw = getattr(model, "_orig_mod", model)
-    for m in raw.modules():
-        if isinstance(m, BitLinear):
-            m.invalidate_packed()
+    global _bitlinear_cache
+    if not _bitlinear_cache:
+        raw = getattr(model, "_orig_mod", model)
+        _bitlinear_cache = [m for m in raw.modules() if isinstance(m, BitLinear)]
+    for m in _bitlinear_cache:
+        m.invalidate_packed()
 
 
 def create_scheduler(optimizer, max_steps, warmup_steps=200):
@@ -359,7 +377,11 @@ def apply(model, max_steps=10000, lr=0.02, weight_decay=0.01,
         print(f"[P15] Token Triage (60%→full, 40%→10%, applied to base+MTP)")
 
     # P16
-    extras["plateau"] = PlateauBreaker(patience=100, variance_threshold=0.005,
+    # 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)")
@@ -387,7 +409,7 @@ _nan_count = 0
 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:
+                  mtp_weight=0.1) -> float:
     """
     Data flow (verified cumulative):
 
@@ -401,7 +423,7 @@ def training_step(model, batch, optimizer, scheduler,
          ├─ base_loss = weighted_mean(per_token_loss, combined)
          │
          ├─ P13: mtp_loss = MTP(hidden, labels, tok_weights)  ← triage applied!
-         ├─ total_loss = base + 0.3 × mtp
+         ├─ total_loss = base + 0.1 × mtp
          │
     backward(total_loss) → param.grad for ALL params (model + MTP heads)
          │
@@ -435,12 +457,16 @@ def training_step(model, batch, optimizer, scheduler,
             ).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
-                seq_weights = torch.sigmoid(z) * 1.5 + 0.5  # [0.5, 2.0]
+                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]
 
             # P15: Token Triage — per-token informativeness weights
             triage = extras.get("triage")
@@ -484,10 +510,12 @@ def training_step(model, batch, optimizer, scheduler,
         total_loss = total_loss / grad_accum_steps
     total_loss.backward()
 
-    # 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)
+    # 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)
 
     # P18: Grokfast on 1D params only (2D handled by Muon NS)
     grokfast = extras.get("grokfast")