The learning rate is the hyperparameter that breaks production

Part 8 ended with a gradient: for every weight, which way to move it. The optimizer decides how far. That decision is a handful of numbers in a config file, and nothing in the stack checks them.

I read torch’s Adam step and nanoGPT’s schedule, then trained Part 8’s character model, four blocks that read one character at a time, 21 ways: seven learning rates, each with no warmup, with warmup, and with warmup plus a cosine decay. None of them crashed.

Three optimizers, five lines

Plain gradient descent subtracts the gradient times a step size. Momentum keeps a running total of past gradients, shrunk a little each step, and moves along that instead. A weight pushed the same way twice picks up speed. Adam adds a second average, of the squared gradient, and divides by its square root. Here is torch’s step, condensed from torch 2.14’s adam.py:456-546:

exp_avg.lerp_(grad, 1 - beta1)                    # average of the gradient
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)   # average of its square
step_size = lr / (1 - beta1 ** step)              # bias correction, early steps
denom = exp_avg_sq.sqrt() / (1 - beta2 ** step) ** 0.5 + eps
param.addcdiv_(exp_avg, denom, value=-step_size)  # one step per weight

The division is the whole idea. A weight whose gradient has been large lately gets divided by a large number. A weight with small gradients gets divided by a small one. Every weight ends up with its own step size, set by its own recent history. lr only sets the scale. Two config numbers, beta1 and beta2, set how fast each average forgets what it saw; nanoGPT uses 0.9 and 0.95.

AdamW adds one line, adam.py:419: param.mul_(1 - lr * weight_decay), which shrinks every weight a little on every step. Where that line sits is the last section.

The loss scale cancels, and memory pays for it

Part 8 multiplied a loss by 40 and nothing changed. Here is why. Multiply every gradient by 40. Then exp_avg grows 40 times and exp_avg_sq grows 1,600 times, so its square root grows 40 times too. The factor cancels in the division. It cancels everywhere except eps. That constant is added to the denominator to keep it away from zero. Scaling the loss by 40 divides it by 40. If you’ve changed a metric’s units and watched the alerts hold because every threshold was relative, this is that.

On three steps of the character model, the median weight under AdamW moved 0.0022 whether the loss was scaled or not, and the two update vectors pointed the same way; under plain SGD the same change multiplied the median step by 17.5 and turned the update 30 degrees away. The cancellation is not identical weight by weight: the largest disagreement was 0.00134, over half the size of a typical step.

That robustness is bought with memory. Adam keeps two numbers per weight, which torch calls exp_avg and exp_avg_sq; SGD with momentum keeps one and plain SGD none. I read that off torch’s optimizer state on the character model and got exactly those counts. For GPT-2 small in fp32 that is 996 MB of optimizer state, next to 498 MB of weights and 498 MB of gradients. Call it 2 GB before a single activation from Part 8’s graph.

Read the three bars that don’t move with the batch: on AdamW they come to the 2 GB above. Switch to SGD and watch a third of the total disappear. Then raise the batch and see what activations do to the same budget.

The first step is the whole learning rate

Look again at step_size = lr / (1 - beta1 ** step). On the first step exp_avg is the gradient times 1 - beta1, and the correction divides by exactly that. The two cancel. The first update is the gradient divided by its own size, times lr.

On step one the median weight moved 0.001 with lr set to 0.001. Exactly the learning rate, whatever the gradient was. The gradient’s size never reaches the update. Only its sign survives, weight by weight. From the second step on, what survives is each gradient measured against its own recent size.

That exactness has two conditions, and both are config. A gradient close to eps (10⁻⁸ by default) breaks it: at 10⁻⁹ the first step is a tenth of the rate, at 10⁻¹² a ten-thousandth. And torch’s AdamW applies a default decay of 0.01 unless you pass your own. That is why my largest first step was 0.00101, just over the rate.

That is the argument for warmup. Start the rate near zero and ramp it up linearly over the first few hundred steps, and those first exact-size updates are small ones. nanoGPT’s schedule, train.py:231-242, is eight lines:

if it < warmup_iters:                       # 1) linear warmup
    return learning_rate * (it + 1) / (warmup_iters + 1)
if it > lr_decay_iters:                     # 2) then a floor
    return min_lr
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
assert 0 <= decay_ratio <= 1
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))   # 3) cosine in between
return min_lr + coeff * (learning_rate - min_lr)

With 100 warmup steps my first update was 0.0000099 instead of 0.001, a hundred times smaller. It grew from there. Everything after warmup then walks the rate back down to a tenth of the peak, which is what nanoGPT’s min_lr at train.py:68 sets. Its one assert checks the arithmetic of the ratio, not the value you passed in.

What the sweep says

I trained the character model for 600 steps at seven learning rates, from 0.0001 to 0.1, times three schedules: none, 100 steps of warmup, and warmup plus the cosine decay. Gradients were clipped at 1.0 throughout, as nanoGPT does. The sweep cannot separate warmup from clipping, and clipping was already doing work. At 0.003 without warmup, turning it off took the score from 2.18 to 2.38 and let the gradient norm reach 79. Frequency-only guessing scores 3.23 on held-out text.

Nothing produced a NaN, including a learning rate ten times the best one.

The best run scored 1.73 at a rate of 0.01, with warmup and cosine. The same rate scored 1.84 with warmup alone and 2.61 with neither, and 2.02, at 0.001, was the best any rate managed without warmup. So warmup didn’t rescue a broken run. It unlocked a rate ten times higher, and that rate is where the good score was.

At a constant 0.1 the loss flattened out near 3.0, barely better than guessing by letter frequency. The loss went down. It just stopped going down. You’d only call that value wrong if you knew what the same model does at 0.01. With the cosine walking the rate back to 0.01, the 0.1 run kept improving, to 2.71, still almost a nat behind the best run.

Start at 0.01 and tap no warmup, which drops the cosine with it, and the held-out score rises from 1.73 to 2.61. Then drag up to 0.1 and watch the curve go shallow instead of exploding, with the NaN readout still saying never.

Be fair about what a laptop can show. A four-block model on one is not a pretraining run, and every number here is one seed, so nothing under about 0.05 nats is worth reading. At real scale a bad rate does produce the loss spike and the NaN. What carries over is the shape. The same rate is fine with warmup and wasteful without, and the failure you should plan to catch is not a crash.

Decay, and where you put it

Weight decay shrinks weights towards zero. The original way was to add weight_decay * param to the gradient. That sends it through Adam’s division, so it gets rescaled per weight like everything else. AdamW’s param.mul_(1 - lr * weight_decay) keeps it out of that path, which is what the W stands for.

The difference is not subtle. Trained the same way for 400 steps with decay 0.1 on every parameter, the decoupled version scored 2.22 on held-out text; the coupled one scored 3.03, barely better than frequency guessing. It ended with a weight norm, the length of all its weights taken as one vector, of 21 against 41. The decay Adam rescaled was strong enough to hollow the model out.

Two things about that number. The same run with no decay at all scored 2.21. Over 400 steps the decoupled decay bought nothing, and the whole gap is damage from the other path. Decaying every parameter likely makes the coupled case worse: the shrinking gradient also lands on the norm scales, which start at 1.0 and hold most of the weight norm at initialization.

nanoGPT decays selectively, model.py:269-274: matrices get decay, biases and LayerNorm scales don’t. Two lines of filtering, and no error if you skip them.

What the ML engineer watches

Four habits, and the first two are just numbers you keep.

Log the effective learning rate, not the configured one. The number that matters is what the schedule returned this step. A schedule that restarted on resume ran warmup again, starting 70 times below a full resume’s rate and ending the run over three times above it. The config file said the same thing in both runs. Over the same 150 steps it scored 2.275 against 2.230, inside one-seed noise: only the logged rate would have shown it.

Checkpoint the optimizer state, not just the weights. AdamW’s two averages are 996 MB for GPT-2 small. Drop them and the averaging silently starts over. The first step after resume cost 2.64 where a full resume cost 2.38, on the same batch. Mine recovered inside 150 steps, to within 0.007. At this scale the state was worth almost nothing. The transient is the argument, and I did not test a rate large enough for it to be fatal.

Size the machine at 16 bytes per parameter, before activations. In fp32, weights, gradients and AdamW’s two averages are four bytes per parameter each, so sizing by the weights alone is wrong by three quarters. Mixed precision doesn’t shrink the sum: 2-byte weights and gradients plus a 4-byte master copy and the two averages still come to 16 bytes, and weights-only sizing is then wrong by seven eighths.

Sweep the rate on a small model, and set an abort rule. The small model gives you the shape, not the value: the best rate falls as the model grows, unless you use something like μP. Nothing raises, so the rule has to be yours. Mine was a loss that doubles its own best. It fired on four of the 21 runs. Three were at 0.1; the fourth went on to finish fine. A rule that cheap will kill a good run now and then. That is the price of catching the bad ones in the first ten steps.

None of this is modelling work either. Five lines of arithmetic, an eight-line schedule, and a config file nothing validates. Between my best run and my worst sat 1.3 nats, on the same model, the same data and the same seed. The only difference was three numbers in the config.