The backward pass as a kernel chain

Parts I–III taught the forward pass: the GEMMs, the attention, the fused norms that turn tokens into logits. Training runs that chain backwards. This lesson is the hinge: what a backward pass actually is at the kernel level, why it costs ~2× the forward in FLOPs but is bottlenecked by memory, and what every op must stash on the way forward so its gradient kernel can run on the way back.

The question this lesson answers

You can write a fused forward kernel. Now someone hands you a model that has to learn. What new kernels appear, what do they cost, and where does the GPU actually choke? The surprising answer — and the reason training infra is its own discipline — is that the backward pass roughly doubles your FLOPs but can quadruple or worse your memory, and the memory is what you run out of first.

Autograd is a tape, not magic

A forward pass is a straight-line program: each op reads tensors, writes a tensor. Reverse-mode autodiff records that program as a tape — an ordered list of the ops and the inputs each one will need to compute its gradient. The backward pass walks the tape in reverse. For every forward op y = f(x) there is a vector-Jacobian product (vjp): given dy = ∂L/∂y flowing in from downstream, produce dx = ∂L/∂x = Jᵀ · dy to pass upstream.

The whole backward pass is just the forward kernel chain reversed, with each forward kernel replaced by its vjp kernel:

Two facts fall out of this picture and they drive everything in Part IV:

1. Backward needs forward's intermediates. The matmul vjp needs X and W; GELU's needs its pre-activation; attention's needs the softmax statistics. Those tensors are produced early in the forward and consumed late in the backward, so they must live in HBM for the whole step. That is the activation-memory peak.
2. The dependency edges reverse. A forward op whose output fed three consumers becomes a backward op whose gradient is the sum of three incoming gradients — gradient accumulation is a reduction, and reductions are lesson 07's template all over again.

The FLOP arithmetic: the 6N rule

One linear layer Y = X·W with X ∈ ℝ^T×K, W ∈ ℝ^K×N costs 2·T·K·N FLOPs forward (a multiply-add is 2). The backward produces two matmuls of the same size:

So backward is 2× the forward FLOPs, and a full training step is forward + backward ≈ 3× a forward-only (inference) pass. Counted per parameter per token this is the famous rule: ~2 FLOPs forward + ~4 backward = ~6 FLOPs per parameter per token. For an N-parameter model on D tokens, training is ≈ 6·N·D FLOPs — the number behind every compute-budget estimate you have seen.

But the bottleneck is memory, not those FLOPs

Here is the twist that makes training infra hard. The forward pass of a transformer block produces a stack of activations — the input to every matmul, every norm, every attention — and the backward needs almost all of them. They cannot be freed until their gradient kernel has run. So peak memory is dominated not by weights or gradients but by stored activations, and that grows linearly with batch×sequence and with depth.

The weights/grads/optimizer terms are fixed per model (the "16P+" that Adam mixed-precision training famously needs). The activation term is the one you control at the kernel level, and at long sequence it dwarfs the rest. Two kernel-side levers shrink it, and both are lessons in this Part:

• Recompute instead of save (lessons 28–30). If recomputing a tensor in the backward is cheaper than the HBM it would occupy, throw it away on the forward and recompute. FlashAttention does this with the N² score matrix; gradient checkpointing does it for whole blocks.
• Fuse so intermediates never hit HBM (lesson 27, and everything from Part III). The fused cross-entropy kernel is the extreme case: the [B·T, V] logits and their gradient are the single largest tensors in the model, and a fused kernel keeps them off HBM entirely.

What each forward op must save — the "save_for_backward" decision

Every op faces the same question: to run my vjp, do I save a forward tensor or recompute it? The answer is a bandwidth-vs-compute trade you can compute, and it is the recurring decision of this whole Part.

Three kernel-shaped facts about backward kernels

1. Layouts transpose. dX = dY·Wᵀ and dW = Xᵀ·dY touch transposed operands. The coalescing trap from lesson 04 and the tiling from lesson 05 reappear — a backward kernel that ignores the transpose tanks its bandwidth. Lesson 26.
2. Gradients accumulate, so backward has extra reductions. A weight used by every token gets a gradient summed over all B·T rows; a bias gradient is a column reduction. These are lesson 07 reductions, and doing them with cross-block atomics breaks determinism (lesson 30).
3. Accumulate gradients in fp32. A bf16 grad accumulator loses the low bits exactly when many small contributions sum — the same trap as the RMSNorm sum-of-squares in lesson 19. Every backward reduction upcasts.

Interactive · the training-step budget

Set the model and batch shape. The widget computes forward FLOPs, backward FLOPs (the 2× term), the fixed weight+grad+optimizer memory, and the activation peak — then shows how turning on full-block recompute trades the activation peak for extra backward FLOPs. Find the batch×seq where activations, not weights, decide whether the step fits.

What this sets up

You now have the map: training is the forward chain reversed, ~2× the FLOPs, bottlenecked by the activations the backward must consume, with every op making a save-vs-recompute call. The next five lessons build the kernels that fill in this map, starting with the atom they are all made of — the matmul backward, which is two GEMMs, not one.