Computational Graphs
Nodes are operations, edges carry values forward and gradients backward. Reverse mode is cheap for many parameters and one loss, the framework builds the graph as you call it, and the activations it stores are the memory bill.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
Why does a framework record a graph of every operation, why is reverse mode the right direction for training, and why does the memory cost of training scale with the activations rather than the parameters?
A team training a transformer runs out of GPU memory at a sequence length the parameter count says should fit easily. "The weights are a few gigabytes. Why does a longer input blow the memory, and what is this activation checkpointing flag everyone says to turn on?"
Memory is the model size. If the weights fit with room to spare, the batch fits. If it runs out, buy a bigger GPU.
The backward pass needs the forward value of every node whose local derivative depends on it — the input of every matrix multiply, the pre-activation of every ReLU, the attention weights of every head. The framework keeps all of them from the forward pass until the backward pass has consumed them. For a deep network on a long batch that is many times the parameter memory (Memory Bandwidth & VRAM).
- The backward pass needs the forward value of every node whose local derivative depends on it — the input of every matrix multiply, the pre-activation of every ReLU, the attention weights of every head. The framework keeps all of them from the forward pass until the backward pass has consumed them. For a deep network on a long batch that is many times the parameter memory (Memory Bandwidth & VRAM).
- Doubling the sequence length doubles most activation tensors and quadruples the attention ones, while the weights are unchanged; the "model size" estimate was off by the entire activation term (Self-Attention).
- A bigger GPU moved the wall, and the next experiment with a longer context hit it again. The cost is structural: it is the price of reverse-mode differentiation, not of the hardware.
- An engineer turned on activation checkpointing everywhere and the step time rose by a third; nobody had said that checkpointing trades memory for a second forward computation.
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- The surrounding system trains a sequence model with a cross-entropy loss over tokens; the target of this lesson is the graph the framework records during the forward pass and what it must keep to run the backward pass.
- The decision the team faces is memory against compute: store every activation and run fast, or recompute some and fit a longer sequence.
- One training example is a token sequence; a batch is several. Every layer's output for every token in the batch is an activation tensor, and there are dozens of layers.
- The parameters are fixed in size; the activations scale with batch × sequence length × width × layers, and for attention with sequence length squared.
- The team profiled peak memory and found it dominated by tensors that exist only between the forward and backward passes.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- A computational graph is a directed acyclic graph whose nodes are operations — matrix multiply, add, ReLU, softmax, log — and whose edges are tensors. The forward pass evaluates nodes in topological order and produces each edge's value; the backward pass visits nodes in reverse topological order and produces each edge's gradient, ∂L/∂edge, by multiplying the consumer's gradient by the local derivative (Backpropagation).
- Reverse mode computes the gradient of one output with respect to every input in one backward sweep, at a cost comparable to the forward pass. Forward mode would need one sweep per input. Training has one scalar loss and millions of parameters, so reverse mode is the only affordable direction; that asymmetry is the whole reason backpropagation is the training algorithm and not a curiosity.
- Frameworks build the graph dynamically: each tensor operation you call records itself and its inputs, so by the end of the forward pass the graph exists for exactly the operations that ran — including data-dependent branches.
backward()then walks that record. Nothing is derived symbolically; the local derivative of each primitive is code the framework ships. - Because a node's local derivative usually depends on its forward inputs — ∂(Wx)/∂W needs x, ∂relu(z)/∂z needs the sign of z, ∂softmax needs the softmax — the framework must keep those inputs alive until the backward pass reaches them. That is the activation memory: proportional to every intermediate tensor in the graph, which for a deep model over a long batch dwarfs the parameters. Activation checkpointing drops some of them and recomputes them during the backward pass, trading memory for a partial second forward pass (Checkpointing).
Nodes, edges, and the two directions
Every tensor operation in the forward pass becomes a node; the tensors between them are edges. The forward pass produces values along the edges in topological order. The backward pass produces gradients along the same edges in reverse — each node receives ∂L/∂output from its consumers, multiplies by its local derivative, and hands ∂L/∂input back to its producers. A node with two consumers sums the two gradients it receives.
The graph is a DAG, and the backward sweep is a topological sort run backwards; Algorithms owns the general algorithm and this domain uses one instance of it. What makes it automatic is that each primitive ships with its local derivative, so the framework never needs to know what the whole function is.
1class Node:2 def __init__(self, value, parents=(), local_grads=()):3 self.value, self.parents, self.local_grads = value, parents, local_grads4 self.grad = 0.05 6def mul(a, b): return Node(a.value * b.value, (a, b), (b.value, a.value)) # d(ab)/da = b7def add(a, b): return Node(a.value + b.value, (a, b), (1.0, 1.0))8def relu(a): return Node(max(0.0, a.value), (a,), (1.0 if a.value > 0 else 0.0,))9 10def backward(loss):11 order, seen = [], set()12 def visit(n): # topological order by DFS13 if id(n) in seen: return14 seen.add(id(n)); [visit(p) for p in n.parents]; order.append(n)15 visit(loss)16 loss.grad = 1.017 for n in reversed(order): # reverse topological order18 for p, g in zip(n.parents, n.local_grads):19 p.grad += n.grad * g # chain rule, one edge at a timeEvery Node keeps the forward values its local derivatives need — mul stores both operands. Multiply that by every operation in a deep model over a long batch and you have the activation memory that the parameter count never mentioned.
Why reverse mode, and what it costs
The chain rule can be applied in either direction. Forward mode carries a derivative alongside each value and answers "how does everything change if this one input moves" — one sweep per input. Reverse mode carries a gradient backward and answers "how does this one output change with every input" — one sweep per output. Training has one output, the loss, and millions of inputs, the parameters. The direction is forced.
The price is storage. Reverse mode cannot begin until the forward pass is complete, and it needs the forward inputs of every node whose derivative depends on them. Those tensors are live from the moment they are computed until the backward sweep consumes them — which for the first layer is the entire step.
One full pass through the graph per parameter. Exact (forward mode) or approximate (finite differences); for a million parameters, a million passes per step.
One forward pass that records values, one backward pass that produces every parameter's gradient; total cost a small multiple of the forward pass, memory proportional to stored activations.
Training differentiates a single scalar with respect to many parameters, and reverse mode's cost is per output while forward mode's is per input. The activation memory is the price of that asymmetry, and checkpointing is the knob that trades it back for compute.
Activation checkpointing and what it assumes
Checkpointing drops the stored activations inside chosen blocks after the forward pass and recomputes them from the block's input when the backward sweep arrives. Memory for those blocks falls to the block inputs; compute rises by the cost of re-running the block's forward pass. Applied to every block it can cut activation memory by a large factor at roughly a third more step time; applied selectively it costs less.
The assumption it introduces is that recomputation reproduces the forward pass exactly. A block with dropout must re-seed to draw the same mask; a block with a data-dependent branch must take the same branch. A framework that gets this wrong hands the backward pass activations that do not match the forward values, and the gradient is wrong without any error.
Every activation the backward pass uses — stored or recomputed — is the value the forward pass produced, and the graph it walks is the graph that ran.
holds when Checkpointed blocks are deterministic or re-seeded per block; the batch shape and precision at training time are the ones the memory estimate and the equivalence test used.
breaks when A checkpointed block contains un-seeded randomness; mixed precision stores an activation in a lower precision than the forward computation used; a static export traces one branch of a dynamic graph.
respond Disable checkpointing on the offending block and confirm gradients match; fix the seeding; re-estimate memory with all four terms before choosing the batch shape.
How to build it
Most important first.
- Estimate memory as parameters plus optimiser state plus gradients plus activations, and compute the activation term from batch × sequence × width × layers; it is usually the largest and the only one that moves with the input (Training Cost).
- Apply activation checkpointing per block, not everywhere: checkpoint the blocks whose activations are largest and whose recomputation is cheapest, and measure the step-time cost.
- Reduce activation memory before adding hardware — smaller micro-batches with gradient accumulation, mixed precision for the stored tensors, attention variants that do not materialise the full attention matrix.
- When the model itself no longer fits, move to model or pipeline parallelism, which is a distributed-training decision with its own costs (Model, Tensor & Pipeline Parallelism).
- Profile peak memory by tensor, in the training loop, on the real batch shape; the number from a parameter count is a lower bound that is wrong by the activation term.
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- Peak memory per step broken down into parameters, gradients, optimiser state and activations; the activation share is what checkpointing and batch shape control.
- Step time with and without checkpointing on each block; the recompute cost is the number the memory saving is bought with.
- Activation memory as a function of sequence length, measured, so the wall is predicted before an experiment hits it.
- Do not measure "model size in gigabytes" and call it the memory requirement of training. It is the requirement of inference, at batch one, before the optimiser state.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- The graph the framework records during the forward pass is the graph the backward pass walks — recomputed blocks reproduce the same values, including any randomness.
- Peak memory on the production batch shape, with the chosen checkpointing and precision, stays within the device; a change to sequence length or batch size is a memory change.
- The optimiser state and gradients for every parameter fit alongside the activations; the estimate included all four terms.
- A static-graph export for serving covers every branch the dynamic graph can take at training time.
- Offline: a memory profile on the real batch shape with per-tensor attribution; a unit test that the checkpointed and un-checkpointed backward passes produce identical gradients on a fixed batch.
- During training: peak memory and step time as logged metrics; an alert when either moves after a configuration change.
- Over time: a smoke test that trains a few steps at the maximum intended sequence length on the target device before a long run is launched (Training Smoke Tests).
What can go wrong
- Checkpointing recomputes a block that contains a stochastic operation — dropout — with a different random draw, and the backward pass uses activations that do not match the forward pass; the gradient is silently wrong unless the framework re-seeds per block.
- Gradient accumulation over micro-batches changes the effective batch statistics of normalisation layers, and the model trained at micro-batch eight is not the model validated at batch sixty-four (Normalisation Layers).
- A data-dependent branch in the model makes the graph differ between examples; a serving-time export that assumes a static graph traces one branch and silently serves it for every input.
- Mixed precision halves the stored activations and introduces the overflow that produces NaN gradients on rare batches (Vanishing and Exploding Gradients).
- Reverse mode gives every gradient for the price of one backward sweep and charges for it in stored activations; there is no version of backpropagation that avoids the storage without recomputation.
- Checkpointing recovers memory at the cost of recomputing the forward pass for the checkpointed blocks — a step-time increase that has to be weighed against the longer context or larger batch it enables.
- Dynamic graph construction gives flexibility — branches, loops, data-dependent shapes — and makes the graph an artefact of one execution, which complicates export, optimisation and reproducibility.
- "The weights fit, so training fits." Training holds the weights, their gradients, the optimiser state and every activation the backward pass will need. The last term scales with the input and is usually the largest.
- "Turn on checkpointing; it is free memory." It is memory bought with a partial second forward pass. On a compute-bound step that is a real slowdown, and on a block with randomness it can be a wrong gradient.
- "Reverse mode is just the efficient implementation of the chain rule." It is the efficient direction for one output and many inputs. For a function with many outputs and one input, forward mode wins; training happens to be the first case.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- GENERALReverse-mode differentiation over a recorded graph is how every major framework trains every differentiable model; the graph, its topological order and the stored activations are the same mechanism from a two-unit network to a transformer.
- FRAMEWORK-SPECIFICWhether the graph is recorded dynamically per execution or compiled ahead of time, how checkpointing handles randomness, and what the memory profiler attributes to each tensor differ by framework; the accounting is the same, the tooling and the defaults are not.
- SCALE-SPECIFICFor a small dense model the activation term is negligible and none of this matters; for a deep sequence model the activation term dominates memory, and for attention it grows with sequence length squared until an attention variant that avoids materialising the matrix is used.
Where the depth lives
This domain teaches the model and hands the rest off by name.
- — Programming Languages & Runtime Internals — a dynamically recorded graph is a trace of one execution, and the difference between tracing and compiling a model is a language-runtime question this domain names and does not own.