The shape of a training step
A training step is four phases with very different cost profiles. Forward: run the model, produce a scalar loss, and — the part people forget — retain every intermediate tensor the backward pass will need. Backward: walk the graph in reverse, turning each upstream gradient into an input gradient and a parameter gradient. Reduce: accumulate or synchronize those parameter gradients. Update: hand them to the optimizer.
The arithmetic is lopsided in a predictable way. For a linear layer y = x W the forward is one matmul; the backward is two (∂L/∂x = (∂L/∂y) Wᵀ and ∂L/∂W = xᵀ (∂L/∂y)). That is the origin of the familiar rule of thumb: a training step costs about 3× a forward-only pass, or roughly 6 · P FLOPs per token for a model with P parameters (2 forward, 4 backward). Everything below is an attempt to keep that 3× from becoming 5× while staying inside a fixed memory budget.
Why activations, not weights, break your memory budget
Count what a training step holds resident. Weights: P values. Gradients: another P. AdamW state: two more moments, so 2P. With fp32 master weights you easily reach 16 bytes per parameter — a large but fixed number.
Activations are the variable one, and they scale with your data, not your model. Every tensor the backward pass will consume must survive from the moment it is produced until its gradient flows back through it — roughly O(B × N × d × L) for batch B, sequence length N, width d, and L layers, with a constant of maybe 10–20 tensors per block. A concrete case: B=8, N=2048, d=1024, L=24, bf16 (2 bytes), 16 saved tensors per layer gives 8 · 2048 · 1024 · 24 · 16 · 2 ≈ 12.9 GB — for a model whose weights are under a gigabyte. Double the sequence and it doubles. That asymmetry is why every technique below is really about activations.
Gradient accumulation: a batch-size simulator
Suppose the recipe calls for a batch of 512 sequences and your device holds 8. One observation resolves it: the loss is a mean over examples and differentiation is linear, so the gradient of the mean is the mean of the gradients. Run K = 64 micro-batches, add each one’s gradients into a persistent buffer, and step the optimizer once at the end.
for k in 1..K: # K micro-batches of size b
loss_k = forward(batch_k) / K # scale so the sum is a mean
backward(loss_k) # grads ACCUMULATE into .grad
optimizer.step(); optimizer.zero_grad()
effective_batch = K · b
peak activation memory ∝ b (NOT K · b)The /K is the step people get wrong: autograd’s .grad field adds by default, so without it you get a gradient K× too large and an effective learning rate K× too high. The accumulator costs one extra buffer of P values. What it does not buy is speed: K small matmuls are less efficient than one large one, so accumulation trades throughput for the ability to run a large-batch recipe at all.
Activation checkpointing: recompute instead of remember
If activations are the problem, keep fewer of them and regenerate the rest. Checkpointing saves activations only at chosen boundaries — typically each block’s input — and discards everything inside. When the backward pass reaches a block, it replays that block’s forward from the saved input, uses the regenerated intermediates for its VJPs, then frees them again.
The trade is clean. Checkpointing every block keeps only the block-boundary tensors — often a 5–10× reduction in practice — for one extra forward pass, so step cost goes from ~3× forward to ~4×: about +33% compute. Checkpoint at √L-spaced intervals instead and the classic result is O(√L) memory for the same single recompute. Choose boundaries deliberately: a block with an expensive forward but small saved tensors is a bad candidate, while attention blocks — large saved tensors, moderate recompute — are usually the best.
Precision: bf16 compute, fp32 accumulate, fp32 master weights
Mixed precision is not one decision but three, and conflating them causes most precision bugs. Storage: activations and weight copies in bf16 or fp16, halving memory and bandwidth. Accumulation: matmul inner products summed in fp32 inside the kernel, because adding thousands of low-precision terms loses bits fast. Master weights: an fp32 copy of the parameters that the optimizer updates.
The master copy exists because of a swallowing problem. bf16 carries about 8 bits of mantissa, so 1.0 + 0.0001 rounds straight back to 1.0: a typical update of relative size 10^-4 or smaller simply vanishes, and training silently stalls. Applying updates in fp32 and casting down for the next forward preserves them. The fp16-vs-bf16 split follows from range, not precision: fp16’s smallest normal is around 6×10^-5, so small gradients underflow to zero and you need loss scaling (multiply the loss by S ≈ 2^15, unscale before the step, skip steps whose gradients went non-finite). bf16 keeps fp32’s exponent range and needs no scaler — which is why it won.
Fused backward kernels: the bandwidth argument
A naive backward pass through bias-add, then GELU, then dropout launches three kernels. Each reads its input from main memory, does a few operations per element, and writes its output back. The arithmetic intensity is miserable — on the order of one FLOP per byte moved — so the hardware waits on memory while the arithmetic units idle.
Fusion merges those steps into one kernel that reads once, keeps intermediates in registers, and writes once, cutting memory traffic by roughly the number of ops fused. The wins concentrate in exactly the places you would guess: the elementwise chains after each matmul, layernorm/RMSNorm backward (reduction terms computed in one pass over the row rather than three), fused attention backward that never materializes the [N, N] score matrix, and fused optimizer steps. The principle: for memory-bound backward ops, the number of passes over the tensor matters more than the number of FLOPs. On a CPU the same logic holds with cache lines standing in for global memory.
Embedding gradients are sparse: gather forward, scatter-add backward
One layer breaks the dense pattern entirely. The embedding table is a matrix E: [V, d] with vocabulary V often 32k–256k, and its forward pass is not a matmul at all — it is a gather: y_t = E[id_t], a row lookup. The backward pass is the transpose of a gather, which is a scatter-add: each token’s upstream gradient is added into the row it read from, and untouched rows get exactly zero.
forward : y[t, :] = E[id_t, :] (gather)
backward: ∂L/∂E[v, :] = Σ_{t : id_t = v} ∂L/∂y[t, :] (scatter-add)
A batch of B·N = 16k tokens touches at most 16k distinct rows.
With V = 128k, at least 87% of ∂L/∂E is structurally zero.Two consequences follow. First, the add must be atomic or serialized — repeated tokens map to the same row, and a race there corrupts the gradient silently. Second, materializing a dense [V, d] gradient wastes hundreds of megabytes on zeros, which is why frameworks offer sparse gradient modes. The catch: AdamW’s moments are dense and its weight decay touches every row, so a sparse gradient with a stock AdamW quietly changes the update semantics.