diff options
| -rw-r--r-- | ep_run/bp_lm.py | 17 | ||||
| -rw-r--r-- | ep_run/dipfarm_freezer.py | 38 | ||||
| -rw-r--r-- | ep_run/lt_ep_train.py | 7 | ||||
| -rw-r--r-- | ep_run/newstop_gate.log | 5 | ||||
| -rw-r--r-- | ep_run/newstop_gate.py | 32 | ||||
| -rwxr-xr-x | ep_run/runs/bp_sweep.sh | 8 | ||||
| -rw-r--r-- | ep_run/runs/tol_sweep.sh | 24 |
7 files changed, 128 insertions, 3 deletions
diff --git a/ep_run/bp_lm.py b/ep_run/bp_lm.py index 7fe0bdd..c01cd4f 100644 --- a/ep_run/bp_lm.py +++ b/ep_run/bp_lm.py @@ -33,14 +33,27 @@ def main(): 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('--stdinit', action='store_true') # standard transformer init (EQBlock's is tuned for relaxation) + ap.add_argument('--beta2', type=float, default=0.999) + ap.add_argument('--sched', choices=['cos', 'const'], default='cos') 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) + if cfg.stdinit: # GPT-style: N(0,0.02), scaled residual projections + with torch.no_grad(): + for W in (blk.WQ, blk.WK, blk.WV, blk.fc, blk.Wh, blk.tok): + W.normal_(0, 0.02) + for W in (blk.WO, blk.pj): + W.normal_(0, 0.02 / (2 ** 0.5)) + blk.pos.normal_(0, 0.01) + opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd, betas=(0.9, cfg.beta2)) + if cfg.sched == 'cos': + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05) + else: + sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1.0) ema = [p.detach().clone() for p in blk.allp] best, t0 = float('inf'), time.time() for step in range(1, cfg.steps + 1): diff --git a/ep_run/dipfarm_freezer.py b/ep_run/dipfarm_freezer.py new file mode 100644 index 0000000..49cdead --- /dev/null +++ b/ep_run/dipfarm_freezer.py @@ -0,0 +1,38 @@ +"""Consolidated trajectory freezer for the 4-seed dip farm (runs anywhere on the shared NFS): +copies runs/dipfarm_sN.pt -> runs/dipfarm_sN_traj/s<step>.pt whenever a new step line appears. +Feeds the ARPACK dip-screening (stability-dip statistics: rate, width, screenable yield).""" +import time, os, re, shutil + +os.chdir("/home/yurenh2/ept/ep_run") +SEEDS = (1, 2, 3, 4) +seen = {s: set() for s in SEEDS} +for s in SEEDS: + os.makedirs(f"runs/dipfarm_s{s}_traj", exist_ok=True) +t0 = time.time() +while time.time() - t0 < 24 * 3600: + time.sleep(20) + done = 0 + for s in SEEDS: + log, ck = f"runs/dipfarm_s{s}.log", f"runs/dipfarm_s{s}.pt" + try: + ls = [l for l in open(log) if l.startswith("step")] + except Exception: + continue + if not ls: + continue + m = re.search(r"step +(\d+)/", ls[-1]) + if not m: + continue + step = int(m.group(1)) + if step not in seen[s] and os.path.exists(ck) and os.path.getsize(ck) > 1e6: + try: + shutil.copy2(ck, f"runs/dipfarm_s{s}_traj/s{step}.pt") + seen[s].add(step) + print(f"froze s{s}/{step}", flush=True) + except Exception: + pass + if step >= 4000 or "DONE" in ls[-1]: + done += 1 + if done == len(SEEDS): + break +print("dipfarm freezer done:", {s: len(v) for s, v in seen.items()}, flush=True) diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py index 1307702..2807be3 100644 --- a/ep_run/lt_ep_train.py +++ b/ep_run/lt_ep_train.py @@ -192,7 +192,10 @@ def ep_step(blk, idx, y, T1, T2, eps, beta, jacreg=0.0, holo=0, hr=0.02, t1max=0 acc = None for _ in range(K): if getattr(blk, 'track', False): # common-mode-tracking AEP (loose-tolerant) - _tr = holo_a_track_fast if getattr(blk, 'holofast', False) else holo_a_track + if getattr(blk, 'holoavg', False): # trend-aware stop + plateau averaging (semi-convergence fix) + from holo_ep import holo_a_track_avg as _tr + else: + _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)) @@ -429,6 +432,7 @@ def main(): 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('--sdpa', action='store_true') # fused flash attention in the no_grad relax loop + ap.add_argument('--holoavg', action='store_true') # trend-aware stop + plateau-avg track (gate: 0.913->0.936 @t2sel160) 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 @@ -505,6 +509,7 @@ def main(): blk.track = cfg.track blk.holofast = cfg.holofast blk.sdpa = cfg.sdpa + blk.holoavg = cfg.holoavg blk.nbrake = cfg.nudge_brake blk.qknorm = cfg.qknorm if cfg.resinit != 1.0: # near-identity block at init (contractive) -> stable big-width start diff --git a/ep_run/newstop_gate.log b/ep_run/newstop_gate.log new file mode 100644 index 0000000..22b46c4 --- /dev/null +++ b/ep_run/newstop_gate.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 + argmin(fast): cos=0.9667 0.8773 0.8941 mean=0.9127 + trend+avg: cos=0.9726 0.8893 0.9448 mean=0.9356 +NEWSTOP_GATE_DONE diff --git a/ep_run/newstop_gate.py b/ep_run/newstop_gate.py new file mode 100644 index 0000000..dc01e17 --- /dev/null +++ b/ep_run/newstop_gate.py @@ -0,0 +1,32 @@ +"""Ship-gate for holo_a_track_avg (trend-aware stop + plateau averaging, the semi-convergence fix): +cos(EP,BPTT) at t2sel=160 on s2000 — the regime where argmin-t_best got fooled (t2_probe: batch2 +degraded 0.956->0.894 at 160). PASS = avg recovers the t2sel-80 level (~0.93+) at 160 without +hurting the healthy batches.""" +import torch +import lt_ep_train as L +from diag_cos import cos_ep_bptt +import holo_ep + +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 +blk.holofast = True +torch.manual_seed(11) +batches = [L.get_batch('train', 24, 256) for _ in range(3)] + +orig_fast = holo_ep.holo_a_track_fast +for name, fn in (('argmin(fast)', holo_ep.holo_a_track_fast), ('trend+avg', holo_ep.holo_a_track_avg)): + holo_ep.holo_a_track_fast = fn # ep_step routes track via holofast -> this symbol + import lt_ep_train + lt_ep_train.ep_step.__globals__ # (route happens at import inside ep_step) + cs = [] + for idx, y in batches: + c, _ = cos_ep_bptt(blk, idx, y, 150, 20, 0.1, 0.02, holo=2, hr=0.02, t2sel=160, bsub=4) + cs.append(c) + print(f"{name:>13}: cos={' '.join(f'{c:.4f}' for c in cs)} mean={sum(cs)/len(cs):.4f}", flush=True) +holo_ep.holo_a_track_fast = orig_fast +print("NEWSTOP_GATE_DONE", flush=True) diff --git a/ep_run/runs/bp_sweep.sh b/ep_run/runs/bp_sweep.sh new file mode 100755 index 0000000..e98c79c --- /dev/null +++ b/ep_run/runs/bp_sweep.sh @@ -0,0 +1,8 @@ +#!/bin/bash +cd /home/yurenh2/ept/ep_run +for cfg in "qk3 --qknorm --lr 3e-3" "qk1 --qknorm --lr 1e-3" "nq3 --lr 3e-3" "nq1 --lr 1e-3"; do + set -- $cfg; name=$1; shift + CUDA_VISIBLE_DEVICES=0 python3 bp_lm.py "$@" --steps 32000 --log 1000 --ckpt runs/bp_$name.pt > runs/bp_$name.log 2>&1 + grep "DONE" runs/bp_$name.log +done +echo BP_SWEEP_ALL_DONE diff --git a/ep_run/runs/tol_sweep.sh b/ep_run/runs/tol_sweep.sh new file mode 100644 index 0000000..0536f59 --- /dev/null +++ b/ep_run/runs/tol_sweep.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# TOLERANCE SWEEP (hardware gate, task #15/#16): warm-start s2000, 1500 steps per config. +# GATED on the hr2 verdict (EP+regs from a TRAINED state must be Pascal-safe). 2 GPUs x 5 configs. +# Grid: device noise fnoise x weight-quantization wq_bits (core cells; wmis later). +cd /home/yurenh2/ept/ep_run +PY=/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3 +BASE="--mode ep --attn_mode thick --B 24 --C 512 --H 16 --T 256 --c 1.0 --jacreg 0.1 --jr_floor 0.1 --jr_max 0.1 --resreg 0.2 --holo 2 --hr 0.02 --t2sel 40 --track --pema 0.999 --t1max 300 --res_est 1e-4 --res_gate 0 --qknorm --init_ckpt runs/redx_traj/s2000.pt --warmup 50 --T1 150 --T2 20 --lr 6e-4 --wsd 0.25 --steps 1500 --log 100 --save_every 1500 --abort_res 0.3 --data data/tinystories_bpe" +run() { local g=$1 n=$2; shift 2; env CUDA_VISIBLE_DEVICES=$g $PY lt_ep_train.py $BASE "$@" --ckpt runs/tol_$n.pt --state runs/tol_$n.state > runs/tol_$n.log 2>&1; grep -E "DONE|abort" runs/tol_$n.log | tail -1; } +( + run 6 base + run 6 fn1e3 --fnoise 1e-3 + run 6 fn3e3 --fnoise 3e-3 + run 6 fn1e2 --fnoise 1e-2 + run 6 fn3e2 --fnoise 3e-2 +) & +( + run 7 wq8 --wq_bits 8 + run 7 wq6 --wq_bits 6 + run 7 wq4 --wq_bits 4 + run 7 fn3e3wq8 --fnoise 3e-3 --wq_bits 8 + run 7 wm1e2 --wmis 1e-2 +) & +wait +echo TOL_SWEEP_ALL_DONE |