Part V — The hard parts

Autodiff as a compiler pass

Lesson 11 taught the compiler to survive variable shapes and data-dependent control flow, so it can now capture, fuse, plan, and tune a real serving graph that changes size every request. But everything up to here has compiled the forward graph — inference. Training needs the backward graph too, and that backward must match a forward you have already fused, re-laid-out, and rewritten beyond recognition. Hand-writing backward kernels for an optimized graph is hopeless: the optimizer is allowed to delete, merge, and reorder ops you never see. The fix is the cleanest idea in the track — make differentiation itself a graph-to-graph pass, so the compiler builds the backward and then optimizes it with the exact same machinery it used on the forward.

1 · Reverse-mode AD is a graph transform, not runtime magic

Automatic differentiation (AD) computes exact derivatives of a program by composing the known derivatives of its primitive ops via the chain rule — not symbolically (no expression blowup) and not numerically (no finite-difference error). For deep learning we want the gradient of one scalar loss with respect to millions of parameters, so we use reverse mode: one forward sweep to compute values, one backward sweep that propagates the gradient of the loss back through every op. Reverse mode costs ≈1 backward for all inputs at once, versus forward mode's one pass per input — for θ ∈ ℝ⁷ᴮ → loss ∈ ℝ, that is the difference between two sweeps and seven billion.

The atom of reverse mode is the vjp — the vector-Jacobian product. Every primitive op y = f(x) ships a rule that, given the incoming gradient ḡ = ∂loss/∂y, returns ∂loss/∂x = Jᶠᵀ · ḡ without ever forming the Jacobian J. The vjp is a small graph of ops, so differentiation is just substitution:

To differentiate a whole graph, walk it in reverse topological order: seed the output with ∂loss/∂loss = 1, then for each op apply its vjp to the gradient flowing back into its outputs, summing contributions where a value fanned out to multiple consumers. The result is a brand-new subgraph of vjp ops — the backward graph. Critically, this transformation happens at compile time on the IR: you give it a forward graph, it hands back a forward+backward graph. There is no Python running, no objects, no interpreter in the loop.

Contrast eager PyTorch. Eager autograd is a runtime tape: as each op executes, it pushes a backward closure (a Python/C++ object holding saved tensors) onto a graph, and loss.backward() later walks that tape calling each closure. That works, but the tape is built and torn down every iteration, it is opaque to any graph optimizer, and the saved tensors it pins are whatever the op author chose — there is no pass that can fuse a backward op into a forward one or trade memory for recompute, because there is no graph to rewrite. Turning AD into a compile-time graph transform is exactly what unlocks the rest of this lesson: once the backward is in the IR, fusion, planning, and layout treat it like any other ops.

eager autograd                          AD as a compiler pass
--------------                          ---------------------
run op  → push backward closure         forward graph  (IR)
run op  → push backward closure                 |  apply vjp per op, reverse sweep
   ...     (a runtime tape)                      v
loss.backward(): walk tape, call         forward + backward graph  (IR)
  each closure in reverse                        |  fuse / plan / layout / tune
                                                  v
opaque to the optimizer                  one compiled artifact, backward optimized too

2 · AOTAutograd: trace once, differentiate, partition, compile both halves

In torch.compile, the pass that does this is AOTAutograd (ahead-of-time autograd). After Dynamo (lesson 02) captures a forward FX graph, AOTAutograd takes over and runs a fixed pipeline that produces two compiled functions — one for the forward, one for the backward — wired together by a small set of saved tensors:

The payoff is that the backward is a first-class graph. When Inductor fuses bias+gelu into the matmul epilogue in the forward, it independently fuses the corresponding gelu′ ⊙ ȳ and Wᵀ-matmul cluster in the backward. The reason hand-writing backward kernels for an optimized forward is hopeless — the forward op boundaries are gone — is exactly why generating it works: the compiler differentiates the clean graph first, then optimizes forward and backward in lockstep. JAX takes the same stance with grad/vjp as program transformations over jaxprs, lowered through XLA.

3 · Functionalization: kill mutation and aliasing so the passes are legal

Reverse-mode AD and every rewrite from lessons 03–06 share one precondition: value semantics. A vjp rule for add assumes its inputs still hold their forward values when the backward runs; CSE assumes that two reads of the same value return the same thing; buffer reuse assumes a value is dead after its last read. In-place mutation and view aliasing violate all three. If x.add_(1) overwrites x, then a later op that wanted the old x for its vjp silently gets the wrong value. If y = x[2:] is a view and someone writes through y, x changes under the rewriter's feet.

Functionalization is the pass that removes this. It rewrites every mutating or aliasing op into a pure op that returns a new value, threading the new value forward to all subsequent uses:

This is the lesson 03 SSA motif paying back exactly when promised: lesson 03 argued value semantics make rewriting local and sound and flagged that mutation would have to be removed before AD. Functionalization is that removal. After it runs, the joint graph is pure SSA, so AD's vjp substitution is correct and fusion/CSE/DCE/planning all apply unchanged. The in-place ops are not gone forever — memory planning (06) reintroduces safe in-place updates later, but only where it can prove via liveness that no other consumer needs the old value. The discipline is: be pure for analysis, be in-place for execution, and let the compiler decide which is provably safe.

4 · The partitioner: save-vs-recompute as a min-cut

The backward needs forward intermediates — the matmul vjp needs x and W, gelu's vjp needs its input. The naive choice is to save every forward activation in HBM and read it back in the backward. That is the activation-memory wall: for a deep model the saved activations dwarf the parameters, and you OOM (the same ledger as cs336 · 05 accounting). The opposite extreme — recompute everything, saving only the inputs and regenerating each activation from scratch in the backward — costs almost a second full forward pass in FLOPs. Neither extreme is right; the answer is a per-activation choice, and the compiler makes it.

The partitioner formalizes it as a graph min-cut. Picture the joint graph with the forward on one side and the backward on the other; every forward value the backward consumes is an edge crossing the boundary. Each crossing edge can be paid for two ways: save it (cost = its bytes in HBM, carried across the cut) or recompute it (cost = the FLOPs of the ops that regenerate it from things already on the backward side). The partitioner picks the cut — the set of values to save — that minimizes total cost under the memory budget. Rematerialization is the general name for recompute-instead-of-store; activation checkpointing is the manual version (cs336 · 05), and the compiler partitioner is the automatic one.

The heuristics the real min_cut_rematerialization_partition uses are intuitive once you see the cost model. Cheap-and-big → recompute: elementwise ops (gelu, dropout mask, bias) are nearly free in FLOPs but their outputs are large, so regenerating them is far cheaper than the bytes. Expensive-and-shared → save: a matmul output is costly to recompute and often feeds several backward ops, so save it once. This is precisely FlashAttention's trade made by a compiler instead of by hand: cs336 · 07 kernels derives why attention recomputes the T×T scores in the backward rather than storing the O(T²) matrix — cheap FLOPs to save scarce HBM. The partitioner generalizes that single hand-tuned decision to every activation in the graph, automatically.

5 · The same middle-end runs on the joint graph

The quiet, important fact: once AD has produced the joint graph and the partitioner has split it, nothing downstream needs to know it is a backward. Fusion (05) sees a chain of memory-bound elementwise ops and fuses them whether they compute activations or gradients. Memory planning (06) computes liveness and reuses buffers across forward and backward alike — and the partitioner's save/recompute choice is literally input to that liveness analysis (a recomputed value is born late and dies fast; a saved value lives across the whole backward). Layout (07) propagates the same tensor-core-friendly formats. This is the reuse dividend of treating AD as a graph transform: you wrote the optimization passes once, for forward inference, and they pay off again, unchanged, for training. The compiler did not learn calculus; it learned to rewrite a graph, and gradients are just more graph.

6 · Drive it: the recompute partitioner

The widget is a small forward chain — x → matmul → gelu → dropout → matmul → norm → loss. Each box is an activation the backward will need; click it to toggle save (kept in HBM, green) vs recompute (regenerated in the backward, blue). The KPIs are the three numbers the partitioner trades: saved-activation memory, extra backward FLOPs from recompute, and peak memory (saved activations + the working set). The dashed budget is one device. Try the presets: save everything minimizes FLOPs but blows past the memory budget (OOM); recompute everything fits easily but pays nearly a second forward in FLOPs; the min-cut preset saves the expensive matmul outputs and recomputes the cheap-but-big elementwise ones — fitting the budget at a tiny FLOPs premium.

The shape of the trade is the lesson: peak memory and extra FLOPs pull in opposite directions, and the partitioner finds the cut that fits the budget for the fewest extra FLOPs. Save-everything is the eager-PyTorch default that OOMs at scale; recompute-everything is gradient checkpointing taken to its wasteful extreme; the min-cut is what the compiler picks for you — and it picks it again, automatically, every time shapes change.

Failure modes & checklist

Checkpoint

Where this points next

You now have the whole training graph compiled: AD generated the backward as a graph transform, functionalization made it sound, the partitioner balanced memory against recompute, and the same middle-end fused and planned both halves. But every number in that pipeline assumed one device. The partitioner fought the activation-memory wall on a single GPU and won some headroom — yet a frontier model's parameters alone (tens to hundreds of GB), plus optimizer state, plus the activations the partitioner couldn't recompute away, do not fit on one 80 GB device, and even when they fit, one device can't finish the step in acceptable time. The min-cut bought slack; it did not buy a second order of magnitude. The only move left is to spread the graph — parameters, activations, and the batch — across many devices, which means partitioning and inserting communication becomes another compiler pass. That is lesson 13, Distributed compilation.

Interview prompts

• Why generate the backward graph instead of hand-writing backward kernels for a compiled model? (§intro/2 — fusion/layout/planning have dissolved the forward's op boundaries, so there is nothing for a human to differentiate; AD runs on the clean graph first, then the same passes optimize the backward.)
• What is a vjp and why does reverse mode use it? (§1 — the vector-Jacobian product Jᵀ·ḡ computed without forming J; reverse mode gives the gradient of one scalar loss w.r.t. all inputs in a single backward sweep, vs forward mode's one pass per input.)
• Contrast eager autograd's tape with AD as a compiler pass. (§1 — the tape is a runtime list of backward closures, rebuilt every step and opaque to optimizers; the pass produces a backward in the IR at compile time, so fusion/planning/layout can rewrite it.)
• What does functionalization do and why is it required before AD? (§3 — rewrites in-place and aliasing ops into pure value-semantics ops; AD's vjp rules and the 03–06 rewrites all assume a value is defined once and still available, which mutation breaks.)
• Explain the save-vs-recompute partition as a min-cut. (§4 — each forward value the backward needs is an edge across the fwd/bwd boundary, payable as its HBM bytes (save) or its recompute FLOPs (recompute); the partitioner picks the cut minimizing total cost under the memory budget — rematerialization.)
• Which activations should you save and which should you recompute? (§4 — recompute cheap-and-big (elementwise: gelu, dropout, bias); save expensive-and-shared (matmul/attention outputs). It's FlashAttention's hand-tuned trade generalized to every op.)
• Why do the forward optimization passes work on the backward unchanged? (§5 — once AD emits a joint graph, the backward is just more typed SSA ops; fusion/liveness/layout don't care whether a value is an activation or a gradient, so the passes written for inference pay off again for training.)