From d5158e64cad8a008794b0e61270c3b18821b41b7 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Wed, 5 Aug 2026 06:40:24 -0500 Subject: =?UTF-8?q?EqOLion=20R4(=E7=9C=9F=E5=AE=9E=E8=BD=A8=E8=BF=B9?= =?UTF-8?q?=E5=88=A4=E5=86=B3):=20=E5=86=BB=E7=BB=93=E6=9E=81=E5=9B=A0?= =?UTF-8?q?=E5=AD=90=E5=87=A0=E5=8D=81=E6=AD=A5=E9=80=80=E7=9B=B8=E5=85=B3?= =?UTF-8?q?=E8=87=B3~0.53=3D=E6=85=A2=E6=BC=82=E7=A7=BB=E5=89=8D=E6=8F=90?= =?UTF-8?q?=E5=AF=B9=E7=AC=A6=E5=8F=B7=E6=A8=A1=E5=BC=8F=E4=B8=8D=E6=88=90?= =?UTF-8?q?=E7=AB=8B(=E6=9E=81=E5=9D=87=E8=A1=A1=20=E6=94=BE=E5=A4=A7?= =?UTF-8?q?=E7=97=85=E6=80=81=E5=B0=BE=E9=83=A8,=E5=B0=8F=CF=83=E5=A5=87?= =?UTF-8?q?=E5=BC=82=E5=90=91=E9=87=8F=E9=AB=98=E9=80=9F=E6=97=8B=E8=BD=AC?= =?UTF-8?q?);=20=E5=AD=90=E6=BB=A1=E7=A7=A9=E6=8E=A2=E9=92=88NaN,=20?= =?UTF-8?q?=E6=BB=A1=E7=A7=A9=E5=8D=95=E8=BF=AD=E4=BB=A3=E4=BB=85=E5=B9=B3?= =?UTF-8?q?=E5=9F=BA=E7=BA=BF;=20=E5=AD=98=E6=B4=BB=E7=A9=BA=E9=97=B4?= =?UTF-8?q?=E6=94=B6=E7=AA=84=E5=88=B0=E6=B8=A9=E5=90=AF=E5=8A=A8=20?= =?UTF-8?q?=E5=85=A8NS(3-5=C3=97)=E6=88=96Occam=E8=87=82(=E8=A1=8C/?= =?UTF-8?q?=E5=88=97AGC)=E5=8D=87=E4=B8=BB=E5=80=99=E9=80=89;=20=E9=99=90?= =?UTF-8?q?=E5=AE=9A=3D=E6=9C=80=E6=97=A9=E6=9C=9F=E8=BD=A8=E8=BF=B9?= =?UTF-8?q?=E4=B8=BA=E6=9C=80=E9=9A=BE=E5=B7=A5=E5=86=B5,=20=E6=99=9A?= =?UTF-8?q?=E6=9C=9F=E9=87=8Ddump=E5=90=8E=E5=86=8D=E5=88=A4=E3=80=82=20?= =?UTF-8?q?=E9=99=84:=20--dump=5Fmom=E6=8E=A2=E9=92=88=E5=AE=9E=E8=A3=85(6?= =?UTF-8?q?01=E6=AD=A5=C3=973=E7=9F=A9=E9=98=B5/276MB)=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 | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) (limited to 'ep_run') diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 7ded9d9..4a436d1 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -61,6 +61,9 @@ 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_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) ap.add_argument('--read_lin', action='store_true') # linear-form theta-read: cotangent = the stored d tensor (full # precision) instead of (z - o) (an fp32-ROUNDED copy of d); # algebraically identical via the read identity z - o = d @@ -1048,6 +1051,20 @@ for step in range(start_step, args.steps + 1): WSNAP['p'] = [p.detach().clone() for p in all_params] WSNAP['o'] = _clone_state(opt.state_dict()) 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 + # representative block matrices (low / mid / top), CPU-cloned after each step. + om = opt.optimizers[0] + _mats = om.param_groups[0]['params'] + _sel = [_mats[1], _mats[len(_mats) // 2], _mats[-1]] + if 'MDUMP' not in globals(): + MDUMP = {'shapes': [tuple(p.shape) for p in _sel], 'traj': []} + MDUMP['traj'].append([om.state[p]['mom'].detach().float().cpu().clone() + for p in _sel if p in om.state and 'mom' in om.state[p]]) + if step >= args.dump_mom: + torch.save(MDUMP, f'runs/momdump_{args.tag}.pt') + print(f'[momdump] DONE {len(MDUMP["traj"])} steps -> runs/momdump_{args.tag}.pt', flush=True) + import sys; sys.exit(0) if args.qup_bits > 0: # STAGE-0 HW GATE: finite conductance levels. Snap every weight to an ABSOLUTE # per-tensor grid (range/2^bits) with stochastic rounding (unbiased) — emulates -- cgit v1.2.3