Fused optimizers & the training step

Every gradient kernel in this Part produces one thing — gradients — and the optimizer consumes all of them. The optimizer is pure bandwidth, the launch-overhead problem at its worst, and the home of the mixed-precision dtype dance. This closing lesson builds the fused optimizer kernel, adds loss scaling and fp32 master weights, brings in recomputation as the last memory lever, and assembles the full training step back to lesson 25's peak.

Builds on

Consumes the gradients from lessons 26–29. Reuses the launch-overhead tax (lesson 18), the bandwidth roofline (lesson 10), and the recompute-for-memory trade introduced in lesson 25 and used in 27–29. This is the synthesis lesson of Part IV.

The question this lesson answers

You profile a training step and the optimizer — a few elementwise multiply-adds — eats a startling slice of the time, and the model barely fits in memory despite the weights being small. Why is the cheapest math in the model a performance problem, and where did all the memory go?

The optimizer is pure bandwidth

AdamW, per parameter, given gradient g:

m ← β₁·m + (1−β₁)·g v ← β₂·v + (1−β₂)·g² p ← p − lr·( m̂ / (√v̂ + ε) ) − lr·wd·p

Count the memory: to update one parameter you read p, g, m, v and write p, m, v — seven parameter-sized tensor touches for a handful of FLOPs. Arithmetic intensity is ~1 FLOP/byte: deeply memory-bound (lesson 10's roofline puts it far left of the ridge). There is no matmul to make fast and no reduction to tile; the only levers are "move fewer bytes" and "launch fewer kernels." That reframes the optimizer as a bandwidth kernel, not a math kernel.

The launch-overhead trap: thousands of tiny tensors

A transformer has thousands of separate parameter tensors (every weight matrix, norm scale, bias). A textbook optimizer loops in Python and launches one kernel per tensor. Each launch pays the 10–30 µs CPU-side dispatch tax from lesson 18, and each kernel moves so few bytes it finishes before the next is even queued — the GPU sits idle between launches.

Implementation	Launches	What dominates
Per-tensor (loop in Python)	~thousands	launch overhead — GPU starved between kernels
`foreach` / multi-tensor-apply	~tens (batched by dtype/shape group)	bandwidth — the right bottleneck
Fused (one kernel, whole Adam math)	~tens	bandwidth, minimal intermediates

Multi-tensor-apply (the foreach path) packs a list of tensors into one kernel launch: a metadata array of pointers+sizes lets a single grid stride across all tensors, so thousands of launches collapse to a handful. The fused optimizer goes further — it does the entire m/v/p update in one pass per element instead of several _foreach ops, so m and v never round-trip to HBM between sub-steps. Both turn a launch-bound kernel back into a clean bandwidth-bound one. This is the exact same "stop launching tiny kernels, fuse the elementwise chain" lesson as fused activation (lesson 19) — applied to the optimizer.

Mixed precision: the dtype dance

Matmuls want bf16/fp16 for tensor-core throughput (lesson 09), but the optimizer update is a tiny number added to a large weight. In bf16 (only 7 mantissa bits) a small lr·update added to a big p rounds away entirely — "swamping." So mixed-precision training keeps an fp32 master copy of the weights and fp32 optimizer states; the optimizer updates the fp32 master, then casts to bf16 for the next forward.

Loss scaling — and why bf16 mostly skips it

fp16 has a narrow exponent range; small gradients underflow to zero in its tail and the model stops learning. The fix is loss scaling: multiply the loss by S before backward so every gradient is scaled up by S into fp16's representable range, then unscale (divide by S) before the optimizer step. Dynamic loss scaling raises S until it sees an inf/NaN, then backs off. bf16 shares fp32's 8-bit exponent (same range, fewer mantissa bits), so it rarely underflows and usually needs no loss scaling — the single biggest practical reason modern training prefers bf16 over fp16.

Where all the memory went: the 16-bytes-per-parameter floor

Now close lesson 25's table. Mixed-precision Adam holds, per parameter:

Buffer	dtype	bytes/param
weight (compute copy)	bf16	2
gradient	bf16 (or fp32)	2 (–4)
fp32 master weight	fp32	4
Adam m	fp32	4
Adam v	fp32	4
total (the famous floor)		16 (–18)

So a 7B model needs ~112 GB before a single activation — which is why it doesn't fit on one 80 GB H100 and why optimizer-state sharding (ZeRO/FSDP) exists. The weights you can name are a quarter of the real footprint; the optimizer states are the rest. This is the fixed term in lesson 25's widget.

The last lever: activation recomputation (gradient checkpointing)

The fixed 16P is sharded away across GPUs; the activation term is what a single GPU still chokes on, and it's the term you control with kernels. FlashAttention already recomputes its scores (lesson 29); gradient checkpointing generalizes that to whole transformer blocks: on the forward, store only each block's input (the boundary), discard everything inside; on the backward, recompute the block's forward to regenerate its internal activations, then run the block's backward. The trade:

Memory: activations drop from "every layer's internals" to "one block's worth at a time" — for L layers checkpointed at block boundaries, roughly an L-fold cut in stored activations.
Compute: one extra forward per checkpointed block ≈ +33% step FLOPs (recall the 6N rule: forward is ~1/3 of fwd+bwd; recomputing it adds that third).

It's the same compute-for-memory bargain as every other kernel in Part IV, now at the granularity of a whole block. You reach for it exactly when lesson 25's widget says you're activation-bound.

The full training step, assembled

Putting Part IV together, one step is:

Forward (Parts I–II kernels) — save what each op needs (lesson 25), recompute the rest.
Loss — fused cross-entropy keeps the giant logits off HBM (lesson 27).
Backward — matmul's two GEMMs (26), norm/activation gradients with cross-token reductions (28), FlashAttention backward's recompute (29); all accumulating gradients in fp32; recompute checkpointed blocks here.
Optimizer — one fused multi-tensor Adam pass over fp32 master + states, loss-scale unscale, cast to bf16 (this lesson).

Peak memory = (16P optimizer floor, shardable) + (activations, controllable by fusion + recompute). Every kernel in this Part attacks one of those terms. That is the engineering line of thinking the whole track has been building toward: you don't memorize kernels, you locate the bottleneck term and apply the kernel that shrinks it.

Interactive · optimizer time and the memory floor

Set parameter count, the number of separate tensors, and the per-launch overhead. The widget compares per-tensor vs fused-foreach optimizer time (launch-bound vs bandwidth-bound) and shows the 16P memory floor that fusion does not change.

Where this leaves you

That closes Part IV: you can read or write the gradient kernel of every op in a transformer, predict its bottleneck (bandwidth, reduction, recompute, or launch), and assemble a training step that fits. Part V (lessons 31–33) drills the Part I CUDA primitives under interview conditions — the declarations, the optimized snippets, and the algorithm-and-data-structure thinking you reach for on a whiteboard.