Learning-rate schedules & warmup

01 · Concept

Concept

Halfway through a long pretraining run, the cluster grows. Twice as many nodes join, the global batch doubles from four million tokens per step to eight million, throughput improves beautifully, and everyone is pleased. The run finishes on schedule in token terms, the loss curve looks reasonable, and the final checkpoint is worse than it should be — noticeably worse on downstream evaluation than a smaller pilot predicted. The cause is not the data, the architecture, or the extra nodes. It is that the learning-rate schedule was written in steps.

The optimizer turns gradients into parameter updates, but the learning rate sets their scale, and one constant rate rarely serves an entire run. Early training has uncalibrated activations and unreliable optimizer moments; the middle wants substantial progress; the end benefits from small, refining updates. A schedule encodes that changing risk.

Warmup starts below the intended peak and rises over an initial interval. Linear warmup at step $t$ is simply

for $t\le T_{warmup}$. The first updates stay restrained while normalization scales, gradient magnitudes, and Adam’s moment estimates settle. Warmup does not repair an excessive peak; it only changes how that peak is approached.

After warmup, cosine decay is the modern default. Writing $p$ for the fraction of the decay horizon completed,

Evaluate it once by hand with a peak of 3e-4 and a floor of 3e-5, a typical ten-percent minimum. At the quarter mark, $\cos(\pi/4)=0.7071$, so

At the halfway point, $\cos(\pi/2)=0$, so

The curve is gentle near both ends and steepest in the middle, which is why it spends a long time near the peak and then a long time near the floor.

Now return to the opening scenario with those numbers in hand. Suppose the plan was a two-trillion-token budget at four million tokens per step, giving 500,000 steps, and the scheduler’s horizon was configured as 500,000 steps. Doubling the global batch to eight million tokens means the same two trillion tokens are consumed in 250,000 steps. The run ends when the data ends — at $p=0.5$ — with the learning rate still at 1.65e-4, having never entered the decay phase at all. Every refinement that the second half of a cosine provides simply did not happen, and nothing in the logs looked abnormal, because the rate followed its configured curve faithfully.

The alternatives to cosine are worth knowing. Linear decay falls at a constant pace toward a chosen endpoint and makes remaining progress trivial to read off. Inverse-square-root decay, proportional to $t^{-1/2}$ after a warmup construction, was the original Transformer recipe and has the useful property of not requiring a declared horizon. Constant-with-warmup keeps the peak after ramp-up, which suits an uncertain total horizon but can leave late updates noisier than a nearly converged model wants. Endpoints matter as much as shapes: decaying exactly to zero assumes the run ends where planned, and extending it afterwards produces no learning unless the schedule is rewritten, while a nonzero floor preserves adaptation but keeps disturbing a settled model. Warm restarts, popular elsewhere, introduce deliberate rises and are rarely appropriate for an expensive single-pass pretraining run.

AdamW separates gradient-based optimization from weight decay, but the scheduler and the decay still interact, and parameters such as normalization scales and biases are usually excluded from decay entirely. Gradient clipping bounds an update driven by an outsized gradient before the optimizer step; it is a guardrail, not a substitute for a sane schedule. Scaling the global batch can justify adjusting the peak rate, but linear and square-root rules are heuristics with regime boundaries — larger batches reduce gradient noise and may want longer warmup or different optimizer tuning, so run pilot sweeps rather than extrapolating a rule across orders of magnitude.

Log the realized rate at every step alongside token count, loss, gradient norm, optimizer statistics, and skipped mixed-precision updates. On resume, restore the scheduler’s position: accidentally restarting warmup after a checkpoint changes the optimization path, and jumping to a wrong late-run rate manufactures exactly the kind of spike that lesson 5.12 has to diagnose. Overlaying the rate on the loss curve helps correlate instability with update scale, though correlation alone never proves the schedule caused a spike.

The durable model is controlled momentum for learning. Warmup avoids wrenching an uncalibrated network, the central phase moves quickly through useful parameter space, and decay narrows the step as the run approaches its planned end. The schedule is part of the experiment specification and part of checkpoint state — not cosmetic optimizer configuration, and not something to express in a unit that changes when your cluster does.