Optimizers & Training Tricks Primer

What w ← w − η · ∇L hides is that ∇L for one weight, measured on one mini-batch, is mostly noise. Adam stops trusting any single gradient and keeps two running summaries of the stream instead — one of it, one of its square.

The first is an exponentially weighted mean: it keeps β₁ of what it had and takes 1 − β₁ of each new gradient, so every past step is still in there at a geometrically decaying share. Each disc is one gradient's share of the register — drag β₁ and watch the newest step hand its share back:

Notice that at β₁ = 0 one disc carries the whole 100% and the register reads 0.769 — the newest gradient, unchanged. An average over a window one step long is the gradient. At β₁ = 0.9, a ten-step window, the newest step is down to 13.9%, the oldest of the twelve still supplies 4.4%, and the register reads 0.470. Momentum is exactly this and nothing more.

The second summary averages g² rather than g, and the useful thing is its square root: √v̂ estimates how big this parameter's gradients are, with the sign thrown away. Raise β₂ and watch the second register fill beside the first:

Watch √v̂ at β₂ = 0: it is the newest sample's own magnitude, 0.769, so it reports one draw and nothing about the parameter. By β₂ = 0.95 the newest step supplies 10.9% of it and the register settles at 0.677 — this parameter's typical gradient size, for one multiply-add and one stored number per weight.

Divide one register by the other and the units cancel. m̂ / √v̂ is a pure number of order one whatever scale the gradients live at. Shrink the whole stream fifty-fold and watch both registers collapse while the gauge along the floor does not move:

Notice that m̂ falls from 4.70e−1 to 9.40e−3 and √v̂ from 6.77e−1 to 1.35e−2 — fifty times smaller, both of them — while the gauge holds at 0.69. So η stops being a compromise between the steepest parameter and the shallowest, and becomes one distance every parameter travels.

One correction is still missing. Both registers start empty, so while the window is filling they are sums over steps that have not happened — dividing each by 1 − βᵗ removes that exactly. Walk the window back and watch the slots that are still empty open up under the raw total:

Because 1 − 0.9¹ is 0.1, at step 1 the raw total is 0.135, a tenth of what it should be, while the corrected register already reads 1.353 — the gradient itself. By step 12 the raw total is 0.337, 0.718 of the corrected 0.470. It matters far more for v than for m, which is the thread §03 picks up.

Adam hands every parameter a step of about the same size for as long as the gradient points one way, and nothing in m̂ / √v̂ shrinks as the weight grows. So a weight the loss finds useful grows without bound — and a network with enormous weights is one bad batch from an activation that overflows.

The cure is a term that grows with the weight: subtract η · λ · w after the adaptive step, every step, whatever the gradient did. Raise λ and watch the weight stop climbing:

Notice it settles rather than shrinks. Push and pull balance at E|m̂/√v̂| / λ, so at λ = 0.10 the weight holds near 1.70 against that ceiling while the undecayed one has reached 3.51 and is still going. Weight decay is not a pull toward zero; it is a ceiling whose height is 1/λ.

Now the part the field took three years to notice. Adam's original paper folded decay in as L2 — add λ · w to the gradient — and the sum then goes through the same 1/√v̂ divisor as everything else. Switch between the two and watch what one λ does to six parameter groups:

Under AdamW every group loses the same 3e-5 of itself per step. Under L2 in the gradient the realised decay is η · λ / √v̂: the embedding, whose gradients are large, loses 1e-3; a bias, whose gradients are smaller, loses 7.5e-1 — three quarters of itself, every step. One λ, a 750× spread, and no error anywhere.

That is the fix that made AdamW the default everywhere. What it costs is memory — twelve of every sixteen bytes a parameter occupies are the optimizer's, not the model's. Drag the model size and watch the column stand against one H100's 80 GiB:

Notice which cells are grey. Four bytes of bf16 weight and gradient are the model; the other twelve are the optimizer's. At 7B that is 104 GiB against 80, so a quarter of the column stands clear of the rule; a 70B wants 1,043 GiB, before any activation.

So stop replicating it. FSDP shards all sixteen bytes across the data-parallel group, so each GPU holds 16P/N: split it two ways and 104 GiB becomes 52.2 GiB, under the rule. The bill is bandwidth, which is what §07 is made of.

§01 ended on bias correction being exact, and it is — in expectation. What no algebra fixes is that at step 1 the second average holds one sample, and Adam is about to divide by its square root.

Work it through: at t = 1, m̂ = g₁ and v̂ = g₁², so the update is ±1 for every parameter in the model whatever its gradient. Change β₂ and watch the opening step refuse to move:

Notice the opening value is 1.00 at every β₂ the slider reaches. There is nothing to average yet, so the ratio is the sign of one noisy number — with no ramp, a full η step in a random direction for all seven billion parameters at once. Thirty of the first forty updates still exceed 0.5.

One sample is the extreme case; the useful question is how many steps it takes before the divisor is worth trusting. Follow twelve independent parameters and watch the band of √v̂ close on the truth:

At step 1 the twelve are spread from 0.01 to 2.86 times the truth — the worst of them is 186% away, and one is 99% below — and they open there at every β₂, because at step 1 each one simply is its own single sample. β₂ earns its keep later: by step 1,000 the band has closed to ±2% at β₂ = 0.999 and is still ±19% at 0.95, because a short window never stops resampling.

So the fix is not a better divisor. It is a smaller step while the divisor is bad. Ramp η linearly from zero and watch the distance a step actually moves the weights flatten:

With no warmup the first step moves each weight by the full 3.00e-4; with 2,000 warmup steps it moves 1.50e-7, smaller by exactly the ratio of the two. Real LLM warmups run 500–2,000 steps — 0.1–2% of a run, and, not by coincidence, about the 1,000-step window that β₂ = 0.999 averages over.

Warmup covers the first fraction of a percent. The other ninety-nine is a decay of η, and every training report from every lab gives the same three numbers for it: the peak, the shape, and the floor.

The four shapes in use differ less than their names suggest. Walk the cursor through a hundred thousand steps and switch the curve underneath it:

Watch what the curves integrate to. Cosine and linear both come to 16.47 over this run — identical budgets spent differently, cosine holding high early and dropping fast at the end. WSD is the outlier at 27.05 because it holds the peak until the last fifth, which is also what makes it resumable: you need not know the total step count when you start.

Where the curve ends is a bigger lever than its shape. Lift the floor off zero and watch the area under the curve grow:

Because that area is the total distance the optimizer may travel, a floor at 10% of peak buys 9.8% more of it — spent entirely in the tail, where the model is polishing rather than exploring. Decay to zero when the token budget is fixed; keep a floor when you may want to continue.

The peak is the number worth agonising over, and it is not free to pick: it falls as the model grows. Step through GPT-3's eight models and read both axes:

Both axes are logarithmic, so the straight line is a power law and not a proportion — each step right is a tenfold model, and the line drops by a fixed fraction across it. Fitting the eight gives η ∝ N^−0.31, and over the 1,400× from 125M to 175B the peak falls exactly 10×, 6.0e-4 to 0.6e-4 — a rule of thumb with GPT-3's own Table 2.1 behind it.

Adam bounds the update per parameter. It does not bound the update's total length: one bad batch can make every coordinate of the gradient large at once, and the loss curve never comes back from that leap. Clipping bounds the whole vector.

The rule is a line of arithmetic — if ‖g‖ exceeds the threshold, scale the whole gradient by c/‖g‖. This one already does. Drag the gradient around the ball and watch what comes back:

Notice what does not change. The clipped vector lies on the same ray, at 0.0°: clipping changes how far the step goes and never where it goes. Everything else in this section depends on that.

In a real run the threshold spends most of its life doing nothing at all. Lower it across two hundred steps of gradient norm and watch how many it catches:

At the standard threshold of 1.00 it fires on 4 of 200 steps — the spikes, the largest of them 7.39 — and leaves the other 196 exactly as they were. Drag it to 0.10 and it fires on 180 of 200: every step now has the same length, the effective rate has become c·η/‖g‖, and the run is quietly a much slower one that no error message will mention.

There is one way to get this wrong and it is one identifier away in every framework. This scene opens on clip_grad_value_ — a box, not a ball — so drag the gradient around it and watch the direction move; switch back to the norm and the turn goes to zero:

clip_grad_value_ clamps each coordinate separately, so it shortens the large one and leaves the small one alone: the opening gradient (2.0, 1.0) becomes (1.0, 1.0), turned 18.4° off the direction the backward pass computed. Not a smaller step — a different one, and it fails silently: a run descending 18.4° off the gradient still descends. Use clip_grad_norm_.

Everything above assumed the numbers were exact. They are not. Modern training keeps weights, activations and gradients in two bytes, which halves the memory and roughly doubles the throughput on the same units: an H100 SXM does 989 dense bf16 TFLOPS on its tensor cores against 495 in TF32 on those same cores. Its 67 TFLOPS of vector fp32 is a different datapath, and a 15× measured against that is a number no matrix multiply sees.

A float spends its bits on two jobs: exponent bits buy range and mantissa bits buy precision, and with sixteen to divide, every bit one gets is one the other loses. Slide the split and watch the reachable interval stretch:

fp16 and bf16 are two points on this one line. fp16 takes 5 exponent bits and 10 mantissa; bf16 takes 8 and 7, which is fp32's exponent field exactly. So bf16 reaches 3.4e38 where fp16 stops at 6.6e4, and pays 2.4 decimal digits against fp16's 3.3.

Precision is not an abstraction: it is the gap between the two values a format can actually hold. Drag a number along the axis and watch where the cast puts it:

Because that gap is a fixed fraction of the value — 2⁻⁷, or 0.78%, everywhere in bf16 — the relative error is the same at 1e-30 as at 1e30, and it is eight times fp16's 0.098%. Every number lands on one of the two ticks. Nothing lands between them.

Range is what breaks fp16 in practice. Gradients live several decades below one and fp16's smallest subnormal is 5.96e-8. Slide the loss scale and watch it carry the distribution between the two walls:

With no scaling, 8.0% of this modelled spread is under fp16's floor — flushed to zero, silently, and those parameters do not train. Multiply the loss by 2¹⁰ before the backward pass and it is 0.05%; push to 2²⁴ and 5.9% overflows to inf at the top instead. bf16 keeps the whole spread at every setting, which is most of why it won.

One thing stays fp32 whatever the format, and this is the figure that says why. Compare what a weight of 1.0 actually gains when an update is added to it:

An update of 1e-3 lands as 9.77e-4 in fp16 and as nothing at all in bf16 — it is under half of bf16's 0.0078 gap, so the addition returns the weight unchanged. With η at 3e-4 that is most of the updates in a real run. So the working copy is bf16 for the speed and the master copy is fp32, so that training accumulates at all.

§02 counted 104 GiB of state against an H100's 80 GiB and answered it by sharding that state across the data-parallel group — memory bought with bandwidth, which is the currency of this whole section. Every frontier run is cut along three axes that compose, and what sizes each one is what it has to say on a wire.

The easiest is data parallel: every GPU holds the whole model, runs its own slice of the batch, and then they average gradients. The averaging is the interesting part — step the ring and watch each chunk gather one contribution per hop:

Notice nobody is a hub. Over the first N−1 phases each chunk walks one step round the ring and picks up one GPU's contribution at every stop, so at phase 3 every GPU owns one fully reduced chunk; the next N−1 phases hand those finished chunks back round. Nothing ever passes through a coordinator.

That matters because of what the alternative costs. Grow the GPU count and watch the ring and the central server separate on a log–log frame:

A ring's busiest link carries 2(N−1)/N · M, which climbs toward 2M and stops: 13 GiB of bf16 gradients cost 22.8 GiB per link at 8 GPUs and 26.0 GiB at 512. A central server's link carries 2NM — 209 GiB at 8 and 13.0 TiB at 512. One is flat in N and the other is a straight rising line on a log–log frame, which is what a power of N looks like. That is why all-reduce is a ring.

Data parallel needs the model to fit on one GPU. When it does not, cut the weight matrices instead: each GPU keeps a slab of every layer. Split the layer and watch what the cut costs:

A 4096-wide layer is 201M parameters; split eight ways each GPU holds 25.2M. What it buys in memory it pays in traffic: four all-reduces of the activation per layer per iteration, not two — Megatron splits each block into an f/g pair, and a layer has two such blocks, so two all-reduces go forward and two come back. That is 448 MiB at eight-way for one micro-batch of 8,192 tokens, whose activation is 64 MiB to begin with. Every layer, every step — which is why tensor parallelism stays inside one NVLink node and does not cross the network.

The third cut runs along the layers: stage 0 owns the first blocks, stage 1 the next. That one has a cost with no communication in it at all. Add micro-batches and watch the empty slots close:

With one micro-batch, four stages are idle 75% of the time: a stage cannot start until its predecessor has finished. The bubble is exactly (p−1)/(m+p−1), so it is fought with micro-batches — sixteen takes four stages to 16%. Real runs use all three cuts at once: a 1,024-GPU job might be 32-way data × 8-way tensor × 4-way pipeline, and those three numbers multiply to the GPU count.

Read a training script and the machinery is unremarkable: a constructor, a loop, and seven calls inside it. What makes it a working run rather than a diverging one is that each call answers a specific failure — and the answers have to arrive in the right order.

Step through one iteration and read the knob each operation owns — and how many bytes per parameter it moves through memory:

Notice which row is long. The optimizer step moves 28 of the 56 bytes per parameter the whole iteration moves, against the forward pass's 2 — which is why an Adam step is bandwidth-bound and why fusing or sharding it buys as much as the arithmetic inside it. The order is not arbitrary either: unscaling comes after the all-reduce and before the clip, because the norm compared to 1.0 must be the true one — clipping a scaled gradient clips at the wrong threshold, silently.

opt = AdamW(p, lr=3e-4, betas=(.9, .95),
            eps=1e-8, weight_decay=0.1)

for step in range(1, TOTAL + 1):
    set_lr(opt, lr_at(step))     # §03, §04
    with autocast(dtype=bfloat16):  # §06
        loss = model(batch).loss
    loss.backward()              # §07 ring
    clip_grad_norm_(p, 1.0)      # §05
    opt.step(); opt.zero_grad()  # §01, §02

Those are the numbers a frontier lab publishes and a fine-tune inherits.