diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 17 |
1 files changed, 17 insertions, 0 deletions
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 |