diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_eq_train.py | 4 | ||||
| -rw-r--r-- | ep_run/muon.py | 6 |
2 files changed, 7 insertions, 3 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, diff --git a/ep_run/muon.py b/ep_run/muon.py index 1289620..7bad679 100644 --- a/ep_run/muon.py +++ b/ep_run/muon.py @@ -177,7 +177,11 @@ def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ 'signline': lambda: SignLine(mats, lr=lm, wd=wd), 'conslion': lambda: ConsensusLion(mats, lr=lm, wd=wd), 'ditherlion': lambda: DitherLion(mats, lr=lm, wd=wd), - 'cautlion': lambda: CautiousLion(mats, lr=lm, wd=wd)}[opt_name]() + 'cautlion': lambda: CautiousLion(mats, lr=lm, wd=wd), + 'olionns': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=0), + 'olionk1': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=1), + 'olionk2': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=2), + 'olionk3': lambda: OLion(mats, lr=lm, wd=wd, ns_steps=3)}[opt_name]() oa = torch.optim.AdamW(rest, lr=lr, weight_decay=1e-4) opts = [om, oa] if total_steps > 0: |