The chain rule tells you what to compute in a backward pass. It says nothing about where the numbers live, and that second question decides whether your training run fits in memory or dies at step 40 with an OOM. Backprop’s formulas are settled mathematics; backprop’s architecture — the accumulator, the activation cache, the precision policy, the kernel fusion, the gradient reduction — is engineering, and it is where almost all real training failures come from. This article takes the gradient formulas as given (the companion piece derives them) and looks at the machine they run on: why activations rather than weights blow your memory budget, how gradient accumulation buys a large effective batch on small hardware, what checkpointing and fusion actually trade, and how these choices shift when the device is a CPU and the model is small.
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.
Reduction and clipping: between backward and step
Two operations sit between the last VJP and the optimizer. Under data parallelism each replica’s gradient is a partial mean and must be all-reduced into a single average; implementations overlap that communication with the tail of the backward pass, reducing each bucket as soon as its layer finishes. With accumulation, synchronize only on the final micro-batch or you pay K× the network cost for an identical result.
Then gradient clipping, which is global, not per-tensor. Compute g = √(Σ_p ||∂L/∂p||²) across every parameter, and if g > c scale all gradients by c / g. Clipping tensor-by-tensor changes the gradient’s direction, not just its magnitude — a different and worse algorithm. That norm is also your best diagnostic: log it every step, before clipping. A spike before the loss moves buys you a step of warning; a flat zero for some parameter group means gradient flow is broken upstream, not that the model converged.
Reading the failure modes
Backward-pass failures cluster into a few signatures, each pointing at a specific piece of the architecture above.
| Symptom | Likely cause | First move |
|---|---|---|
| OOM at step 1, forward was fine | Activation cache, not weights | Checkpoint; halve micro-batch |
| OOM only on long sequences | The [N, N] scores | Fused attention backward |
| Loss plateaus, weights barely move | Updates swallowed by bf16 | fp32 master weights |
| NaN under fp16 | Overflow or a bad loss scale | Switch to bf16 |
Grad norm K× too large | Missing /K | Scale the loss; zero the buffer |
| Nondeterministic gradients | Unordered scatter-add | Deterministic embedding kernel |
The pattern: when a training step fails, the fault is almost never in the gradient formulas — it is in the storage, the scaling, or the order of accumulation around them.
What changes on a CPU with a small model
Scale down to a small language model on a CPU and the priorities reorder. All-reduce disappears — one device, so the reduction is just the accumulator. Memory is comparatively plentiful (system RAM, not a fixed 24 GB card), which makes checkpointing much less attractive: you would pay a real 33% compute penalty on an already compute-starved device to solve a problem you may not have. Prefer a larger micro-batch and keep the activations.
What gets more important is bandwidth. CPUs have far less memory bandwidth per FLOP than accelerators, so the fusion argument sharpens: every avoided pass over an activation tensor is a direct win, and staying inside L2/L3 cache matters more than raw op count. Precision inverts too — bf16 arithmetic is only fast on CPUs with the right instruction support, so fp32 throughout is often the honest default and mixed precision can be a slowdown. Gradient accumulation, meanwhile, matters more than ever: it is how you reproduce a published large-batch recipe on a machine that holds a few sequences at a time.