1 files changed, 56 insertions, 0 deletions
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)