diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 9 |
1 files changed, 6 insertions, 3 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index db6d71a..a57e7b2 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -139,6 +139,8 @@ ap.add_argument('--relax_tol', type=float, default=0.0) # >0: ADAPTIVE relax â # signal), then one final graphed round. 0 = legacy fixed-K. ap.add_argument('--muon_mom', type=float, default=0.95) # Muon momentum (late-SNR arm: 0.99 = ~10x noise averaging) ap.add_argument('--adam_b1', type=float, default=0.9) # AdamW beta1 (late-SNR arm companion) +ap.add_argument('--adam_b2', type=float, default=0.999) # AdamW beta2 (LLM practice often 0.95) +ap.add_argument('--muon_wd', type=float, default=0.0) # decoupled wd on Muon matrices (sigma/drift-peak fix arm) ap.add_argument('--est', choices=['single', 'centered', 'richardson'], default='single') # centered: [g(+b)+g(-b)]/2 (O(b^2) bias, 2x relax cost) # richardson: 2g(b)-g(2b) (O(b^2) bias, large-b friendly) @@ -356,16 +358,17 @@ elif args.opt == 'muon': muon_mom=args.muon_mom, adam_b1=args.adam_b1, total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio, head_param=(W_out if args.untie or args.olmo2 else None), - head_lr_mult=args.head_lr_mult) + head_lr_mult=args.head_lr_mult, muon_wd=args.muon_wd) 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()} | \ {id(p) for p in blocks.parameters() if p.ndim < 2} | {id(p) for p in ln_f.parameters()} opt = torch.optim.AdamW([ {'params': [p for p in all_params if id(p) not in nodecay], 'weight_decay': args.wd}, - {'params': [p for p in all_params if id(p) in nodecay], 'weight_decay': 0.0}], lr=args.lr) + {'params': [p for p in all_params if id(p) in nodecay], 'weight_decay': 0.0}], lr=args.lr, + betas=(args.adam_b1, args.adam_b2)) else: - opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4) + opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4, betas=(args.adam_b1, args.adam_b2)) if args.cosine: def _lrlam(s): if s < args.warmup: return (s + 1) / max(args.warmup, 1) |