bp_lm twin control + env forensics probes + holo_a_track_avg (semi-convergence fix, ungated)

BP+EMA 2x2 (local, healthy env): qknorm {1.9951, 1.9977}, no-qknorm {1.9728,
1.9732} -> parameterization-matched BP twin tops at ~1.97 vs EP 1.7888: the
equilibrium computation's iteration/depth dividend = 0.18 CE from identical
parameters. External anchor (tuned depth-1 BP, 1.7921): EP at parity.
bp_lm on 107/2.3.1 gave 1.9823 ~= local -> plain backward exonerated on the
pascal env; the 107 divergence (pair/floss/resreg dead by step 600) narrows
to the EP reg/estimator loop. Single-step fingerprints all match (a/b/c/d) ->
suspected intermittent kernel issue; discriminators in flight (torch-2.7 env
probe + delay_hr2's step-2000 reg-on transition).

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

4 files changed, 238 insertions, 0 deletions

diff --git a/ep_run/bp_lm.py b/ep_run/bp_lm.py
new file mode 100644
index 0000000..7fe0bdd
--- /dev/null
+++ b/ep_run/bp_lm.py
@@ -0,0 +1,75 @@
+"""bp_lm.py — the plain-BP standard-transformer reference with weight-EMA (task #34), rebuilt for
+metric parity with the EP lineage: SAME parameter tensors as EQBlock (7.48M), SAME data + eval
+protocol, and the SAME pema weight-EMA so both lineages report best-of(raw, ema) val. Forward = one
+standard pre-LN residual block (no relaxation): h = x + attn(LN1 x) + ffn(LN2 x); logits = h @ Wh.
+Historical anchor (no EMA, no qknorm): 'BP standard transformer (7.48M, lr 3e-3) best 1.7921 (20k)'
+(FINDINGS:355). Run with --qknorm for architecture parity with the current EP recipe."""
+import argparse, time, torch, torch.nn.functional as F
+import lt_ep_train as L
+
+
+def fwd(blk, idx):
+    x = blk.embed(idx)
+    h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
+    h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
+    h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
+    return h @ blk.Wh
+
+
+def evaluate(blk, nb=8, B=32):
+    tot = 0.0
+    with torch.no_grad():
+        for _ in range(nb):
+            idx, y = L.get_batch('val', B, blk.T)
+            tot += float(F.cross_entropy(fwd(blk, idx).reshape(-1, L.vocab), y.reshape(-1)))
+    return tot / nb
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--steps', type=int, default=32000)
+    ap.add_argument('--B', type=int, default=24)
+    ap.add_argument('--lr', type=float, default=3e-3)
+    ap.add_argument('--wd', type=float, default=1e-4)
+    ap.add_argument('--pema', type=float, default=0.999)
+    ap.add_argument('--qknorm', action='store_true')
+    ap.add_argument('--log', type=int, default=200)
+    ap.add_argument('--ckpt', type=str, default='runs/bp_lm.pt')
+    cfg = ap.parse_args()
+    torch.manual_seed(0)
+    blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick')
+    blk.qknorm = cfg.qknorm
+    opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05)
+    ema = [p.detach().clone() for p in blk.allp]
+    best, t0 = float('inf'), time.time()
+    for step in range(1, cfg.steps + 1):
+        idx, y = L.get_batch('train', cfg.B, blk.T)
+        loss = F.cross_entropy(fwd(blk, idx).reshape(-1, L.vocab), y.reshape(-1))
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step(); sched.step()
+        with torch.no_grad():
+            for e, p in zip(ema, blk.allp):
+                e.mul_(cfg.pema).add_(p, alpha=1 - cfg.pema)
+        if step % cfg.log == 0:
+            raw = evaluate(blk)
+            bak = [p.detach().clone() for p in blk.allp]
+            with torch.no_grad():
+                for p, e in zip(blk.allp, ema):
+                    p.copy_(e)
+            emav = evaluate(blk)
+            with torch.no_grad():
+                for p, b in zip(blk.allp, bak):
+                    p.copy_(b)
+            m = min(raw, emav)
+            if m < best:
+                best = m
+                torch.save({'allp': [p.detach().cpu() for p in blk.allp], 'step': step, 'best': best}, cfg.ckpt)
+            print(f"step {step:5d}/{cfg.steps} | val {raw:.4f} ema {emav:.4f} (best {best:.4f}) "
+                  f"| {step / (time.time() - t0):.2f} it/s", flush=True)
+    print(f"[bp_lm] DONE best {best:.4f}", flush=True)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/ep_run/grad_env_probe.py b/ep_run/grad_env_probe.py
new file mode 100644
index 0000000..981ee9e
--- /dev/null
+++ b/ep_run/grad_env_probe.py
@@ -0,0 +1,51 @@
+"""Cross-env gradient fingerprint: three gradient paths on the SAME ckpt (s2000) + SAME seeded batch.
+Run under torch 2.10 (local A6000) and torch 2.3.1 (timan107 GTX1080); diverging fingerprints pinpoint
+which path is numerically broken on the pascal env (suspects: double-backward reg paths).
+  (a) plain CE backward through the feedforward block   (bp-style backward)
+  (b) resreg path: tforce graph at z_T1 + autograd.grad (single backward through force graph)
+  (c) jacreg path: autograd.functional.jvp(create_graph) + backward   (double-backward)
+Prints total grad-norm per path + first-3 per-tensor norms."""
+import torch, torch.nn.functional as F
+import lt_ep_train as L
+
+torch.manual_seed(0)
+blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
+ck = torch.load('runs/redx_traj/s2000.pt', map_location=L.dev)
+with torch.no_grad():
+    for p, w in zip(blk.allp, ck['allp']):
+        p.copy_(w.to(L.dev))
+torch.manual_seed(42)
+idx, y = L.get_batch('train', 8, 256)
+xin = blk.embed(idx).detach()
+zs = L.relax(blk, xin.clone(), xin, 150, 0.1)
+
+def report(tag, params, grads):
+    tot = sum(float(g.pow(2).sum()) for g in grads if g is not None) ** 0.5
+    firsts = " ".join(f"{float(g.norm()):.3e}" if g is not None else "None" for g in grads[:3])
+    fin = all(torch.isfinite(g).all() for g in grads if g is not None)
+    print(f"{tag}: total={tot:.4e} finite={fin} first3=[{firsts}]", flush=True)
+
+# (a) plain BP backward
+x = blk.embed(idx)
+h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
+h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
+h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
+loss = F.cross_entropy((h @ blk.Wh).reshape(-1, L.vocab), y.reshape(-1))
+ga = torch.autograd.grad(loss, blk.allp, allow_unused=True)
+report("(a) plain-BP  ", blk.allp, list(ga))
+
+# (b) resreg path
+with torch.enable_grad():
+    Fz = blk.tforce(zs, xin)
+    Rr = (0.1 * Fz).pow(2).sum() / (zs.pow(2).sum() + 1e-9)
+    gb = torch.autograd.grad(Rr, blk.block, allow_unused=True)
+report("(b) resreg    ", blk.block, list(gb))
+
+# (c) jacreg double-backward path
+er = torch.randn_like(zs)
+with torch.enable_grad():
+    Jv = torch.autograd.functional.jvp(blk.nc_force, zs.detach(), er, create_graph=True)[1]
+    R = 0.1 * (Jv ** 2).sum() / (er ** 2).sum()
+    gc = torch.autograd.grad(R, blk.block, allow_unused=True)
+report("(c) jacreg dbb", blk.block, list(gc))
+print("GRAD_ENV_PROBE_DONE", flush=True)
diff --git a/ep_run/grad_init_probe.py b/ep_run/grad_init_probe.py
new file mode 100644
index 0000000..18512c0
--- /dev/null
+++ b/ep_run/grad_init_probe.py
@@ -0,0 +1,56 @@
+"""Cross-env gradient fingerprint: three gradient paths on the SAME ckpt (s2000) + SAME seeded batch.
+Run under torch 2.10 (local A6000) and torch 2.3.1 (timan107 GTX1080); diverging fingerprints pinpoint
+which path is numerically broken on the pascal env (suspects: double-backward reg paths).
+  (a) plain CE backward through the feedforward block   (bp-style backward)
+  (b) resreg path: tforce graph at z_T1 + autograd.grad (single backward through force graph)
+  (c) jacreg path: autograd.functional.jvp(create_graph) + backward   (double-backward)
+Prints total grad-norm per path + first-3 per-tensor norms."""
+import torch, torch.nn.functional as F
+import lt_ep_train as L
+
+torch.manual_seed(0)
+blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
+# RANDOM-INIT probe: no ckpt load (the sick regime is step 0-400)
+torch.manual_seed(42)
+idx, y = L.get_batch('train', 8, 256)
+xin = blk.embed(idx).detach()
+zs = L.relax(blk, xin.clone(), xin, 150, 0.1)  # init operator: res ~e-9 regime
+
+def report(tag, params, grads):
+    tot = sum(float(g.pow(2).sum()) for g in grads if g is not None) ** 0.5
+    firsts = " ".join(f"{float(g.norm()):.3e}" if g is not None else "None" for g in grads[:3])
+    fin = all(torch.isfinite(g).all() for g in grads if g is not None)
+    print(f"{tag}: total={tot:.4e} finite={fin} first3=[{firsts}]", flush=True)
+
+# (a) plain BP backward
+x = blk.embed(idx)
+h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
+h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
+h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
+loss = F.cross_entropy((h @ blk.Wh).reshape(-1, L.vocab), y.reshape(-1))
+ga = torch.autograd.grad(loss, blk.allp, allow_unused=True)
+report("(a) plain-BP  ", blk.allp, list(ga))
+
+# (b) resreg path
+with torch.enable_grad():
+    Fz = blk.tforce(zs, xin)
+    Rr = (0.1 * Fz).pow(2).sum() / (zs.pow(2).sum() + 1e-9)
+    gb = torch.autograd.grad(Rr, blk.block, allow_unused=True)
+report("(b) resreg    ", blk.block, list(gb))
+
+# (c) jacreg double-backward path
+er = torch.randn_like(zs)
+with torch.enable_grad():
+    Jv = torch.autograd.functional.jvp(blk.nc_force, zs.detach(), er, create_graph=True)[1]
+    R = 0.1 * (Jv ** 2).sum() / (er ** 2).sum()
+    gc = torch.autograd.grad(R, blk.block, allow_unused=True)
+report("(c) jacreg dbb", blk.block, list(gc))
+# (d) the FULL ep_step gradient (holo track estimator + AEP + regs) — the actual training signal
+blk.track = True
+ge, res = L.ep_step(blk, idx, y, 150, 20, 0.1, 0.02, jacreg=0.1, holo=2, hr=0.02,
+                    t1max=300, res_est=1e-4, t2sel=40, corr_every=1, res_gate=0.0, resreg=0.2)
+gs = [g for g in ge.values() if g is not None]
+tot = sum(float(g.pow(2).sum()) for g in gs) ** 0.5
+fin = all(torch.isfinite(g).all() for g in gs)
+print(f"(d) full ep_step: total={tot:.4e} finite={fin} res={res:.2e} n={len(gs)}", flush=True)
+print("GRAD_ENV_PROBE_DONE", flush=True)
diff --git a/ep_run/holo_ep.py b/ep_run/holo_ep.py
index 5485bc7..d8d2be5 100644
--- a/ep_run/holo_ep.py
+++ b/ep_run/holo_ep.py
@@ -298,6 +298,62 @@ def holo_a_track_fast(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
     return a_best.detach(), t_best
 
 
+def holo_a_track_avg(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
+    """track_fast + the semi-convergence fix (t2_probe 2026-07-05): the adjoint iteration on a
+    near-marginal operator SEMI-converges — error dips at a batch-dependent optimum then grows, and the
+    plain argmin-of-increment t_best gets fooled by rotating slow modes. Two changes:
+    (1) trend-aware stop: break after the increment rises on 2 consecutive checks past 2x inc_min
+        (instead of the blunt exit_mult=5 single-shot);
+    (2) plateau averaging: return the MEAN of the a_t whose increment <= 1.5x inc_min (the flat bottom
+        of the semi-convergence curve) — averages out the rotating error component around the optimum."""
+    import torch.func as tf
+    B = zs.size(0)
+    Z = torch.cat([zs, zs], 0)
+    X2 = torch.cat([xin, xin], 0)
+    y2 = torch.cat([y, y], 0)
+    sg = torch.cat([torch.full((B, 1, 1), r, device=zs.device), torch.full((B, 1, 1), -r, device=zs.device)], 0)
+    fnc = lambda zz: blk.nc_force(zz)
+    a_prev = None
+    hist = []                                              # (inc, a_t) at each K-checkpoint
+    inc_min, rise = float('inf'), 0
+    zs2a = torch.cat([zs, zs], 0)
+    kappa = getattr(blk, 'nbrake', 0.0)
+    for t in range(1, T2max + 1):
+        with torch.no_grad():
+            zbar = 0.5 * (Z[:B] + Z[B:])
+            f = rforce(blk, Z, X2) - sg * rgrad_ce(blk, Z, y2, denom=y.numel())
+            if kappa > 0:
+                f = f - kappa * (Z - zs2a)
+            v0 = (Z[:B] - zbar).contiguous()
+            _, Jv0 = tf.jvp(fnc, (zbar,), (v0,))
+            JTv0 = tf.vjp(fnc, zbar)[1](v0)[0]
+            corr0 = Jv0 - JTv0
+            Z = Z + eps * (f - torch.cat([corr0, -corr0], 0))
+        if t % K == 0 or t == T2max:
+            a_t = (Z[B:] - Z[:B]) / (2 * r)
+            if not torch.isfinite(a_t).all():
+                break
+            if a_prev is not None:
+                inc = (a_t - a_prev).norm().item()
+                hist.append((inc, a_t))
+                if inc < inc_min:
+                    inc_min, rise = inc, 0
+                elif inc > 2.0 * inc_min and t >= 3 * K:
+                    rise += 1                              # trend-aware: need 2 consecutive rising checks
+                    if rise >= 2:
+                        break
+                else:
+                    rise = 0
+            a_prev = a_t
+    if not hist:
+        return (a_prev if a_prev is not None else (Z[B:] - Z[:B]) / (2 * r)).detach(), T2max
+    flat = [a for inc, a in hist if inc <= 1.5 * inc_min]  # the semi-convergence plateau
+    if not flat:
+        flat = [min(hist, key=lambda p: p[0])[1]]
+    a_avg = torch.stack(flat).mean(0)
+    return a_avg.detach(), len(hist) * K
+
+
 def holo_a_lockin(blk, zs, xin, y, r, P, ncyc, eps):
     """True oscillatory EP / lock-in estimator (Laborieux–Zenke taken literally) — the
     noisy-physics form: ONE trajectory, sinusoidal nudge beta(t)=r·sin(2πt/P), in-phase