From 6604a012b96ea9b18d95fabf152072c8728df31b Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 6 Jul 2026 21:15:45 -0500 Subject: dp_ep smoke PASSED: 2-GPU lockstep clean, ~107% linear scaling (allreduce negligible vs seconds-long EP steps) val 6.0->5.83/100 steps at effective B48, res nominal, ranks never diverged. Speed package final ledger: pack 1.47x X DP ~linear (X optional AA-v2 1.3x). dip_screen.py: ARPACK screener for the dip-farm trajectories (running on freed 1080s, 2 chains). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/dip_screen.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 ep_run/dip_screen.py (limited to 'ep_run/dip_screen.py') diff --git a/ep_run/dip_screen.py b/ep_run/dip_screen.py new file mode 100644 index 0000000..5af4c14 --- /dev/null +++ b/ep_run/dip_screen.py @@ -0,0 +1,47 @@ +"""ARPACK dip-screening of a dip-farm trajectory (the governor's dataset): for every frozen snapshot +>= --from_step, top-3 leading eigenvalues of the forward map at the 400-deep state. Output: one line +per snapshot -> dips (all |lam|<1) marked DIP. Usage: dip_screen.py --seed N [--from_step 800]""" +import argparse, numpy as np, torch, scipy.sparse.linalg as sla +from pathlib import Path +from torch.autograd.functional import jvp +import lt_ep_train as L + +ap = argparse.ArgumentParser() +ap.add_argument('--seed', type=int, required=True) +ap.add_argument('--from_step', type=int, default=800) +cfg = ap.parse_args() +EPS, B, C = 0.1, 6, 1.0 +traj = Path(f'runs/dipfarm_s{cfg.seed}_traj') +steps = sorted(int(p.stem[1:]) for p in traj.glob('s*.pt') if int(p.stem[1:]) >= cfg.from_step) +print(f"[seed {cfg.seed}] screening {len(steps)} snapshots: {steps}", flush=True) + +for s in steps: + torch.manual_seed(0) + blk = L.EQBlock(512, 16, 256, 256, c=C, attn_mode='thick'); blk.qknorm = True + ck = torch.load(traj / f's{s}.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, _ = L.get_batch('train', B, 256) + xin = blk.embed(idx).detach() + z = L.relax(blk, xin.clone(), xin, 400, EPS) + sh, n = z.shape, z.numel() + kk = 1.0 - EPS * (1.0 + C) + + def mv(x, z=z, sh=sh): + v = torch.from_numpy(np.asarray(x, dtype=np.float32)).to(L.dev).view(sh) + with torch.no_grad(): + Mv = kk * v + EPS * jvp(blk.nc_force, z, v)[1] + return Mv.reshape(-1).double().cpu().numpy() + + A = sla.LinearOperator((n, n), matvec=mv, dtype=np.float64) + try: + vals = sorted(sla.eigs(A, k=3, which='LM', return_eigenvectors=False, maxiter=1500, tol=2e-4), + key=lambda x: -abs(x)) + stable = all(abs(l) < 1.0 for l in vals) + out = " ".join(f"|l|={abs(l):.5f}" for l in vals) + print(f"s{cfg.seed}/{s:<5} {out} {'DIP' if stable else ''}", flush=True) + except Exception as e: + print(f"s{cfg.seed}/{s:<5} ARPACK-fail {type(e).__name__}", flush=True) +print(f"DIP_SCREEN_S{cfg.seed}_DONE", flush=True) -- cgit v1.2.3