Cosine LR with Linear Warmup

A cosine-with-warmup schedule is a learning rate that climbs linearly from zero to a peak, then decays along the upper half of a cosine curve to a floor. It has three regions. From step zero to step warmup_steps the rate scales linearly from zero to the configured peak lr_max. From warmup_steps to total_steps it follows the cosine, decaying from lr_max to lr_min. After total_steps it stays pinned at lr_min, so a misconfigured trainer that overshoots does not silently walk off the end of the schedule.

The warmup exists because the first thousand training updates are the loudest. The weights are still close to initialization, the optimizer's running second-moment estimate has not stabilised, and the gradient norm is large and noisy. Run those updates at peak learning rate and the model either diverges outright or settles into a loss plateau it never escapes. Gradient clipping is the other well-known fix, and Phase 19 lesson 45 covers it. The schedule is the one that starts small and ramps up.

The hard part of building this is that schedules are easy to get wrong by one step. The error surfaces six hours into a run, as a learning rate 1 percent too high or too low just as the model starts overfitting. You will not see it unless you test the schedule at its boundaries.

Warmup formula

For step in [0, warmup_steps] with warmup_steps > 0, the learning rate is lr_max * step / warmup_steps. The degenerate case warmup_steps = 0 means no warmup: the schedule starts at lr_max on step zero and goes straight into cosine decay. Some test harnesses pass zero to confirm the schedule still produces a usable curve.

Cosine formula

For step in (warmup_steps, total_steps] the learning rate is lr_min + 0.5 * (lr_max - lr_min) * (1 + cos(pi * progress)), where progress = (step - warmup_steps) / max(1, total_steps - warmup_steps). At step = warmup_steps the cosine evaluates to cos(0) = 1, which gives lr_max and matches the warmup endpoint exactly. At step = total_steps it evaluates to cos(pi) = -1, which gives lr_min and matches the decay endpoint exactly.

That continuity is deliberate, and it is why the schedule is one function of step instead of three functions glued together. A glued schedule loses a boundary the first time someone changes lr_max.

Floor after total steps

For step > total_steps the rate stays at lr_min. The schedule does not raise and does not extrapolate. It pins at the floor and logs nothing, so the trainer must watch its own step count. A trainer that needs more steps changes the schedule's total_steps, not the loop.

Gradient norm logging alongside the rate

The schedule is half of training health, and the gradient norm is the other half. The loop logs both on every step. A divergent run shows its gradient norm spike before the loss does. A well-tuned warmup keeps the norm rising linearly with the rate, while a peak set too aggressively shows up as a norm that stays high after warmup. write_step_log_csv writes that log with the columns step, lr, grad_l2_norm, loss, and the CSV is the durable record once you call it (the demo only prints the rows).

code/main.py implements:

• CosineWithWarmup: a stateless function lr(step) -> float over the configured schedule.
• TrainState: wraps a model, an AdamW optimizer, and the schedule into a single step function.
• TrainState.step: runs one forward pass and one backward pass, logs the gradient L2 norm, and applies lr(step) to the optimizer.
• plot_schedule_ascii: renders the schedule as a text plot you can read in a terminal.
• write_schedule_csv: writes one row per step with the learning rate.
• write_step_log_csv: writes the per-step training log (step, lr, grad_l2_norm, loss). The demo does not call it.
• Comparison and utility helpers: InverseSqrtWarmup, LinearWarmupConstant, EWMA, split_decay_groups, and summarize_step_log.

CosineWithWarmup raises ValueError on an invalid config, such as warmup_steps >= total_steps or lr_max <= 0.

The demo at the bottom of the file builds a tiny two-layer model (nn.Sequential of Linear, GELU, Linear) and trains it for 20 steps on a fixed input batch. It prints the learning rate, gradient norm, and loss for each step, then renders the schedule as a text plot for a visual sanity check.

Run it:

python3 code/main.py

The script exits zero and prints a per-step training log plus the schedule plot.

Four patterns turn the schedule into a production artifact.

The schedule lives in a config, not in code. The trainer reads warmup_steps, total_steps, lr_max, and lr_min from a YAML or JSON config committed to git. Because the config is content-addressed, the schedule is reproducible. Because it shows up in the PR diff, it is auditable.

The step counter is monotonic and decoupled from epochs. Some frameworks confuse step and epoch when the dataset is sharded or the dataloader restarts. Read global_step from the trainer's checkpoint instead of a local counter, and a resumed run picks up at the right point in the schedule. The step counter is the durable axis.

The schedule plot goes in the run directory. Each training run writes outputs/lr_schedule.png (a text plot, in this lesson) next to its other artifacts. A reviewer skimming the directory can sanity-check the schedule without re-running anything, which catches misconfigured schedules at PR time.

The log row schema is fixed. The columns are step, lr, grad_l2_norm, loss, in that order. Downstream notebooks and dashboards read that schema, and renaming a column without bumping a version breaks every existing dashboard.

Production patterns:

• Sweep peak before sweeping anything else. lr_max is the most sensitive knob. Sweep it on a small model first. The optimal lr_max scales weakly with model size, so the small-model sweep is a strong prior.
• Set warmup as a fraction of total steps, not an absolute count. A 200-million-step run with 2,000 warmup steps reaches peak almost immediately, while a 20,000-step run with the same count warms up for 10 percent of training. Configure warmup as a fraction (typically 1-3 percent) so the schedule scales with training duration.
• Keep lr_min non-zero on purpose. A floor at 10 percent of lr_max keeps the optimizer learning through the long tail. With lr_min = 0 the training curve looks great on a plot, but the model has not finished training.

On a real project, outputs/skill-cosine-warmup.md would record which config carries the schedule, which trainer step the global counter is read from, and which lr_max sweep produced the deployed value. This lesson ships the engine.

1. Compare the existing InverseSqrtWarmup against the cosine on a 200-step toy run, then extend it with an lr_min floor. Which curve reaches the lower final loss?
2. Add a --restart flag that inserts a second warmup at total_steps / 2. Decide whether warm restarts help or hurt on the toy run, and defend the call.
3. Add a unit test that the schedule is continuous: for every step in [0, total_steps], |lr(step+1) - lr(step)| is bounded by lr_max / warmup_steps.
4. Wire the schedule into a torch.optim.lr_scheduler.LambdaLR so it composes with framework code. The lesson uses a plain step function, so what does the wrapper change?
5. Add a --plot-png flag that writes a real plot via matplotlib. Decide whether the text plot or the PNG is the better default for CI runs, and defend it.

• Loshchilov and Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts (arXiv 1608.03983): the cosine schedule's reference paper
• Loshchilov and Hutter, Decoupled Weight Decay Regularization (arXiv 1711.05101): AdamW's reference paper
• PyTorch torch.optim.lr_scheduler: how step functions compose with framework schedulers
• Phase 19 · 42: the downloader whose corpus this schedule consumes
• Phase 19 · 43: the dataloader the schedule co-evolves with
• Phase 19 · 45: gradient clipping and AMP, the next layer in the loop