beta-cap-by-loop-gain controller (--beta_cap_rho): rho^ meter from sweep residuals, cap OVERRIDES floor near the wall; smoke passed

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

1 files changed, 18 insertions, 2 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 5214311..7b9780c 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -51,6 +51,9 @@ ap.add_argument('--data', default='tinystories_bpe')     # dataset dir under ep_
 ap.add_argument('--sync_check', type=int, default=500)   # DDP: verify bitwise param sync every N steps (0=off)
 ap.add_argument('--ddp_backend', default='nccl', choices=['nccl', 'gloo'])  # gloo = correctness tests on shared GPUs
 ap.add_argument('--ddp_grad_test', action='store_true')  # one-step grad equivalence test vs single-GPU big batch, then exit
+ap.add_argument('--beta_cap_rho', type=float, default=0.0)  # >0: LOOP-GAIN CAP on beta — if per-sweep residual
+                                                          # ratio rho^ exceeds this, bscale *= 0.8 (beta backs off
+                                                          # under the wall-2 ceiling); recovers x1.02 when rho^ low
 ap.add_argument('--relax_tol', type=float, default=0.0)   # >0: ADAPTIVE relax — sweep until rel. state change < tol
                                                           # (or --kmax), geta backtracks x0.6 on residual GROWTH (rho>=1
                                                           # signal), then one final graphed round. 0 = legacy fixed-K.
@@ -341,10 +344,13 @@ def relax(z0, zs, ins, outs, y, beta, K, x):
         ins, outs = n_ins, n_outs
         return rnum / max(rden, 1e-9)
 
-    if not adaptive:                     # legacy fixed-K path (bit-identical to before)
+    if not adaptive:                     # legacy fixed-K path (bit-identical update semantics)
+        rlist = []
         for k in range(K):
             forces(k % args.dtop_every == 0)
-            rebuild(k + 1 == K)
+            rlist.append(rebuild(k + 1 == K))
+        if len(rlist) >= 2 and rlist[-2] > 1e-12:
+            GOV['rho'] = rlist[-1] / rlist[-2]   # per-sweep contraction ratio = live loop-gain meter
         GOV['kuse'] = K
         return zs, outs
 
@@ -353,6 +359,8 @@ def relax(z0, zs, ins, outs, y, beta, K, x):
         forces(k % args.dtop_every == 0)
         res = rebuild(False)
         k += 1
+        if prev_res is not None and prev_res > 1e-12:
+            GOV['rho'] = res / prev_res
         if prev_res is not None and res > prev_res and res > args.relax_tol:
             geta_l = max(0.2, geta_l * 0.6)   # residual GREW: local rho>=1 -> damp harder
         prev_res = res
@@ -396,6 +404,8 @@ def ep_step(x, y):
     fl = args.beta_floor
     if args.bf_late > 0.0 and GOV.get('step', 0) >= args.bf_late_at: fl = args.bf_late
     if fl > 0.0: beta_t = max(beta_t, fl)
+    beta_t = beta_t * GOV.get('cap', 1.0)   # wall-2 loop-gain cap OVERRIDES the floor (the ceiling
+                                            # can sit below the floor near the wall; survival first)
     if args.bsign_rand and torch.rand((), generator=BGEN).item() < 0.5: beta_t = -beta_t
     z0, zs, ins, outs = free_states_graphed(x)
     zs_free = [z.clone() for z in zs]
@@ -422,6 +432,12 @@ def ep_step(x, y):
             for p in all_params: p.grad = None
             return free_ce, beta_t, GOV.get('kuse', GOV['K']), False
     GOV['drift'] = gdrift
+    if args.beta_cap_rho > 0 and GOV.get('rho') is not None:
+        rho_g = ddp_bcast_scalar(GOV['rho'])   # rank0's meter rules (identical control on all ranks)
+        if rho_g > args.beta_cap_rho:
+            GOV['cap'] = max(GOV.get('cap', 1.0) * 0.8, 0.02)   # back beta off under the wall-2 ceiling
+        elif rho_g < 0.5 * args.beta_cap_rho:
+            GOV['cap'] = min(GOV.get('cap', 1.0) * 1.02, 1.0)   # slow recovery toward schedule
     E = 0.0
     for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach().float() - o.float()) ** 2).sum()   # fp32 accumulation (bf16-safe; no-op in fp32)
     if args.est == 'single':