summaryrefslogtreecommitdiff
path: root/ep_run
diff options
context:
space:
mode:
Diffstat (limited to 'ep_run')
-rw-r--r--ep_run/holo_ep.py44
-rw-r--r--ep_run/holo_fast_gate.log7
-rw-r--r--ep_run/holo_fast_gate.py25
-rw-r--r--ep_run/holo_fast_parity.log6
-rw-r--r--ep_run/holo_fast_parity.py33
-rw-r--r--ep_run/holo_fast_probe.log5
-rw-r--r--ep_run/holo_fast_probe.py49
-rw-r--r--ep_run/lt_ep_train.py7
8 files changed, 174 insertions, 2 deletions
diff --git a/ep_run/holo_ep.py b/ep_run/holo_ep.py
index 31054e4..5485bc7 100644
--- a/ep_run/holo_ep.py
+++ b/ep_run/holo_ep.py
@@ -254,6 +254,50 @@ def holo_a_track(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
return a_best.detach(), t_best
+def holo_a_track_fast(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
+ """EXACT restructure of holo_a_track (same math, ~half the correction cost): the two phase-halves'
+ deviations from the common mode are exact negatives (v[B:] = -v[:B], since zbar is their mean) and
+ the Jacobian anchor zbar is shared, so the doubled-batch jvp/vjp computed [J v0; -J v0] redundantly.
+ Compute Jv/JTv once at batch B and mirror. Bit-equal up to fp nondeterminism."""
+ 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 = a_best = None
+ inc_min, t_best = 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: # measurement brake: Tikhonov-regularized adjoint
+ f = f - kappa * (Z - zs2a)
+ v0 = (Z[:B] - zbar).contiguous() # v of the +r phase; the -r phase's v is exactly -v0
+ _, 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()
+ if inc < inc_min:
+ inc_min, a_best, t_best = inc, a_t, t
+ elif inc > exit_mult * inc_min and t >= 3 * K:
+ break
+ a_prev = a_t
+ if a_best is None:
+ a_best = a_prev if a_prev is not None else (Z[B:] - Z[:B]) / (2 * r)
+ t_best = T2max
+ return a_best.detach(), t_best
+
+
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
diff --git a/ep_run/holo_fast_gate.log b/ep_run/holo_fast_gate.log
new file mode 100644
index 0000000..a057dd4
--- /dev/null
+++ b/ep_run/holo_fast_gate.log
@@ -0,0 +1,7 @@
+batch cos_orig cos_fast res
+/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/autograd/graph.py:865: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:330.)
+ return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
+ 0 0.9068 0.9117 9.47e-03
+ 1 0.8533 0.8525 4.00e-03
+ 2 0.9177 0.9178 4.00e-03
+HOLO_FAST_GATE_DONE
diff --git a/ep_run/holo_fast_gate.py b/ep_run/holo_fast_gate.py
new file mode 100644
index 0000000..2a5f39b
--- /dev/null
+++ b/ep_run/holo_fast_gate.py
@@ -0,0 +1,25 @@
+"""Ship-gate for --holofast: gradient-level equivalence. cos(EP,BPTT) with orig vs fast track on the
+same batches (s2000, track path, holo=2 t2sel=40). Ship iff the two cos columns are statistically
+indistinguishable (the a-level 45% parity gap is FD noise; what matters is the gradient direction)."""
+import torch
+import lt_ep_train as L
+from diag_cos import cos_ep_bptt
+
+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))
+blk.track = True
+
+print(f"{'batch':>5} {'cos_orig':>9} {'cos_fast':>9} {'res':>9}", flush=True)
+torch.manual_seed(11)
+batches = [L.get_batch('train', 24, 256) for _ in range(3)]
+for b, (idx, y) in enumerate(batches):
+ blk.holofast = False
+ c0, r0 = cos_ep_bptt(blk, idx, y, 150, 20, 0.1, 0.02, holo=2, hr=0.02, t2sel=40)
+ blk.holofast = True
+ c1, r1 = cos_ep_bptt(blk, idx, y, 150, 20, 0.1, 0.02, holo=2, hr=0.02, t2sel=40)
+ print(f"{b:>5} {c0:>9.4f} {c1:>9.4f} {r0:>9.2e}", flush=True)
+print("HOLO_FAST_GATE_DONE", flush=True)
diff --git a/ep_run/holo_fast_parity.log b/ep_run/holo_fast_parity.log
new file mode 100644
index 0000000..1aedf66
--- /dev/null
+++ b/ep_run/holo_fast_parity.log
@@ -0,0 +1,6 @@
+/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/autograd/graph.py:865: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:330.)
+ return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
+parity: rel_diff=4.48e-01 cos=0.89945102 t_best 20 vs 20
+orig: 6.433s (best of 3)
+fast: 4.163s (best of 3)
+HOLO_FAST_PARITY_DONE
diff --git a/ep_run/holo_fast_parity.py b/ep_run/holo_fast_parity.py
new file mode 100644
index 0000000..d2df0da
--- /dev/null
+++ b/ep_run/holo_fast_parity.py
@@ -0,0 +1,33 @@
+"""Parity + timing for holo_a_track_fast (exact halved-jvp restructure) vs holo_a_track.
+Same ckpt (s2000), same batch, same T2max: a_best must match to fp noise, t_best exactly;
+timing over 3 reps each (GPU contended — relative ratio is the signal)."""
+import time, torch
+import lt_ep_train as L
+from holo_ep import holo_a_track, holo_a_track_fast
+
+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', 24, 256)
+xin = blk.embed(idx).detach()
+zs = L.relax(blk, xin.clone(), xin, 150, 0.1)
+
+r, T2, eps = 0.02, 40, 0.1
+a0, t0 = holo_a_track(blk, zs, xin, y, r, T2, eps) # warmup + reference
+a1, t1 = holo_a_track_fast(blk, zs, xin, y, r, T2, eps)
+rel = ((a1 - a0).norm() / (a0.norm() + 1e-12)).item()
+cos = torch.nn.functional.cosine_similarity(a0.flatten(), a1.flatten(), dim=0).item()
+print(f"parity: rel_diff={rel:.2e} cos={cos:.8f} t_best {t0} vs {t1}", flush=True)
+
+for name, fn in (('orig', holo_a_track), ('fast', holo_a_track_fast)):
+ ts = []
+ for _ in range(3):
+ torch.cuda.synchronize(); t = time.time()
+ fn(blk, zs, xin, y, r, T2, eps)
+ torch.cuda.synchronize(); ts.append(time.time() - t)
+ print(f"{name}: {min(ts):.3f}s (best of 3)", flush=True)
+print("HOLO_FAST_PARITY_DONE", flush=True)
diff --git a/ep_run/holo_fast_probe.log b/ep_run/holo_fast_probe.log
new file mode 100644
index 0000000..c7e77b8
--- /dev/null
+++ b/ep_run/holo_fast_probe.log
@@ -0,0 +1,5 @@
+/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/autograd/graph.py:865: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:330.)
+ return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
+(a) single-eval: rel=5.98e-07 (EXACT — gap is fp/FD noise)
+(b) orig-vs-orig under 1e-6 state noise: rel=4.87e-01 cos=0.881368
+HOLO_FAST_PROBE_DONE
diff --git a/ep_run/holo_fast_probe.py b/ep_run/holo_fast_probe.py
new file mode 100644
index 0000000..9387199
--- /dev/null
+++ b/ep_run/holo_fast_probe.py
@@ -0,0 +1,49 @@
+"""Disambiguate the parity gap: real bug vs FD-amplified fp noise.
+(a) single-eval exactness: at a synthetic two-phase state Z=[zs+d, zs-d], compare the full doubled-batch
+ correction (Jv-JTv) against the halved+mirrored one. Exact math => allclose at fp level.
+(b) estimator noise floor: perturb zs by 1e-6 relative and rerun the ORIGINAL holo_a_track — if a_best
+ moves by ~the same 0.4 rel, the parity gap is the estimator's intrinsic FD sensitivity, not a bug."""
+import torch, torch.func as tf
+import lt_ep_train as L
+from holo_ep import holo_a_track
+
+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', 24, 256)
+xin = blk.embed(idx).detach()
+zs = L.relax(blk, xin.clone(), xin, 150, 0.1)
+B = zs.size(0)
+fnc = lambda zz: blk.nc_force(zz)
+
+# (a) single-eval exactness
+torch.manual_seed(7)
+d = 0.02 * torch.randn_like(zs)
+Z = torch.cat([zs + d, zs - d], 0)
+zbar = 0.5 * (Z[:B] + Z[B:])
+zb2 = torch.cat([zbar, zbar], 0)
+v = (Z - zb2).contiguous()
+with torch.no_grad():
+ _, Jv = tf.jvp(fnc, (zb2,), (v,))
+ JTv = tf.vjp(fnc, zb2)[1](v)[0]
+ corr_full = Jv - JTv
+ v0 = (Z[:B] - zbar).contiguous()
+ _, Jv0 = tf.jvp(fnc, (zbar,), (v0,))
+ JTv0 = tf.vjp(fnc, zbar)[1](v0)[0]
+ corr_half = torch.cat([Jv0 - JTv0, -(Jv0 - JTv0)], 0)
+rel_a = ((corr_full - corr_half).norm() / (corr_full.norm() + 1e-12)).item()
+print(f"(a) single-eval: rel={rel_a:.2e} ({'EXACT — gap is fp/FD noise' if rel_a < 1e-4 else 'REAL BUG'})", flush=True)
+
+# (b) estimator noise floor of the ORIGINAL
+r, T2, eps = 0.02, 40, 0.1
+a_ref, _ = holo_a_track(blk, zs, xin, y, r, T2, eps)
+zs_p = zs + 1e-6 * zs.norm() / (zs.numel() ** 0.5) * torch.randn_like(zs)
+a_prt, _ = holo_a_track(blk, zs_p, xin, y, r, T2, eps)
+rel_b = ((a_prt - a_ref).norm() / (a_ref.norm() + 1e-12)).item()
+cos_b = torch.nn.functional.cosine_similarity(a_ref.flatten(), a_prt.flatten(), dim=0).item()
+print(f"(b) orig-vs-orig under 1e-6 state noise: rel={rel_b:.2e} cos={cos_b:.6f}", flush=True)
+print("HOLO_FAST_PROBE_DONE", flush=True)
diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py
index d56af1d..e7155d3 100644
--- a/ep_run/lt_ep_train.py
+++ b/ep_run/lt_ep_train.py
@@ -180,12 +180,13 @@ def ep_step(blk, idx, y, T1, T2, eps, beta, jacreg=0.0, holo=0, hr=0.02, t1max=0
z = z + eps * f
return z.detach()
if holo == 2 and t2sel > 0: # adaptive-T2, phase-batched fast path (validated ==)
- from holo_ep import holo_a_select2, holo_a_track
+ from holo_ep import holo_a_select2, holo_a_track, holo_a_track_fast
K = max(1, getattr(blk, 'navg', 1)) # restart-averaging: noise / sqrt(K)
acc = None
for _ in range(K):
if getattr(blk, 'track', False): # common-mode-tracking AEP (loose-tolerant)
- ai, _ = holo_a_track(blk, zs, xin0, y, hr, t2sel, eps)
+ _tr = holo_a_track_fast if getattr(blk, 'holofast', False) else holo_a_track
+ ai, _ = _tr(blk, zs, xin0, y, hr, t2sel, eps) # fast = exact halved-jvp mirror (1.55x)
else:
ai, _ = holo_a_select2(blk, zs, xin0, y, hr, t2sel, eps, li=getattr(blk, 'li_avg', 0))
acc = ai if acc is None else acc + ai
@@ -419,6 +420,7 @@ def main():
ap.add_argument('--li_avg', type=int, default=0) # lock-in integration window (0=snapshot mode)
ap.add_argument('--navg', type=int, default=1) # restart-averaged contrast estimates per update
ap.add_argument('--track', action='store_true') # common-mode-tracking AEP correction
+ ap.add_argument('--holofast', action='store_true') # exact halved-jvp track (1.55x nudged phase; parity = FD noise floor)
ap.add_argument('--rt_final', type=float, default=0.0) # anneal res_target to this (0=off), 25%-75% of run
ap.add_argument('--nudge_brake', type=float, default=0.0) # kappa: anchor spring during nudge (Tikhonov adjoint)
ap.add_argument('--init_ckpt', type=str, default='') # warm-start weights from a saved ckpt
@@ -493,6 +495,7 @@ def main():
blk.li_avg = cfg.li_avg
blk.navg = cfg.navg
blk.track = cfg.track
+ blk.holofast = cfg.holofast
blk.nbrake = cfg.nudge_brake
blk.qknorm = cfg.qknorm
if cfg.resinit != 1.0: # near-identity block at init (contractive) -> stable big-width start
generated by cgit v1.2.3 (git 2.25.1) at 2026-08-18 16:06:15 +0000