diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 6bd83f8..5a3e640 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -175,7 +175,8 @@ if args.resume: print(f'[resume] loaded {args.resume} at step {start_step}', flush=True) if args.opt == 'muon': from muon import build_hybrid - opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup) + opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup, + total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio) else: if args.wd >= 0: # OLMo2-style grouped decay: linear weights + head decay; embeddings/norm-gains none nodecay = {id(p) for p in tok.parameters()} | {id(p) for p in pos.parameters()} | \ |