summaryrefslogtreecommitdiff
path: root/ep_run
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-08-05 06:40:24 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-08-05 06:40:24 -0500
commitd5158e64cad8a008794b0e61270c3b18821b41b7 (patch)
treed0f63dcd1897eb8915dbf2fef062ed46c015d23b /ep_run
parent044973e2a78e0faa1296ac370cefbfc2f6feadd6 (diff)
EqOLion R4(真实轨迹判决): 冻结极因子几十步退相关至~0.53=慢漂移前提对符号模式不成立(极均衡
放大病态尾部,小σ奇异向量高速旋转); 子满秩探针NaN, 满秩单迭代仅平基线; 存活空间收窄到温启动 全NS(3-5×)或Occam臂(行/列AGC)升主候选; 限定=最早期轨迹为最难工况, 晚期重dump后再判。 附: --dump_mom探针实装(601步×3矩阵/276MB)。 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run')
-rw-r--r--ep_run/casc_eq_train.py17
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