diff options
Diffstat (limited to 'ep_run/casc_bp_train.py')
| -rw-r--r-- | ep_run/casc_bp_train.py | 11 |
1 files changed, 10 insertions, 1 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 8047091..7ee84e8 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -17,6 +17,8 @@ ap.add_argument('--wandb', default=''); ap.add_argument('--wandb_run', default=' ap.add_argument('--opt', choices=['adamw', 'muon'], 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 std (GPT-standard 0.02) +ap.add_argument('--cosine', action='store_true') # warmup then cosine decay to lr_min_ratio*lr over --steps (long runs) +ap.add_argument('--lr_min_ratio', type=float, default=0.1) args = ap.parse_args() torch.manual_seed(args.seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' @@ -54,7 +56,14 @@ if args.opt == 'muon': opt, sched = build_hybrid(blocks, params, args.lr, args.muon_lr, args.warmup) else: opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=1e-4) - sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(args.warmup, 1))) + if args.cosine: + def _lrlam(s): + if s < args.warmup: return (s + 1) / max(args.warmup, 1) + p = min(1.0, (s - args.warmup) / max(1, args.steps - args.warmup)) + return args.lr_min_ratio + 0.5 * (1 - args.lr_min_ratio) * (1 + math.cos(math.pi * p)) + sched = torch.optim.lr_scheduler.LambdaLR(opt, _lrlam) + else: + sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(args.warmup, 1))) def fwd(x): z = tok(x) + pos(torch.arange(args.T, device=dev))[None] |