What Autograd Actually Does

One number, a hundred billion gradients

Training a language model ends, every single step, with one number. A batch of text goes in, a hundred billion parameters do their work, and out drops a scalar — the loss. Then the optimiser asks a question that ought to be absurd:

Every one of those hundred billion parameters wants to know how much it personally contributed to that one number. And we need the answer within roughly the same amount of compute it took to produce the number in the first place, or training does not happen at all.

Consider what the obvious method would cost. Finite differences: nudge one parameter, run the whole model again, see how the loss moved.

That is two forward passes per parameter. For a 100B-parameter model at roughly \(2N\) FLOPs per token per forward pass, you would need something like \(4 \times 10^{22}\) FLOPs to differentiate a single training example. An H100 at 1000 TFLOP/s would take about 1.3 million years. And the answer would be numerically garbage besides, because subtracting two nearly equal floats in bf16 destroys most of the precision you were trying to measure.

The remarkable fact — the fact this entire article exists to explain — is that the right method costs roughly twice the forward pass. Not twice per parameter. Twice, total. All hundred billion derivatives, for the price of a couple of extra function evaluations.

The plan for this article: build the machinery from a single scalar chain rule, watch it generalise to graphs and then to Jacobians, discover reverse mode by asking a question about matrix multiplication order, derive the backward pass of every layer in a transformer, and then follow the consequences into memory — because the price of reverse mode is not compute, it is that you have to remember the forward pass.

The chain rule, taken literally

Start where everyone starts, but stay there longer than usual, because the whole edifice is this one identity applied with discipline.

Most people learn this as a symbol-manipulation rule. It is more useful as a statement about sensitivity, so let me state it that way.

Derivatives are exchange rates

Suppose \(v = h(x)\) and \(L = g(v)\). Wiggle \(x\) by a tiny amount \(\varepsilon\). Then \(v\) moves by approximately \(h'(x)\varepsilon\) — the derivative is the exchange rate between a perturbation in \(x\) and a perturbation in \(v\). Now that perturbation in \(v\) causes a perturbation in \(L\) of about \(g'(v) \cdot (h'(x)\varepsilon)\), because \(g'\) is the exchange rate from \(v\) to \(L\).

Converting dollars to euros to yen: you multiply the rates. The chain rule is currency conversion, and a deep network is a very long chain of currency conversions. Nothing more complicated than that is ever going to happen in this article — everything else is bookkeeping about the order you multiply the rates in.

This framing pays for itself immediately, because it tells you what a vanishing gradient is. If every rate in a chain of fifty conversions is 0.9, the end-to-end rate is \(0.9^{50} \approx 0.005\). A dollar's worth of perturbation at the input becomes half a cent at the output. The chain has not broken; it has simply been multiplied into irrelevance.

Two things that go wrong the moment you leave the textbook

The single-variable chain rule is not quite enough, for two reasons that both show up in the first layer of any real model.

A variable can influence the loss through several routes. In a residual block, \(x\) reaches the output both through the sublayer and around it. The rule that covers this is the multivariable chain rule, which says the routes add:

summed over every \(u_k\) that \(x\) feeds directly. This is not a technicality. It is the reason a skip connection helps: it adds a route whose exchange rate is exactly 1, and a sum containing a 1 cannot vanish no matter what the other terms do.

The functions are not scalar. Real layers map \(\mathbb{R}^{n} \to \mathbb{R}^{m}\), so the "exchange rate" becomes a matrix. That generalisation is where all the interesting cost structure lives, and I will get to it after the graph machinery is in place — because the graph is what tells you which matrices to multiply, and the shapes are what tell you in which order.

Everything is a graph

Here is the move that turns calculus into an algorithm. Any composite function — any program, really — can be flattened into a sequence of elementary operations, each of whose derivative you already know. That sequence is a directed acyclic graph. Nodes are intermediate values; edges are dependencies.

Take a deliberately tiny example, small enough to hold in your head and rich enough to contain every phenomenon that matters:

That is a one-unit network with a residual connection and a squared-error loss. Four operations, six leaves, one scalar out. Notice that \(x\) appears twice — once inside the linear term and once in the skip. That fan-out is the interesting part.

The forward pass is just evaluation with note-taking

Running the graph left to right is nothing but computing the function. The only thing autograd adds is that it writes down what it did: which operation produced each node, and which inputs that operation will need again when the derivative is requested. That record — the tape, the graph, the trace, depending on whose framework you are reading — is the entire data structure.

Two design philosophies exist for building it. Define-and-run (Theano, TensorFlow 1) constructs a static graph once and then executes it many times, which permits aggressive whole-graph optimisation but makes control flow that depends on data awkward. Define-by-run (autograd, Chainer, PyTorch) builds the graph as a side effect of executing the forward pass, so the graph is whatever the Python actually did on this particular batch. Loops with data-dependent trip counts, branches on tensor values, early exits — all free. The price is that there is no graph to optimise until after you have run it, which is exactly the gap that torch.compile and JAX's tracing exist to close.

Local derivatives and the adjoint

Give every node \(v\) in the graph a shadow value:

This is the adjoint of \(v\) — how much the final loss moves per unit move in \(v\), holding the graph structure fixed. It is the thing backpropagation computes, and it is worth being precise that \(\bar v\) is a property of the whole graph downstream of \(v\), not of \(v\) itself.

Seed the output: \(\bar{L} = \partial L/\partial L = 1\). Then work backwards. If a node \(v\) feeds into children \(u_1, \ldots, u_k\), the multivariable chain rule gives the backward recurrence:

Read that carefully, because it contains the entire algorithm and one profound engineering consequence.

The term \(\partial u_k / \partial v\) is a local derivative. It depends only on the operation that produced \(u_k\) and the values that went into it. The tanh node knows how to differentiate a tanh; it knows nothing about the loss, the optimiser, the layers above it, or whether it is inside a transformer or a physics simulator. The term \(\bar{u}_k\) is everything global, and it arrives from outside, already computed, as a number.

The algorithm, in full:

# forward: evaluate in topological order, recording the trace
for node in topological_order(graph):
    node.value = node.op(*[p.value for p in node.parents])
    node.saved = node.op.what_backward_will_need(...)

# backward: traverse in reverse topological order, accumulating
adjoint = defaultdict(zeros_like)
adjoint[loss] = ones_like(loss)

for node in reversed(topological_order(graph)):
    g = adjoint[node]                      # the incoming adjoint, already summed
    for parent, local in node.op.backward(g, node.saved):
        adjoint[parent] += local           # += is the multivariable chain rule

Three details in that sketch are doing real work.

Reverse topological order is not optional. A node's adjoint is only complete once every child has contributed. Visit a node too early and you will propagate a partial sum, silently, with no error. This is why frameworks maintain a dependency counter per node and only enqueue a node once all its outgoing edges have reported in.

The accumulation is +=, not =. That plus-equals is the multivariable chain rule. It is also, as we will see, the reason .grad accumulates across backward calls rather than overwriting, which surprises people exactly once.

Nothing here is specific to neural networks. This procedure differentiates any program built from differentiable primitives — ODE solvers, ray tracers, physics engines, other people's gradients. Backpropagation is the special case where the program happens to be a neural network.

Forward accumulation, reverse accumulation

The recurrence above walks the graph backwards. There is a perfectly good recurrence that walks it forwards, and comparing them is the cleanest way to understand why deep learning does the thing it does.

Forward mode: pick an input, push

Fix one input variable \(x_j\). Attach to every node a tangent:

Seed \(\dot{x}_j = 1\) and \(\dot{x}_i = 0\) for \(i \ne j\), then sweep forward in topological order:

The tangents can be computed alongside the values in a single pass — this is exactly dual-number arithmetic, where you carry \((v, \dot v)\) as a pair and every primitive knows how to propagate both. When the sweep finishes, \(\dot L = \partial L / \partial x_j\). One derivative. To get all \(N\) of them you run the sweep \(N\) times.

Reverse mode: pick an output, pull

Fix one output. Seed \(\bar L = 1\) and sweep backwards with the adjoint recurrence. When the sweep finishes, every node carries \(\partial L / \partial v\) — including every leaf. All \(N\) derivatives. One sweep.

The Jacobian view makes the difference obvious

For \(f : \mathbb{R}^{n} \to \mathbb{R}^{m}\), the Jacobian \(J \in \mathbb{R}^{m \times n}\) is the whole derivative. Forward mode, seeded with the basis vector \(e_j\), computes \(J e_j\) — the \(j\)-th column. Reverse mode, seeded with \(e_i\), computes \(e_i^\top J\) — the \(i\)-th row.

So the question "which mode?" is really the question "is your Jacobian tall or wide?"

Training a neural network is the first row, taken to an extreme. \(m = 1\) exactly, and \(n\) is in the hundreds of billions. That ratio, \(n/m \approx 10^{11}\), is the largest asymmetry anyone routinely exploits in numerical computing.

When the wires become buses

Real layers move vectors, not scalars. So promote everything. If \(u = f(v)\) with \(v \in \mathbb{R}^{n}\) and \(u \in \mathbb{R}^{m}\), the local derivative is a matrix:

and the chain rule becomes matrix multiplication. For a network of \(k\) layers, \(L = f_k \circ f_{k-1} \circ \cdots \circ f_1\), the derivative of the whole thing is the product of the layer Jacobians:

Matrix multiplication is associative. The value of that product does not depend on how you parenthesise it. The cost depends on it enormously.

The matrix chain problem, which is the whole thing in disguise

Multiplying an \(a \times b\) matrix by a \(b \times c\) matrix costs \(abc\) multiply-adds. Take a three-layer chain with widths \(n_0 \to n_1 \to n_2 \to n_3\):

Now set \(n_3 = 1\), because the last thing a network computes is a scalar loss. The first ordering costs \(n_2 n_1 + n_1 n_0\). The second costs \(n_2 n_1 n_0 + n_2 n_0\). With \(n_0 = n_1 = n_2 = 4096\), that is 33.6 million operations against 68.7 billion — a factor of about two thousand, from nothing but the placement of parentheses.

Grouping from the output end means starting with the small dimension and never letting the intermediate result grow. That grouping is reverse mode. Grouping from the input end is forward mode. Backpropagation is, at bottom, the observation that when one end of your matrix chain is a scalar, you should start multiplying from that end.

Nobody ever builds these matrices

A word of warning before the next section, because it is the single most common confusion about backpropagation. The Jacobian of a linear layer mapping \(\mathbb{R}^{4096} \to \mathbb{R}^{4096}\) over a batch of 8192 tokens is, formally, a matrix of shape \((8192 \cdot 4096) \times (8192 \cdot 4096)\). That is \(1.1 \times 10^{15}\) entries, four petabytes in fp32.

It is never constructed. It never exists. Not in a compressed form, not in a clever layout — it is simply never a thing. What backpropagation actually manipulates is the action of that matrix on a vector, and the action is cheap even when the matrix is absurd. That distinction is the subject of the next section, and once it clicks, autograd stops being mysterious.

JVP and VJP

Two products. Everything in every autodiff system is one of these two.

The Jacobian-vector product: pushing forward

Take a direction in input space, get the direction it induces in output space. This is the derivative along a direction, computed without ever forming \(J\). Geometrically it is a pushforward: you have a tangent vector attached to a point in the input manifold and you carry it forward through the map.

The vector-Jacobian product: pulling back

Take a covector in output space — a sensitivity, a "how much do I care about each output" — and get the corresponding sensitivity in input space. Geometrically it is a pullback, and pullback is the natural direction for covectors, which is a rather elegant way of saying that gradients live in the dual space and always have.

The adjoint recurrence from earlier is exactly a VJP:

So the backward pass of a layer is a function with the signature (saved_forward_values, upstream_gradient) -> downstream_gradient, and it is required to compute \(J^\top \bar u\) without ever mentioning \(J\).

Concretely: the linear layer

Let \(u = Wv\) with \(W \in \mathbb{R}^{m \times n}\). Then \(J = W\), and:

Both cost \(mn\) multiply-adds, the same as the forward pass. The Jacobian was already sitting there as the weight matrix, so this case looks trivial — but it makes the general shape of the trick visible. For a tanh, \(J\) is diagonal and the VJP is an elementwise multiply. For a reshape, \(J\) is a permutation and the VJP is a reshape back. For a softmax, \(J\) is dense but has a rank-one-plus-diagonal structure that collapses into two cheap operations. In every case the structure of \(J\) is known analytically, and the VJP exploits it. Building the matrix would throw that structure away and then pay to multiply through it.

Why reverse mode wins, precisely

Assemble the pieces into a cost statement. Let \(\mathrm{ops}(f)\) be the cost of one forward evaluation.

with small constants \(c_f, c_r\) that depend on the primitive mix. For a scalar loss, \(m = 1\), and the entire gradient costs a small constant multiple of one forward pass, independent of \(n\). This is the cheap gradient principle, and it is a theorem, not folklore: Baur and Strassen proved in 1983 that the nonscalar arithmetic complexity of computing a rational function together with all of its first partial derivatives is at most three times that of computing the function alone, with a bound near five when all operations are counted. The constant does not grow with the number of variables. That single fact is why gradient-based optimisation of hundred-billion-parameter models is possible at all.

What it costs instead

Nothing is free. Forward mode can compute the value and the tangent in lockstep and throw away each intermediate as soon as it is consumed — its memory footprint is a constant multiple of the forward pass. Reverse mode cannot. It needs \(\partial u_k / \partial v\) evaluated at the forward values, and it needs them in the reverse order they were produced. So it must keep them.

You buy a factor of \(10^{11}\) in compute and you pay for it in storage. That trade is the source of every activation-memory problem in large-model training, and the second half of this article is about what people do about it.

Why the backward pass costs about twice the forward pass

The rule of thumb everyone quotes — forward is \(2N\) FLOPs per token, backward is \(4N\), total \(6N\) — falls out of the linear layer. For \(U = XW\) with \(X \in \mathbb{R}^{B \times n}\), \(W \in \mathbb{R}^{n \times m}\):

One matmul forward, two backward, all the same size — one to send the gradient further down the network, one to land it on the weights. Everything else in a transformer is elementwise or small by comparison, so the 1:2 ratio survives into the whole-model estimate. If a layer has no parameters, its backward is a single matmul and the ratio is 1:1; if you do not need the input gradient, because the layer is the first one, you can skip \(\bar X\) entirely.

The backward pass of every layer you use

Time to be concrete. Below is the VJP for each primitive in a transformer, derived rather than asserted, with an animated figure you can step through operation by operation. Two things to watch for in every case: what the backward pass needs saved from the forward pass, and what the structure of the Jacobian lets you avoid doing.

Linear

The workhorse. \(U = XW\), with \(X \in \mathbb{R}^{B \times n}\), \(W \in \mathbb{R}^{n \times m}\).

Derive \(\bar W\) directly and the shape argument does the work for you: \(U_{ij} = \sum_k X_{ik} W_{kj}\), so \(\partial U_{ij} / \partial W_{k\ell} = X_{ik}\delta_{j\ell}\), and summing over the output indices gives \(\bar{W}_{k\ell} = \sum_i X_{ik}\bar{U}_{i\ell}\), which is \(X^\top \bar U\). If you ever forget the formula, the only arrangement of \(X\), \(\bar U\) and their transposes that produces something of shape \(n \times m\) is the right one. This is a legitimate derivation technique and I use it constantly.

Saved: \(X\) and \(W\). Note that \(X\) is the activation — batch-sized, and therefore the expensive one. \(W\) you already have.

ReLU

The Jacobian is diagonal, so the VJP is an elementwise multiply and a matrix that formally has \(n^2\) entries costs \(n\) operations to apply. The gradient is gated by which units fired. Units that were off receive nothing — which is the dying-ReLU problem stated in one line, since a unit that is off for every example in the batch gets exactly zero gradient forever.

Saved: only the sign pattern. In principle one bit per element; in practice frameworks usually keep the output tensor, which they had anyway.

At \(x = 0\) the derivative does not exist. Everyone picks a convention — PyTorch returns 0 — and the measure-zero handwave is fine in floating point, though it does mean that "the gradient" of a ReLU network is really a choice of element from a set. There is a rigorous theory of this (conservative gradients), and it mostly confirms that the pragmatic answer was fine.

GELU

with \(\Phi\) the standard normal CDF and \(\phi\) its density. Still diagonal, so still an elementwise multiply, but with a real cost difference: the backward pass needs \(\Phi(x)\) and \(\phi(x)\), which means either saving \(x\) and recomputing two transcendentals, or saving both. The derivative is non-monotonic and dips slightly below zero around \(x \approx -0.75\), which is the smooth version of ReLU's hard gate and the reason GELU passes a little gradient to units that were nearly off. The tanh approximation used in most implementations differentiates just as cleanly and is what fused kernels actually run.

LayerNorm: the backward pass is a projection

This is the one worth doing slowly, because the answer is geometrically beautiful and the naive derivation is a mess.

The subtlety is that \(\mu\) and \(\sigma\) both depend on every \(x_i\), so every output depends on every input and the Jacobian is dense. Writing \(g = \bar{y} \odot \gamma\) for the gradient with respect to \(\hat x\), the result collapses to:

and the parameter gradients, summed over the batch, are \(\bar\gamma = \sum \bar y \odot \hat x\) and \(\bar\beta = \sum \bar y\).

Now read the three terms. The first is what you would get if \(\mu\) and \(\sigma\) were constants. The second subtracts the component of \(g\) along the all-ones direction. The third subtracts the component along \(\hat{x}\). Together they are an orthogonal projection: LayerNorm's backward pass removes exactly the two directions its forward pass made the output insensitive to. Shifting \(x\) by a constant does not change \(y\), so the gradient must be orthogonal to \(\mathbf{1}\). Scaling \(x\) does not change \(y\), so the gradient must be orthogonal to \(\hat x\). The algebra is just enforcing a symmetry that was in the forward pass all along.

That is a general principle worth stating on its own: an invariance in the forward pass appears as a projection in the backward pass. Wherever a layer ignores a direction, its gradient annihilates that direction.

Saved: \(\hat x\) (or \(x\) and \(\sigma\)), plus \(\gamma\). Fused kernels recompute \(\hat x\) from \(x\) and the saved statistics rather than storing both.

RMSNorm

The mean-subtraction term is gone, because RMSNorm never subtracted a mean and therefore is not shift-invariant. One projection instead of two. That is the real reason RMSNorm is cheaper — not the missing subtraction in the forward pass, which is trivial, but one fewer reduction in the backward pass, where reductions are the expensive part.

Residual: the gradient highway

The identity term is the whole point. Stack \(k\) residual blocks and the end-to-end Jacobian expands into \(\prod_{i}(I + J_i)\), which contains a term that is exactly \(I\) — a route from the loss to the input along which the gradient is multiplied by 1, \(k\) times in a row. Multiplying a hundred numbers around 0.9 gives you \(10^{-5}\); multiplying a hundred numbers that are exactly 1 gives you 1. The interactive graph earlier in this article shows the same thing in miniature: 1.93 of the 2.13 total gradient at \(x\) arrives along the skip.

Attention

Write the forward pass as three steps and differentiate each:

Count them: two matmuls forward, five backward. That is why attention's backward pass is disproportionately expensive, and it is exactly the count FlashAttention's authors give as the reason the backward kernel is harder to write than the forward one — more live values must be resident in SRAM at once.

The memory story is worse than the compute story. \(P\) is \(N \times N\) per head. At \(N = 8192\) with 32 heads in bf16, that is 8.6 GB of attention probabilities per layer, held from the forward pass until the backward pass reaches it. This is the single line item that made long-context training impossible before people started recomputing it, which is where the last third of this article is going.

Softmax, and the cancellation that saves it

Softmax deserves its own section, because its backward pass contains the most satisfying cancellation in the field and the most important piece of numerical hygiene in the whole training loop.

The Jacobian

With \(p_i = e^{z_i} / \sum_k e^{z_k}\), differentiate directly. For \(i = j\), the quotient rule gives \(p_i(1 - p_i)\); for \(i \ne j\), it gives \(-p_i p_j\). Both cases fit one expression:

A diagonal matrix minus a rank-one matrix. Dense, but with structure you can exploit: the VJP never needs the matrix.

One elementwise product, one dot product, one broadcast subtraction. \(O(n)\) work for an \(n \times n\) Jacobian. That expression — subtract the \(p\)-weighted mean of the incoming gradient, then reweight by \(p\) — is precisely the line \(dS_{ij} = P_{ij}(dP_{ij} - \sum_\ell P_{i\ell}dP_{i\ell})\) that appears in FlashAttention's backward algorithm, and precisely the \(\bar S\) formula in the attention derivation above.

Read it as a projection again, for the same reason as LayerNorm: softmax is invariant to adding a constant to all logits, so its gradient must be orthogonal to \(\mathbf{1}\) under the \(p\)-weighted inner product. Subtracting \(p^\top \bar p\) is exactly that orthogonalisation. Every softmax gradient sums to zero against \(p\): what one logit gains, the others give up.

The cancellation

Now compose with cross-entropy, \(L = -\sum_i y_i \log p_i\) for a one-hot \(y\). The incoming gradient is

which is a numerically alarming object: as the model becomes confident about a wrong answer, \(p_i \to 0\) at the true class and this blows up. Push it through the softmax VJP anyway. The dot product simplifies first:

and then everything falls out:

This is why every framework has a fused cross_entropy rather than making you compose log_softmax and nll_loss yourself, and why the fused version takes logits as its argument rather than probabilities. Computing \(p\) and then dividing by it is not merely wasteful; in bf16, where the smallest normal value is around \(10^{-38}\) but effective precision is about three decimal digits, a confidently wrong prediction produces a \(1/p\) that is meaningless before it is large. The cancellation has to happen symbolically, in the derivation, not numerically, at runtime.

It also has a pleasant interpretation. \(\bar z = p - y\) says: push down every logit in proportion to how much probability mass you put on it, and push up the true one by exactly one unit. Confident and right gives a gradient near zero. Confident and wrong gives a gradient near one. The loss is telling you, in the plainest possible terms, the size of your mistake.

What autograd actually keeps

Reverse mode has to run the forward pass first, and it needs forward values to evaluate local derivatives. So something must be retained. The question of exactly what is the difference between a model that fits and one that does not.

Saved tensors

Each backward rule declares its own dependencies. Look at what each one from the previous section needs:

Two patterns are visible. First, most saved tensors are inputs, which means they are also the outputs of the previous op — frameworks share the storage rather than copying, so the accounting is per-tensor, not per-op. Second, addition needs nothing at all, which is a small quiet reason residual connections are cheap: they add a route to the gradient and consume no extra memory doing it.

Activation lifetime

Here is the awkward shape of the problem. Layer 1's activation is produced at the very beginning of the forward pass and consumed at the very end of the backward pass. Its lifetime is the entire step. Layer \(k\)'s lifetime is shorter, and the last layer's is shortest of all. Memory therefore ramps up monotonically through the forward pass, peaks at the boundary, and drains through the backward pass.

For an actual transformer, the per-layer activation footprint has been worked out carefully. Korthikanti and colleagues give it as \(s\,b\,h\left(34 + 5\,\frac{a\,s}{h}\right)\) bytes per layer for a model with sequence length \(s\), batch \(b\), hidden size \(h\) and \(a\) attention heads, before any parallelism. The two terms tell the whole story: the first is linear in sequence length and comes from all the ordinary per-token tensors; the second is quadratic, and it is the attention probability matrix. At short sequences the linear term dominates and nobody thinks about it. At long sequences the quadratic term eats the GPU.

The API is the graph, exposed

Everything in PyTorch's autograd surface is a direct handle on some part of the machinery above. Once you see which part, the semantics stop needing to be memorised.

requires_grad: which subgraph gets built

A tensor with requires_grad=True is one the system will compute a gradient for. The flag is infectious in the forward direction: any operation with at least one input requiring grad produces an output requiring grad, and records a node. If no input requires grad, no node is recorded and the operation is as cheap as raw tensor math.

This is why freezing a backbone by setting requires_grad=False saves memory as well as compute: the frozen ops stop saving tensors, because nothing downstream will ever ask them for a gradient. And it is why freezing only the early layers saves much more than freezing later ones — the graph is only pruned up to the first parameter that still requires grad.

Leaf tensors and .grad

A leaf is a tensor with no grad_fn: it was created by the user, not produced by an operation. Model parameters are leaves. Inputs are leaves. Everything in the middle is not.

After backward(), gradients are written to .grad on leaves that required grad, and discarded everywhere else. That last part surprises people. The gradient with respect to an intermediate activation absolutely was computed — it had to be, it is how the gradient got past — but it is freed the moment its last consumer is done, because keeping every intermediate gradient would double activation memory for no reason. If you want one, retain_grad() asks for it, or a hook catches it in flight without keeping it alive.

And .grad accumulates. It is +=, the same += as the multivariable chain rule, applied one level up. Two backward calls without an intervening zero_grad() give you the sum of two gradients. This is not a design wart; it is the feature that makes the next item possible.

Gradient accumulation

You want an effective batch of 1024 and 128 fits in memory. Run eight microbatches, call backward() on each, step once:

optimizer.zero_grad()
for i, micro in enumerate(microbatches):          # 8 microbatches of 128
    loss = criterion(model(micro.x), micro.y)
    (loss / len(microbatches)).backward()      # the /K matters — see below
optimizer.step()

The mathematics is exact, not approximate, because the gradient of a mean is the mean of the gradients:

Two things to be careful about. The division by \(K\): if your loss already averages over its microbatch, you need it, or your gradient is \(K\) times too large and your learning rate is silently wrong. And any operation that is not a per-example average — batch normalisation is the notorious case, since its statistics are computed over whatever is present — is not exact under accumulation. Layer and RMS normalisation are per-example, so transformers are fine. In distributed training you also want to skip the all-reduce on every microbatch but the last, which is what no_sync() is for.

detach and no_grad: two different scissors

They are confused constantly and they do different things.

detach() is surgical: the forward value passes through unchanged, and the backward pass stops dead at that tensor. Use it when you want a value but not its history — a target network in RL, the stop-gradient in a self-distillation objective, a running statistic you do not want to differentiate through.

no_grad() is a mode: inside the block, operations do not record nodes and do not save tensors. Use it for inference and for optimiser internals. The memory saving is the real reason it exists; the compute saving is minor by comparison.

A third one exists and is worth knowing: torch.inference_mode(), which is no_grad() plus a promise that the tensors will never be used in autograd later, which lets the runtime skip version counters entirely. Faster, and irreversible — a tensor created inside inference mode cannot later be given a graph.

In-place operations, and why they are dangerous

The backward rule for \(y = x^2\) needs \(x\). If you overwrite \(x\) between the forward pass and the backward pass, the gradient computed will be wrong — not an error, just wrong, quietly, using whatever value happens to be in that memory.

PyTorch prevents this with a version counter on every tensor's storage. Every in-place operation bumps it. Every saved tensor records the version it saw. On backward(), the versions are compared, and if they disagree you get the message everyone has met at least once:

RuntimeError: one of the variables needed for gradient computation
has been modified by an inplace operation: [torch.FloatTensor [128, 512]],
which is output 0 of ReluBackward0, is at version 2; expected version 1

This is a good error. The alternative is silent numerical corruption, and the check costs one integer comparison. Note what it does not say: it does not say in-place operations are forbidden. relu_() is safe, because ReLU's backward needs only the sign of the output, which the in-place version still has. x.add_(1) on a tensor nobody saved is safe. The rule is narrower than the folklore: do not mutate a tensor that some backward rule is holding a reference to. When you hit the error, torch.autograd.set_detect_anomaly(True) will point you at the forward operation whose saved tensor was clobbered, at considerable runtime cost.

Writing your own backward

Autograd will differentiate anything you compose from primitives. Sometimes you should stop it and hand it the answer instead. Four situations where that is right:

You know a cheaper form. The fused cross-entropy is the canonical example: composing softmax and nll_loss gives a correct gradient through a chain that includes a division by \(p\), while the fused rule returns \(p - y\) directly. The autograd-derived version is not wrong on paper; it is worse in floating point and slower besides.

You want to trade memory for compute. A custom backward can recompute intermediates instead of saving them. This is checkpointing, and it is implemented in exactly this way.

The forward pass is not differentiable, or not by autograd. A call into a solver, a sampling step, a straight-through estimator, an implicit layer defined by a fixed point. You supply the mathematics that autograd cannot see.

You want to change the gradient on purpose. Gradient reversal for domain adaptation, straight-through estimators for quantisation, custom clipping. The backward pass need not be the true derivative of the forward pass. It just has to be the thing you want the optimiser to receive — though you should know when you are doing this, because it means you are no longer doing gradient descent on the loss you wrote down.

class FusedThing(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, w):
        y = expensive_forward(x, w)
        ctx.save_for_backward(x, w)      # NOT ctx.x = x — version tracking lives here
        return y

    @staticmethod
    def backward(ctx, grad_y):           # grad_y is the upstream adjoint
        x, w = ctx.saved_tensors
        return grad_y @ w.T, x.T @ grad_y   # one gradient per forward input

Three rules that are easy to get wrong. Use save_for_backward rather than stashing tensors on ctx directly, because that is what participates in version tracking and in reference-count-based freeing. Return exactly as many gradients as forward took inputs, using None for the ones that need none. And test it: torch.autograd.gradcheck compares your rule against finite differences in double precision, which is slow, numerically fussy, and has caught more sign errors than any amount of staring.

Trading memory for recomputation

Reverse mode's bill is now clear. Compute: about twice the forward pass, regardless of parameter count. Memory: every intermediate, alive from when it was produced until the backward pass reaches it. The first number is wonderful. The second is the one that keeps you awake.

Gradient checkpointing is the resolution, and the idea fits in a sentence: do not store what you can cheaply rebuild.

The arithmetic of segments

Split \(n\) layers into \(k\) segments of \(n/k\) layers each. During the forward pass, keep only the \(k\) segment boundaries and discard everything inside. During the backward pass, when you reach a segment, re-run its forward pass from its stored boundary to regenerate the \(n/k\) activations inside, use them, and throw them away again.

Minimise \(k + n/k\) over \(k\). Differentiate, set to zero, and \(k = \sqrt{n}\):

This is Chen and colleagues' result: \(O(\sqrt{n})\) memory for an \(n\)-layer network at the cost of one extra forward pass per minibatch, with a reported runtime overhead around 30% in their experiments. They also note the extreme point of the trade-off: recursive checkpointing reaches \(O(\log n)\) memory for \(O(n \log n)\) forward compute. The idea itself is older than deep learning — it is the checkpointing literature from AD, brought over and generalised.

Uniform checkpointing is the wrong default

The \(\sqrt{n}\) analysis assumes every layer costs the same to store and the same to recompute. In a transformer that is emphatically false, and exploiting the difference is worth a lot.

Korthikanti and colleagues make the point sharply. The tensors that dominate activation memory are not the ones that dominate recomputation cost. The \(N \times N\) attention matrix is quadratic in sequence length to store but comparatively cheap to regenerate; the big matmul outputs are linear to store but expensive to redo. So instead of checkpointing whole layers uniformly, checkpoint selectively: recompute the parts with a bad memory-per-FLOP ratio and store the rest. Combined with sequence parallelism, they report a 5× reduction in activation memory while cutting the execution-time overhead from recomputation from 36% to 4%.

Follow that logic to its conclusion and you get FlashAttention, which is checkpointing pushed inside a single operation. It never materialises \(S\) or \(P\) in high-bandwidth memory at all; it tiles the computation, keeps only the per-row softmax statistics from the forward pass, and recomputes the attention scores on-chip during the backward pass. The result is memory linear rather than quadratic in sequence length — and it is faster, because the recomputation happens in SRAM while the thing it replaced was a trip to HBM. That is the sharpest possible statement of the trade-off: sometimes recomputing is not a cost at all, because arithmetic is cheap and memory movement is not.

Where this lands in a real training step

Everything above has been mathematics with occasional glances at memory. Here is where it meets the machine — and the place to go for the full treatment is the systems companion to this article, Inside a training step, which walks the same step from the hardware's point of view.

The backward pass is not one kernel, it is a schedule

Reverse topological order is a partial order, not a total one. Independent branches can be differentiated in any relative order, or concurrently. Real engines exploit this: PyTorch's autograd engine maintains a ready queue with dependency counts and executes nodes on multiple threads and streams. The visual metaphor of a wave sweeping backwards through the network is only accurate for a plain chain.

Gradients become communication

In data-parallel training, each \(\bar W\) must be all-reduced across ranks before the optimiser step. Because \(\bar W\) for the last layer is ready first, you can start reducing it while the backward pass is still working on earlier layers. Bucketing gradients and launching an asynchronous all-reduce per bucket hides most of the communication behind the remaining compute. The backward pass is where distributed training's overlap opportunity lives, and it exists entirely because of the ordering constraint that makes reverse mode expensive in the first place.

Precision is not uniform

Mixed-precision training runs the forward and backward passes in bf16 or fp16 while keeping an fp32 master copy of the weights and accumulating in fp32. The reason is in the adjoint recurrence: gradients are products of many local derivatives, and products of small numbers underflow. With fp16 this required loss scaling — multiply the loss by a constant before backward(), divide it out before the step — to push the gradient distribution into representable range. bf16's wider exponent mostly removes the need, which is why it displaced fp16 for large-model training. Note that scaling the loss scales every gradient by exactly the same factor, because differentiation is linear; the trick is sound, not a hack.

Non-determinism comes from the accumulation

The += in the adjoint recurrence is floating-point addition, which is not associative. When multiple children contribute to a node's gradient and the order depends on thread scheduling, bitwise reproducibility goes away. This is the source of most "why did my run not reproduce" reports, and it is a real property of the algorithm rather than a bug — deterministic modes exist and cost throughput because they serialise exactly this.

The whole step, in one accounting

Per token, per parameter, for a dense transformer: about \(2N\) FLOPs forward, \(4N\) backward, \(6N\) total. Memory: \(2N\) bytes of bf16 weights, \(4N\) of fp32 master weights, \(8N\) of Adam moments, plus activations that scale with batch and sequence. Reverse mode is responsible for the \(4N\) and for the activations, and it is the reason the other numbers are affordable at all. The optimiser-state side of that ledger is the subject of What an Optimizer Actually Does.

What stays genuinely subtle

A few things that the tidy account above smooths over, and that I think are worth leaving open rather than resolving.

Non-differentiability is everywhere and mostly ignored. ReLU at zero, max at ties, abs, clamping, sorting, indexing by argmax. The standard defence — it is a measure-zero set, floating point will not land there — is true and slightly unsatisfying, since the exact points involved are often the ones a network is being pushed toward. The theory of conservative gradients puts this on firmer ground and extends the cheap-gradient principle to the nonsmooth case, but the practical upshot is that "the gradient" of a ReLU network is a choice, and different frameworks can legitimately make different ones.

The memory-versus-recompute optimum is not knowable in advance. The \(\sqrt{n}\) result is exact for a uniform chain and approximately useless for a real model on real hardware, where the answer depends on kernel fusion, on HBM bandwidth relative to arithmetic throughput, on whether recomputation overlaps with communication, and on what else is resident. Selective recomputation is the current best answer and it is still hand-tuned per architecture. There is no theory that tells you the schedule.

Optimal Jacobian accumulation is NP-complete. Between pure forward and pure reverse lie mixed strategies — eliminate graph edges in any order, and each order has a different cost. Finding the cheapest is, in general, NP-complete (Naumann, 2006). Every framework uses pure reverse mode not because it is optimal but because it is optimal for the specific case of one scalar output, and that case is the one we care about.

Higher-order derivatives are cheap in principle and awkward in practice. A Hessian-vector product is one forward-over-reverse composition and costs about the same as a gradient. But it requires the backward pass itself to be differentiable, which means the backward graph must be retained, which means memory again — and many fused kernels do not implement a double-backward at all. The gap between "mathematically a VJP of a VJP" and "runs on your GPU" is wider here than anywhere else in this article.

What backpropagation is not. It computes gradients. It says nothing about whether those gradients point anywhere useful, whether the loss landscape is navigable, or why the resulting model generalises. Those are entirely separate questions, and confusing "we can compute \(\nabla_\theta L\) efficiently" with "we understand what training does" has been a persistent source of muddle. All this machinery buys is an exchange rate. What you do with it is the rest of the field.

The thing that still strikes me, after all of it, is how little was needed. One identity about composing rates of change. One observation about the order you multiply matrices in. Everything else — the frameworks, the memory hierarchies, the fused kernels, the trillion-parameter training runs — is engineering built on top of a chain rule and a decision about parentheses.

References

Grouped by topic. Every entry links to its primary source.