BP-free formal audit (5.6-sol): EP update verified local (108 bitwise zero-influence checks); fixes — gate_every<=0 = true off switch, governor reaction now opt-in (--gate_govern, default observe-only), dFdtheta restored as self-sealing invariant-test surface; test_bp_free.py 4/4 green

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

1 files changed, 16 insertions, 8 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 4eba267..f1f26a4 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -27,7 +27,8 @@ ap.add_argument('--tok_init', type=float, default=0.0)  # >0: init tok/pos with
 ap.add_argument('--compile', action='store_true')       # torch.compile each block (free speed where supported)
 ap.add_argument('--sig_every', type=int, default=25)    # tok-sigma refresh interval (amortized)
 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
+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)
 args = ap.parse_args()
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -129,14 +130,20 @@ def relax(z0, zs, ins, outs, y, beta, K, x):
     return zs, outs
 
 def dFdtheta(zs, x, y, beta):
-    """dF/dtheta at fixed relaxed states (z0 rebuilt WITH graph so emb gets its E-path grad)."""
+    """theta-readout at FIXED states. Not used by the training loop (relax reuses its own
+    graphs); kept as the INVARIANT-TEST surface for test_bp_free.py. Self-sealing: inputs
+    are detached here so the local-graph property holds for any caller."""
+    zs = [z.detach() for z in zs]
     prev = tok(x) + pos(torch.arange(args.T, device=dev))[None]
     E = 0.0
-    for z, b in zip(zs, blocks): E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum(); prev = z
+    for z, b in zip(zs, blocks):
+        E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum()
+        prev = z   # zs detached at entry => blocks l>0 get detached inputs; block 0 gets the graphed emb
     obj = E / NBT + beta * F.cross_entropy(readout(zs[-1]).reshape(-1, vocab), y.reshape(-1))
     gs = torch.autograd.grad(obj, all_params, allow_unused=True)
     return [g if g is not None else None for g in gs]
 
+
 SIG0 = None
 GOV = {'K': None, 'bscale': 1.0, 'gema': None, 'drift': 0.0, 'gn': 0.0, 'sig': 0.0}
 def ep_step(x, y):
@@ -215,17 +222,18 @@ for step in range(args.steps + 1):
     ce, beta_t, rounds, ok = ep_step(x, y)
     if not ok: skips += 1
     gcos = float('nan')
-    if step % args.gate_every == 0 and ok:
+    if args.gate_every > 0 and step % args.gate_every == 0 and ok:
         gbp = bp_gate(x, y)
         num = den1 = den2 = 0.0
         for p, g in zip(all_params, gbp):
             if p.grad is None or g is None: continue
             num += float((p.grad * g).sum()); den1 += float((p.grad ** 2).sum()); den2 += float((g ** 2).sum())
         gcos = num / max((den1 ** 0.5) * (den2 ** 0.5), 1e-12)
-        if gcos < 0.97:                                   # estimator governor: spend more
-            GOV['K'] = min(GOV['K'] + 2, args.kmax); GOV['bscale'] = max(GOV['bscale'] * 0.7, 0.05)
-        elif gcos > 0.995 and GOV['K'] > args.K:          # relax back when quality is abundant
-            GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0)
+        if args.gate_govern:                              # opt-in: BP-informed control flow
+            if gcos < 0.97:
+                GOV['K'] = min(GOV['K'] + 2, args.kmax); GOV['bscale'] = max(GOV['bscale'] * 0.7, 0.05)
+            elif gcos > 0.995 and GOV['K'] > args.K:
+                GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0)
     torch.nn.utils.clip_grad_norm_(all_params, 1.0)
     opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
     if step % args.log == 0: