diff options
Diffstat (limited to 'ep_run/casc_bp_train.py')
| -rw-r--r-- | ep_run/casc_bp_train.py | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 987277f..d45327b 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -146,7 +146,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, params, args.lr, args.muon_lr, args.warmup) + opt, sched = build_hybrid(blocks, 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 nodecay = {id(p) for p in tok.parameters()} | {id(p) for p in pos.parameters()} | \ |