Add --cosine (warmup->cosine decay to 0.1x lr) to both cascade trainers for full-epoch runs

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]