Add linear warmup before cosine LR decay

Replaces CosineAnnealingLR with a LambdaLR that linearly ramps the LR
from lr/warmup_epochs to lr over the first warmup_epochs steps, then
applies cosine decay for the remainder. Default warmup_epochs=5;
overridable via --warmup-epochs CLI flag.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

scripts/train.py

@@ -17,6 +17,7 @@ def main() -> None:
     parser.add_argument("--epochs", type=int)
     parser.add_argument("--batch-size", type=int)
     parser.add_argument("--lr", type=float)
+    parser.add_argument("--warmup-epochs", type=int, dest="warmup_epochs")
     parser.add_argument("--hidden-dim", type=int)
     parser.add_argument("--n-blocks", type=int)
     parser.add_argument("--emb-dim", type=int)
@@ -57,6 +58,7 @@ def main() -> None:
             "epochs": args.epochs,
             "batch_size": args.batch_size,
             "lr": args.lr,
+            "warmup_epochs": args.warmup_epochs,
             "val_fraction": args.val_fraction,
             "num_workers": args.num_workers,
             "seed": args.seed,