RESULT 6: wall-2 = EP-specific marginal under-convergence at the contractivity edge (A skips16/B skips2/C skips1/D-BP clean); ship --kretry (retry marginal batches at K=kmax); launch OLMo2 parity matrix with parity-gated epoch autolaunch

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 16 insertions, 4 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index f3b8736..5a2677c 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -37,6 +37,7 @@ ap.add_argument('--sig0', type=float, default=-1.0)      # override SIG0 (beta-s
 ap.add_argument('--olmo2', action='store_true')          # OLMo2-standard block: norm-AFTER-sublayer RMSNorm, full-width QK-norm, RoPE(500k), SwiGLU, no-bias, untied head, final RMSNorm, 0.02 init
 ap.add_argument('--wd', type=float, default=-1.0)        # >=0: grouped weight decay (linear weights+head decay; embeddings/norm-gains none). <0 = legacy uniform 1e-4
 ap.add_argument('--zloss', type=float, default=0.0)      # z-loss coefficient on train objective (OLMo2-style logit regularizer); 0 = off
+ap.add_argument('--kretry', type=int, default=0)         # >0: on drift-reject, RETRY the batch once with this many fb rounds (diag B: K8 converges the marginal batches) instead of dropping it
 ap.add_argument('--dtop_every', type=int, default=1)    # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr)
 ap.add_argument('--gate_every', type=int, default=200)  # in-training cos(EP,BP) telemetry; <=0 = fully BP-free (no bp_gate at all)
 ap.add_argument('--gate_govern', action='store_true')   # let gate cos adjust K/bscale (default: observe-only => training control is BP-free)
@@ -290,9 +291,20 @@ def ep_step(x, y):
         drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max(
             sum(float(b.norm()) for b in zs_free), 1e-9)
     if (not math.isfinite(drift)) or (drift > 0.5 and not args.noguard):
-        GOV['skd'] = GOV.get('skd', 0) + 1   # drift-guard reject (relaxation non-convergence)
-        for p in all_params: p.grad = None
-        return free_ce, beta_t, GOV['K'], False
+        ok_retry = False
+        if args.kretry > 0 and math.isfinite(drift) and not args.noguard:
+            GOV['skr'] = GOV.get('skr', 0) + 1   # marginal batch: retry once with deeper relaxation
+            z0, zs, ins, outs = free_states_graphed(x)
+            zs_free = [z.clone() for z in zs]
+            zp, last_outs = relax(z0, zs, ins, outs, y, +beta_t, args.kretry, x)
+            with torch.no_grad():
+                drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max(
+                    sum(float(b.norm()) for b in zs_free), 1e-9)
+            ok_retry = math.isfinite(drift) and drift <= 0.5
+        if not ok_retry:
+            GOV['skd'] = GOV.get('skd', 0) + 1   # drift-guard reject (relaxation non-convergence)
+            for p in all_params: p.grad = None
+            return free_ce, beta_t, GOV['K'], False
     GOV['drift'] = drift
     E = 0.0
     for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach() - o) ** 2).sum()
@@ -368,7 +380,7 @@ for step in range(start_step, args.steps + 1):
         val = evaluate(); best = min(best, val)
         gtag = '' if math.isnan(gcos) else f' cos={gcos:.4f}'
         print(f'step {step:5d}/{args.steps} | train {ce:.4f} val {val:.4f} (best {best:.4f}) '
-              f'| beta={beta_t:.2e} K={rounds} skips={skips}(d{GOV.get("skd",0)}/g{GOV.get("skg",0)}){gtag} '
+              f'| beta={beta_t:.2e} K={rounds} skips={skips}(d{GOV.get("skd",0)}/g{GOV.get("skg",0)}/r{GOV.get("skr",0)}){gtag} '
               f'drift={GOV["drift"]:.3f} gn={GOV["gn"]:.2e} sig={GOV["sig"]:.1f} | {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True)
         if wb is not None:
             try: wb.log({'train_ce': ce, 'val_ce': val, 'best': best, 'beta_t': beta_t,