Part I — Basics

Training — AdamW, schedule, stability

Lesson 03 left you with a clean, modern architecture: pre-norm + RMSNorm, RoPE, SwiGLU, GQA, no biases. It is a differentiable function with a cross-entropy loss — and nothing more. Nothing in it says how to move the parameters θ downhill without the run blowing up, and nothing tells you how to spend a fixed compute budget so the loss actually bottoms out instead of stalling or diverging. This lesson is the optimizer recipe that every frontier run uses, derived one forced choice at a time: AdamW, warmup→cosine, careful init, gradient clipping, and bf16 mixed precision.

The previous step left this broken

You have a model and a loss; you do not have a procedure to minimize it. Plain gradient descent on a Transformer diverges or crawls: the loss surface is wildly anisotropic (the gradient for an embedding row and for an attention projection differ by orders of magnitude), the gradients are heavy-tailed (rare huge values from rare tokens or a bad batch), and at the depth and learning rates needed to finish a billion-token run inside budget, a single bad step can spike the loss to NaN and torch the GPU-hours you have already spent. You need an optimizer that is both stable and budget-efficient, plus the scaffolding (schedule, init, clipping, precision) that keeps it stable.

Linear position

Forced by: a defined model + a fixed loss still need an optimizer that converges fast and never diverges, on a budget with no do-overs.
New idea: the standard recipe — AdamW (per-parameter adaptive step from the 1st/2nd gradient moments, with decoupled weight decay), a linear warmup → cosine learning-rate schedule, careful init, global-norm gradient clipping, and bf16 mixed precision over fp32 master weights.
Forces next: before you launch this you must predict whether it fits in memory and finishes in budget — the optimizer state alone is 12 bytes/param, which lesson 5 turns into the memory wall.

The plan

Eight moves. (1) Why plain SGD fails on a Transformer. (2) The AdamW update equations, with the standard constants. (3) The warmup→cosine schedule and why each half exists. (4) Batch size, gradient accumulation, and the critical batch size. (5) Init that survives depth. (6) Gradient clipping as the spike fuse. (7) bf16 mixed precision and why it beats fp16. (8) Reading a loss curve and what to do when it spikes. Then a widget where you set warmup, peak LR, clip, and β₂ and watch a toy run converge — or diverge.

1 · Why not plain SGD

SGD takes the same step size in every coordinate: θ ← θ − η·g. That is fine when the loss surface is roughly spherical. A Transformer's is not. Different parameter groups see gradients that differ by orders of magnitude — an embedding row touched by one rare token gets a tiny, sparse gradient; an attention output projection gets a dense, large one. A single global η that is safe for the large-gradient group is far too small for the small-gradient group, so most of the network barely moves while you wait. Pick η for the small group instead and the large group overshoots and the loss explodes.

On top of that, language-model gradients are heavy-tailed: across steps, the gradient norm is usually moderate but occasionally enormous — a batch with a pathological document, a long run of repeated tokens, a numerically awkward attention pattern. SGD has no memory of "how big are gradients usually here," so a rare 50× spike becomes a 50× step and the run is gone. The fix is to normalize each coordinate by its own recent gradient magnitude, so every parameter takes a step of comparable, well-scaled size, and to keep a running estimate that a single outlier cannot dominate. That is exactly what Adam does.

2 · AdamW, the update equations

Adam keeps two exponential moving averages per parameter: m (the mean gradient, "momentum") and v (the mean squared gradient, "variance/scale"). At step t, with gradient g_t:

m_t = β₁·m_t−1 + (1 − β₁)·g_t
v_t = β₂·v_t−1 + (1 − β₂)·g_t²

Because both averages start at 0, they are biased toward 0 for the first steps. Divide out that bias:

m̂_t = m_t / (1 − β₁^t) v̂_t = v_t / (1 − β₂^t)

The update divides the smoothed gradient by the smoothed magnitude — a per-coordinate normalized step:

θ_t = θ_t−1 − η · m̂_t / (√v̂_t + ε) − η · λ · θ_t−1

The second term is decoupled weight decay — the "W" in AdamW. The original Adam folded decay into the gradient (g ← g + λθ), which means the decay also got divided by √v̂: parameters with large gradients got decayed less, which is backwards. AdamW applies −η·λ·θ directly to the weights, outside the adaptive normalization, so every parameter shrinks toward 0 at the same proportional rate regardless of its gradient scale. For Transformers this is strictly the right behavior and is now universal.

hyperparameter	typical value	what it controls
β₁	0.9	momentum horizon — averages ≈ last 10 gradients
β₂	0.95	scale horizon — LMs use 0.95, lower than the 0.999 vision default, so the variance estimate reacts faster to spikes
ε	≈ 1e−8	floor on the denominator — stops a near-zero v̂ from blowing up the step
λ (weight decay)	≈ 0.1	decoupled L2 pull toward 0; regularizes, kept off norm/bias params

The β₂ = 0.95 choice is worth dwelling on: with 0.999, v averages over the last ≈1,000 gradients, so a sudden burst of large gradients barely moves the denominator and the step stays huge — straight into a spike. With 0.95 the denominator inflates within ≈20 steps of a gradient burst, automatically shrinking the step and damping the spike. Lower β₂ is the LM community's quiet stability lever; you will feel it in the widget.

# AdamW, one step, the whole thing (per parameter tensor p)
m = b1*m + (1-b1)*g
v = b2*v + (1-b2)*g*g
mhat = m / (1 - b1**t)
vhat = v / (1 - b2**t)
p -= lr * (mhat / (vhat.sqrt() + eps) + wd * p)   # decay decoupled

3 · The learning-rate schedule: warmup → cosine

A constant LR is wrong at both ends of training. The schedule has two halves, each forced by a specific failure.

Linear warmup. For the first few hundred steps, v̂ is built from almost no data — its estimate of each coordinate's gradient scale is garbage. Adam's normalized step m̂/√v̂ is therefore unreliable exactly when the weights are most fragile (random init, every gradient large). Starting at the peak LR here reliably produces an early loss spike or divergence. So you ramp η linearly from 0 to the peak over W steps — typically a few hundred to a few thousand, often ≈1–2% of total steps. By the time you hit peak LR, v̂ is a trustworthy scale estimate and the big steps are safe.

Cosine decay. After the peak, anneal η down a cosine from peak to a floor of about 10% of peak over the remaining steps:

η(t) = η_min + ½(η_peak − η_min)·(1 + cos(π · (t − W) / (T_total − W)))

Early on you want large steps to cover ground; late on you want small steps to settle into a sharp minimum without bouncing out. Cosine spends most of its time near the peak and only decays hard at the end — empirically it beats linear or step decay for LMs. The 10% floor (not 0) keeps the optimizer making progress to the very last step rather than freezing prematurely. One sharp consequence: cosine must be told the total step count up front, because the shape depends on T_total — you cannot honestly resume a cosine run for "a few more steps."

Peak-LR scaling. The peak is not arbitrary. As you widen the model, the right peak LR shrinks (a common rule of thumb ties it to 1/√d); as you grow the batch, it can rise (larger batches give lower-variance gradients that tolerate bigger steps, up to a point — see §4). Typical peaks land around 1e−3 for small models down to ≈1.5–3e−4 for multi-billion-parameter models. The widget's "peak LR" knob is in these units.

Hyperparameter transfer (μP). Tuning the peak LR by trial and error on the target model is exactly the do-over the budget forbids — a single full-scale sweep can cost more than the run itself. Maximal-update parametrization (μP) dissolves the problem: parametrize the network (init scales, per-layer LR multipliers, attention logit scaling) so that the optimal LR is width-invariant — the 1/√d scaling above is baked into the parametrization rather than re-tuned. You then sweep the LR (and other knobs) on a cheap small-width proxy — a few-million-parameter model that trains in minutes — and transfer the winning value straight to the target N without re-tuning. This is what makes "one shot, spend the budget" honest: you de-risk the most fragile knob at 1/1000th the cost before committing the GPU-hours, so the headline run launches with a hyperparameter you already know is good rather than one you're hoping is.

4 · Batch size, gradient accumulation, critical batch size

The gradient you step on is an average over B sequences. Larger B → lower-variance gradient → you can take a larger, more confident step, and you get more parallelism per optimizer step. But two walls appear. First, a large batch may not fit in GPU memory. The fix is gradient accumulation: run k micro-batches of size b, sum their gradients, and step once — an effective batch of B = k·b at the cost of k× the time per step but 1× the memory of one micro-batch. A run might use micro-batch 8, accumulate 32, for an effective batch of 256 sequences (× sequence length = the tokens-per-step that matters).

Second, and deeper: there is a critical batch size. Below it, doubling the batch roughly halves the steps to a target loss — perfect data-parallel scaling, free speedup. Above it, the gradient is already low-variance enough that extra examples buy diminishing returns: you burn 2× the compute for far less than 2× the progress. The critical batch size grows as the loss falls (later in training you can afford bigger batches), and it is what tells you how far you can scale data parallelism before you are wasting FLOPs. Picking a batch far above critical is the most common way to silently waste a budget — the loss curve looks fine, but you paid double for it.

5 · Initialization that survives depth

Init sets the scale of the residual stream at step 0, before the optimizer has done anything. Two rules, both from §03's architecture:

Base init

Most weights drawn from N(0, σ²) with σ ≈ 0.02, or more principled σ = 1/√d (≈0.028 at d=1280). Small enough that L stacked residual additions don't blow up the stream's variance at step 0; the optimizer finds the true scale from there.

Residual-projection downscale

GPT-2's trick: scale the init of the output projection of each block (attn proj, MLP down-proj) by 1/√(2L). With 2L residual branches feeding the stream, this keeps its variance ≈ constant with depth instead of growing ∝ L — so a 48-layer model starts as calm as a 4-layer one.

Worked: at L = 24, the downscale factor is 1/√48 ≈ 0.144 — the residual projections start ~7× smaller than the other matrices. Without it, the stream variance at the final layer is ≈ 2L times the input variance and the first forward pass already produces saturated norms; the warmup then has to dig you out of a worse hole. Init and warmup are partners: good init makes the warmup's job small.

6 · Gradient clipping: the spike fuse

Even with Adam and warmup, heavy-tailed gradients mean an occasional step will see a gradient norm far above normal. Global-norm clipping caps it: compute the L2 norm of the entire gradient (all tensors concatenated), and if it exceeds a threshold c (almost always c = 1.0), rescale every gradient by c / ‖g‖ so the norm is exactly c and the direction is unchanged.

if ‖g‖ > c: g ← g · c / ‖g‖

This is what lets a run survive a loss spike instead of dying on it. A pathological batch produces a 30× gradient; un-clipped, Adam takes a wildly oversized step and the loss jumps from 2.1 to 8.0 (or NaN). Clipped to norm 1.0, the same batch nudges the weights a normal amount, the loss ticks up slightly, and the next clean batch recovers it. Global-norm (not per-parameter) clipping preserves the relative scale across parameters, so the step direction is still meaningful. You will watch this knob save or sink the run in the widget — turning clipping off is one of the three ways to make the toy run diverge.

7 · Mixed precision: bf16 over fp16

Training in fp32 everywhere wastes half the memory bandwidth and tensor-core throughput. Mixed precision keeps a single fp32 master copy of the weights and the optimizer state, but does the forward/backward compute in 16-bit. The split that matters:

bf16 computeForward and backward run in bf16; activations and the gradient are bf16. ~2× the throughput and half the activation memory of fp32.

fp32 master weightsThe authoritative weights live in fp32. The bf16 step is applied to them, so tiny updates (LR · small gradient) aren't rounded away — a bf16-only weight would stop moving once the update fell below its precision.

fp32 optimizer stateAdam's m and v are fp32 — these accumulate over thousands of steps and need the precision.

Why bf16 and not fp16? Both are 16-bit, but they split the bits differently. fp16 has 5 exponent bits → max value ≈ 65,504 and a narrow dynamic range; gradients routinely underflow or a large activation overflows to ∞. fp16 training therefore needs loss scaling — multiply the loss by a large factor before backward to push gradients into fp16's representable range, then unscale — a fiddly, failure-prone dance. bf16 keeps fp32's full 8 exponent bits (same ≈1e38 dynamic range) and sacrifices mantissa precision instead. It almost never overflows or underflows, so no loss scaling is needed — you just compute. The cost is coarser rounding (7 mantissa bits vs fp16's 10), which the fp32 master weights absorb. On any Ampere-or-newer GPU, bf16 is the default and fp16 is legacy.

bf16 is not the end state: fp8 training. On Hopper/Blackwell, the matmuls can run in fp8 for roughly another 2× throughput. Two 8-bit formats split the duty: E4M3 (4 exponent, 3 mantissa — more precision, narrower range) for the forward and weights, and E5M2 (5 exponent, 2 mantissa — more range, less precision) for the gradients, which are heavier-tailed and need the headroom. With only a handful of mantissa bits, fp8 cannot just be dropped in: each tensor needs its own scaling factor to center its values in the representable range, and because recomputing that factor every step is expensive, frameworks use delayed scaling — track each tensor's recent max in a small history and reuse it. fp8 is applied selectively, to the big GEMMs (the attention and MLP projections) where the throughput win lives, while norms, the residual stream, the optimizer state, and the master weights stay in higher precision — the accumulation still happens in fp32. This is the same move lesson 6's roofline makes physical: switching to fp8 doubles the chip's peak compute, lifting the roofline ridge so a matmul that was compute-bound in bf16 has twice the FLOP/s ceiling to climb toward. Getting the scaling right is delicate kernel-level work — ../../gpu_kernels/ covers how the tensor cores actually consume these low-precision formats. The takeaway for this lesson: bf16 is today's safe default, but the precision floor keeps dropping, and the recipe's job is to spend the cheapest bits that still converge.

Preview — the 16 bytes/param wall (lesson 5)

Tally the bytes per parameter that mixed-precision AdamW must hold: fp32 master weight = 4, Adam m = 4, Adam v = 4, plus the bf16 working weight = 2 and bf16 gradient = 2. The optimizer state alone (master + m + v) is 12 bytes/param; with the working copies it is ≈16. For a 7B model that is ≈112 GB of training state before a single activation — which overflows an 80 GB H100. That arithmetic is the hinge of lesson 5 and the reason the whole systems half of the course exists.

8 · Reading a loss curve

The loss curve is your only real-time instrument. Learn to read its three signatures.

Warmup bump

Loss may plateau or wobble during warmup while v̂ stabilizes, then drop fast once peak LR engages. Normal — not a problem to fix.

Spike

A sudden vertical jump (2 → 6), usually from a bad batch or too-high LR. If clipping catches it, the curve recovers within tens of steps. If not, it goes to NaN and the run is dead.

Divergence

Loss climbs and does not come back. Cause: peak LR too high, warmup too short, or no clipping. The run is unrecoverable from here.

When a real run spikes and you have checkpoints, the standard playbook — in order of cost — is: rewind to the last good checkpoint and skip the offending data batches (a specific document often causes it); if it recurs, lower the peak LR (or tighten clipping / lower β₂) and resume. Frontier teams checkpoint frequently precisely so a spike costs hours, not the whole run. The lesson: stability is not luck — it is warmup + clipping + a sane LR + frequent checkpoints, and you tune them by reading this curve.

9 · Drive a run yourself

The widget below is a toy SGD-on-a-bowl simulation dressed as a training run. The optimizer descends a noisy quadratic loss; the noise is heavy-tailed (occasional large gradients = the bad-batch spikes from §1). The knobs are the four levers from this lesson. Watch the curve and the verdict: a sane configuration (warmup > 0, moderate peak LR, clip on, β₂ = 0.95) glides to a low final loss with zero spikes. Now break it. The headline failure is peak LR high with warmup = 0: the first steps land at full LR while v̂ has seen almost no data, so the normalization can't protect the step, the bowl over-steps, and the overshoot compounds to divergence — the verdict flips to DIVERGED. Restore a warmup and the same peak LR is safe, because the LR is ≈0 over exactly those fragile first steps. The second failure is clipping off: the heavy-tailed bursts that clip normally absorbs now pass through and the spike count climbs from ~0 into double digits. The β₂ slider is the variance-smoothing knob: it sets how many recent gradients the denominator v̂ averages over (≈1/(1−β₂) steps). A higher β₂ holds the denominator elevated longer after a burst — a longer memory; a lower β₂ forgets it faster and tracks the most recent gradients. (In real LM training the §2 trade-off bites the other way — see the note below.)

Schedule + spike sim — set the knobs, run the training

The four levers from this lesson drive a toy noisy descent. Set warmup = 0 with a high peak LR and the run diverges in the first few steps — that is the optimizer stepping at full size before v̂ has a trustworthy scale estimate; a warmup keeps the LR ≈0 over exactly those steps and the same peak LR converges. Clip off lets heavy-tailed spikes through. β₂ sets how long the variance denominator remembers recent gradients. Hit run to animate; KPIs report the final loss, the spike count, and whether it diverged.

warmup steps: 40 peak LR (×1e−3): 14 grad clip: β₂: 0.950

Final loss

—

# spikes

—

Diverged?

—

Peak LR reached

—

The β₂ caveat (toy vs real training). In this 1-D toy the gradient magnitude is bounded, so β₂ behaves only as the smoothing horizon — its effect on the spike count is marginal in either direction. In a real LM the bursts are heavy-tailed and unbounded, and the §2 mechanism dominates: β₂ = 0.999 averages v over ≈1,000 gradients, so a sudden burst barely moves the denominator and the full-size step lands — a spike — whereas β₂ = 0.95 inflates the denominator within ≈20 steps and shrinks the step before it does damage. That is why LMs use the lower 0.95: faster-reacting variance, fewer real spikes. The toy shows you the smoothing knob; the §2 reasoning tells you which way to turn it on a real run.

Notice the asymmetry the widget makes physical: warmup and clipping cost you almost nothing when the run is healthy, but they are the only things standing between a heavy-tailed gradient and a dead run. That is why every production recipe ships all four levers on by default.

Failure modes & checklist

Failure modes

No warmup at full peak LR. The first steps use a garbage v̂ and a full-size step. Signal: loss spikes or NaNs within the first few dozen steps, before any real learning.
Coupled weight decay (plain Adam). Decay gets divided by √v̂, so high-gradient params decay less. Signal: regularization "doesn't work like the paper said"; embedding/output norms drift.
Clipping off. A heavy-tailed batch produces a 30× gradient and a 30× step. Signal: isolated vertical spikes in the loss that don't recover; sudden NaN.
β₂ = 0.999 on an LM. The variance denominator reacts too slowly to gradient bursts. Signal: more frequent, larger spikes than the same run at 0.95.
fp16 without loss scaling (or bf16 assumed where only fp16 exists). Signal: gradients silently underflow to 0 → loss flatlines; or activations overflow to inf → instant NaN.
Batch above the critical batch size. Signal: loss curve looks normal but you used ≈2× the FLOPs to reach the same loss as a smaller batch would.
Resuming a cosine schedule for "a few more steps." The shape depended on the original T_total. Signal: LR is already near the floor; the extra steps barely move the loss.

Checklist

AdamW with β₁=0.9, β₂=0.95, ε≈1e−8, wd≈0.1 — decay decoupled, and off norm/bias params.
Warmup a few hundred–few thousand steps, then cosine to ≈10% of peak; set T_total up front.
Scale peak LR down with width (≈1/√d), up with batch (up to the critical batch size).
Init ≈N(0, 0.02) or 1/√d, with residual projections downscaled by 1/√(2L).
Global-norm gradient clip at 1.0, always on.
bf16 compute + fp32 master weights & optimizer state; no loss scaling needed.
Checkpoint often so a spike costs hours, not the run; keep a rewind/skip-batch/lower-LR playbook ready.

Checkpoint

Try it

Open the widget and reproduce each failure deliberately. (1) Set warmup = 0 and push peak LR to ≈40e−3, then hit run a few times: confirm the verdict flips to DIVERGED in the first few steps. Now restore warmup = 40 at the same peak LR and confirm it converges — explain in one sentence why warmup saved it (what is v̂ doing over those first steps, and what is the LR?). (2) Restore a peak LR that converges, then toggle clip off and re-run several times: watch the spike count climb from ~0 into double digits as the heavy-tailed bursts that clipping had been absorbing now pass through. (3) With everything healthy, slide β₂ across its range and read it as the variance-smoothing knob — it sets how many recent gradients v̂ averages over (≈1/(1−β₂) steps). The toy's bursts are bounded, so the spike count barely moves; then re-read §2 to see why on a real LM (unbounded heavy-tailed bursts) the lower 0.95 is what damps spikes — the denominator reacts within ≈20 steps instead of ≈1,000. Then, by hand: for a 7B model, tally the mixed-precision AdamW bytes/param from §7 and confirm the ≈112 GB training-state figure that lesson 5 opens with.

Where this points next

You can now train a model that converges instead of diverging, and spend the budget without wasting it. But this lesson quietly assumed the run fits — that the weights, gradients, optimizer state, and activations all sit in GPU memory, and that D tokens finish in the GPU-hours you have. The preview in §7 already cracked that assumption: mixed-precision AdamW costs ≈16 bytes/param, so a 7B model's training state is ≈112 GB — and an H100 has 80. You cannot launch what you cannot fit, and you cannot promise a deadline you cannot estimate. Before spending a single GPU-hour you need the two back-of-envelope models — compute (the 6ND rule) and memory (the bytes-per-param ledger) — that tell you fit and runtime. That is 05 · Resource accounting — FLOPs and the memory wall, and its punchline (a 7B doesn't fit on one GPU) is what forces the entire systems half of the course.

Takeaway

A modern LM trains on one recipe, each piece forced by a concrete failure of the simpler choice. Plain SGD can't handle a Transformer's anisotropic, heavy-tailed gradients, so you use AdamW — per-coordinate normalization by the 1st/2nd moments (β₁=0.9, β₂=0.95, ε≈1e−8) with decoupled weight decay (λ≈0.1) so decay isn't distorted by the adaptive scaling. Adam's variance estimate is unreliable for the first hundreds of steps, so you linearly warm up the LR before cosine-decaying it to ≈10% of a peak that scales ∝1/√d with width and rises with batch (up to the critical batch size, past which extra batch wastes FLOPs). Init ≈N(0,0.02) with residual projections downscaled by 1/√(2L) keeps the stream calm at depth; global-norm clipping at 1.0 is the fuse that lets a run survive heavy-tailed spikes instead of dying on them. Compute runs in bf16 — chosen over fp16 for its full fp32 exponent range, which removes the need for loss scaling — over fp32 master weights and optimizer state. That fp32 state is 12 bytes/param (≈16 with the working copies), the number that, applied to a 7B model, overflows one GPU and hands the course to resource accounting.

Interview prompts

Why does plain SGD fail on a Transformer where Adam succeeds? (§1 — the loss is anisotropic and gradients are heavy-tailed; a single global step size is wrong for most coordinates and a rare gradient spike becomes a divergent step. Adam normalizes each coordinate by its own running gradient magnitude.)
What does the "W" in AdamW change, and why does it matter? (§2 — decoupled weight decay: −η·λ·θ applied to the weights directly, not folded into the gradient, so it isn't divided by √v̂ and every parameter decays at the same proportional rate regardless of gradient scale.)
Why warm up the learning rate, and why for only a few hundred steps? (§3 — Adam's v̂ is built from too few gradients early on, so the normalized step is unreliable exactly when the weights are fragile; once v̂ stabilizes (a few hundred steps) the peak LR is safe.)
Why bf16 over fp16 for LM training? (§7 — bf16 keeps fp32's 8 exponent bits → ≈1e38 dynamic range, so gradients don't under/overflow and no loss scaling is needed; it trades mantissa precision, which the fp32 master weights absorb.)
What is the critical batch size and why should you care about your budget? (§4 — below it, doubling batch ≈ halves steps to target loss (free speedup); above it, extra batch buys diminishing returns, so you burn ≈2× compute for ≈1× progress. It rises as the loss falls.)
A run spikes to loss 8 at step 50,000. What do you do, in order? (§8 — rewind to the last good checkpoint and skip the offending batches; if it recurs, lower peak LR / tighten clipping / lower β₂ and resume. Frequent checkpoints make this cost hours, not the run.)
Why is β₂ = 0.95 standard for LMs instead of 0.999? (§2 — 0.999 averages v over ≈1,000 gradients, so the denominator barely reacts to a gradient burst and the step stays huge; 0.95 inflates the denominator within ≈20 steps, automatically damping spikes.)
Why scale residual-projection init by 1/√(2L)? (§5 — with 2L residual branches the stream variance grows ∝L at step 0; downscaling each output projection by 1/√(2L) holds the variance ≈constant with depth, so a deep model starts as calm as a shallow one.)