summaryrefslogtreecommitdiff
path: root/ep_run/casc_eq_train.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-08-05 09:09:14 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-08-05 09:09:14 -0500
commit807c6ba3bab8d9afe18ca7d8fb67038bc2cfaf73 (patch)
tree39122335f5c283685838cf3780718d2d85825cb1 /ep_run/casc_eq_train.py
parent68ac1f9d1fa8944a216964479e49e813de6370ac (diff)
RESULT 90: 优化器筛选收敛 — Muon 4.7227 领跑, NS消融判决(摘NS代价0.45, nanoGPT版未调用NS之谜归他们), sign家族聚在5.13-5.25; k价值曲线在跑; 曲率通道死刑+偏置公理入档
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/casc_eq_train.py')
-rw-r--r--ep_run/casc_eq_train.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index e29048c..d19504d 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', 'cautlion'], default='adamw')
+ap.add_argument('--opt', choices=['adamw', 'muon', 'sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion', 'olionns', 'olionk1', 'olionk2', 'olionk3'], 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)
@@ -345,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', 'cautlion'):
+if args.opt in ('sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion', 'cautlion', 'olionns', 'olionk1', 'olionk2', 'olionk3'):
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,