From 68ac1f9d1fa8944a216964479e49e813de6370ac Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Wed, 5 Aug 2026 08:18:11 -0500 Subject: =?UTF-8?q?=E5=81=8F=E7=BD=AE=E6=94=BE=E5=A4=A7=E5=85=AC=E7=90=86(?= =?UTF-8?q?=E7=94=A8=E6=88=B708-06=E6=8C=87=E5=87=BA):=20=E9=80=90?= =?UTF-8?q?=E5=9D=90=E6=A0=87=E5=BD=92=E4=B8=80=E5=8C=96=E6=8A=8A=E6=8C=81?= =?UTF-8?q?=E4=B9=85=E5=B0=8F=E5=81=8F=E7=BD=AE=E6=94=BE=E5=A4=A7=E6=88=90?= =?UTF-8?q?=E6=BB=A1=E5=B9=85=E6=BC=82=E7=A7=BB,=20sign=E5=85=A8=E5=AE=B6?= =?UTF-8?q?=E5=90=8C=E7=BD=AA(=E5=90=A6=E5=86=B3=E5=8F=AA=E6=8C=A1=20?= =?UTF-8?q?=E6=96=B9=E5=B7=AE=E4=B8=8D=E6=8C=A1=E5=81=8F=E7=BD=AE);=20?= =?UTF-8?q?=E6=9B=B2=E7=8E=87=E9=80=9A=E9=81=93=3D=E5=81=8F=E7=BD=AE?= =?UTF-8?q?=E7=A3=81=E9=93=81(E[s=C2=B7=C4=9D]=3D=CE=B2c+=CE=B4/=CE=B2,=20?= =?UTF-8?q?=E6=B1=A1=E6=9F=93=E8=A2=AB1/=CE=B2=E6=94=BE=E5=A4=A7333=C3=97)?= =?UTF-8?q?,=20=E5=AD=98=E6=B4=BB=3D=E4=B8=A4=E6=A1=A3|=CE=B2|=E9=94=81?= =?UTF-8?q?=E7=9B=B8=E5=88=86=E7=A6=BB;=20=E8=AF=84=E5=88=86=E8=A1=A8?= =?UTF-8?q?=E5=A2=9E=E7=AC=AC=E4=B8=89=E8=BD=B4=3D=E5=81=8F=E7=BD=AE?= =?UTF-8?q?=E6=94=BE=E5=A4=A7=E7=B3=BB=E6=95=B0,=20=E5=99=A8=E4=BB=B6?= =?UTF-8?q?=E7=94=B5=E6=B1=A0=E9=A2=84=E6=B3=A8=E5=86=8C=E6=8E=92=E5=BA=8F?= =?UTF-8?q?=E9=87=8D=E6=8E=92=E9=A2=84=E6=B5=8B(per-matrix=E4=BC=98?= =?UTF-8?q?=E9=9B=85=E9=80=80=E5=8C=96/sign=E6=BB=A1=E5=B9=85/Dither?= =?UTF-8?q?=E7=BF=BB=E8=BA=AB)=E3=80=82=20=E9=99=84:=20--dump=5Fgrad?= =?UTF-8?q?=E6=8E=A2=E9=92=88+cautlion=E5=B7=B2=E5=AE=9E=E8=A3=85=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_eq_train.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) (limited to 'ep_run/casc_eq_train.py') diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 6a6865d..e29048c 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -22,7 +22,7 @@ ap.add_argument('--wandb_run', default='') ap.add_argument('--kmax', type=int, default=8) # adaptive fb rounds cap ap.add_argument('--noguard', action='store_true') # diagnosis: skip only non-finite grads ap.add_argument('--untie', action='store_true') # separate readout matrix (untied from tok) -ap.add_argument('--opt', choices=['adamw', 'muon', 'sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion'], default='adamw') +ap.add_argument('--opt', choices=['adamw', 'muon', 'sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion'], default='adamw') ap.add_argument('--muon_lr', type=float, default=0.02) ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos with this std (GPT-standard 0.02) ap.add_argument('--compile', action='store_true') # torch.compile each block (free speed where supported) @@ -61,6 +61,7 @@ ap.add_argument('--gen_new', type=int, default=120) ap.add_argument('--probe_dgspec', type=int, default=0) # >0: M1 spectroscopy, value = n batches; exits before training ap.add_argument('--probe_gains', default='1,2,4,8,16,32,64,128,256') ap.add_argument('--probe_f64', action='store_true') # fp64 states+model in the probe: the fp-floor decisive arm # >1: per-STEP log-uniform dgain_top in +ap.add_argument('--dump_grad', type=int, default=0) # >0: record per-step grad estimates + beta sign ap.add_argument('--dump_mom', type=int, default=0) # >0: record 3 Muon momentum matrices every step # for N steps -> runs/momdump_{tag}.pt, then exit # (EqOLion tracking-falsification input) @@ -344,7 +345,7 @@ if args.bf16: if args.untie: with torch.no_grad(): W_out.data = W_out.data.to(torch.bfloat16) print('[bf16] model cast to bfloat16 (E-accum + sigma stay fp32)', flush=True) -if args.opt in ('sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion'): +if args.opt in ('sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion'): from muon import build_alt opt, sched = build_alt(args.opt, blocks, all_params, args.lr, args.warmup, total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio, @@ -1050,6 +1051,20 @@ for step in range(start_step, args.steps + 1): with torch.no_grad(): WSNAP['p'] = [p.detach().clone() for p in all_params] WSNAP['o'] = _clone_state(opt.state_dict()) + if args.dump_grad > 0 and args.opt == 'muon': + # per-step gradient ESTIMATES + the bsign coin, for the odd/even (curvature-channel) + # decomposition test: E[s*ghat] = beta*c. Captured BEFORE opt.step consumes grads. + om_ = opt.optimizers[0] + _m2 = om_.param_groups[0]['params'] + _s2 = [_m2[1], _m2[len(_m2) // 2], _m2[-1]] + if 'GDUMP' not in globals(): + GDUMP = {'shapes': [tuple(p.shape) for p in _s2], 'sign': [], 'traj': []} + GDUMP['sign'].append(1.0 if beta_t > 0 else -1.0) + GDUMP['traj'].append([p.grad.detach().float().cpu().clone() for p in _s2]) + if step >= args.dump_grad: + torch.save(GDUMP, f'runs/graddump_{args.tag}.pt') + print(f'[graddump] DONE {len(GDUMP["traj"])} steps -> runs/graddump_{args.tag}.pt', flush=True) + import sys; sys.exit(0) opt.step(); sched.step(); opt.zero_grad(set_to_none=True) if args.dump_mom > 0 and args.opt == 'muon': # per-step momentum trajectory for the EqOLion tracking falsification: three -- cgit v1.2.3