Computational Graphs

Every operation your model performs is recorded as a node in a one-way network, and training works by walking that network backwards — which is also why training eats so much memory.

A framework records every operation as a node in a directed acyclic graph, runs it forward caching each value, then walks the same graph in reverse to accumulate gradients

Why Does This Exist?

You write four lines of Python and call .backward(), and correct gradients appear for every parameter. Nobody derived them. There is no symbolic differentiation step you can inspect. What happened?

A graph happened. While your forward code ran, the framework was recording — every operation became a node, every data dependency became an edge, and every intermediate value was stashed at the node that produced it. .backward() walks that recording in reverse.

This is the page that explains three things you otherwise have to accept on faith. Why training needs several times more memory than inference — the forward pass is holding everything for the backward pass. Why gradient checkpointing trades compute for memory and what it actually discards. And why an in-place operation can silently corrupt your gradients: it overwrites a value the graph was keeping.

The chain rule says backpropagation is that rule "applied over a computational graph". This is the graph. The recording-and-replay machinery is automatic differentiation, and the graph vocabulary is graph fundamentals.

Think of It Like This

A kitchen that keeps every intermediate bowl

A chef follows a recipe: chop onions, fry them, add stock, blend, season. A one-way sequence, each step consuming what the last produced.

Now the chef must answer a question after serving: "if I had used 10 percent more onion, how much saltier would the final dish have tasted?" Answering it requires walking backwards through every step, and at each one asking how sensitive that step's output was to its input. But the sensitivity of the blending step depends on how much liquid was in the pot at that moment — you cannot reason about it from the recipe alone. You need the actual state.

So the chef keeps every intermediate bowl. Nothing gets washed until the question is answered. That is the memory cost of training, and it is not optional: the sensitivity of squaring a number depends on the number.

Gradient checkpointing is the chef keeping only every fourth bowl and re-cooking the missing steps when needed. Less shelf space, more cooking.

How It Actually Works

A computational graph is a DAG where each node is an operation and each edge carries a value from producer to consumer. Leaf nodes are inputs and parameters; the single root is usually the scalar loss.

It is a graph in the exact technical sense, and being acyclic is what makes everything work: a DAG has a topological order, so the operations can be sequenced with every input ready before it is needed.

Two passes over one structure

Forward. Traverse in topological order, computing each node's output and caching it at the node. This produces the loss, and it leaves the graph decorated with every intermediate value.

Backward. Traverse in reverse topological order, starting from $\partial \mathcal{L} / \partial \mathcal{L} = 1$ at the root. At each node, multiply the incoming gradient by that node's local derivative and pass the result to its inputs. When a node feeds several consumers, its incoming gradients sum — which is the multiple-paths rule from the chain rule, implemented as accumulation.

Each node only needs to know how to differentiate itself. That locality is why you can write a custom layer without touching the optimiser, and why frameworks ship thousands of ops rather than one giant derivative.

Static versus dynamic graphs

Two designs, and the distinction still shapes the tooling:

Static (define-then-run) — build the whole graph up front, then feed data through. The compiler sees everything, so it can fuse and optimise aggressively, but control flow that depends on data is awkward and debugging is indirect.
Dynamic (define-by-run) — the graph is recorded fresh on every forward pass, as ordinary Python executes. Python if and while just work, stack traces point at your code, and each iteration may have a different graph.

Dynamic won on ergonomics, and then the gap closed from the other side: torch.compile traces a dynamic graph and hands it to a compiler, recovering static-graph optimisation.

Worked example

Take the graph in the diagram: $w = 3$ , $x = 2$ , then $z = w \cdot x$ , $a = z^2$ , $\mathcal{L} = a / 4$ .

Forward, in topological order. $z = 3 \times 2 = 6$ , cached at the multiply node. $a = 6^2 = 36$ , cached at the square node. $\mathcal{L} = 36/4 = 9$ .

Backward, in reverse order. Start at the root with $\partial \mathcal{L} / \partial \mathcal{L} = 1$ .

Divide node: local derivative $1/4$ , so $\partial \mathcal{L} / \partial a = 1 \times 0.25 = 0.25$ .
Square node: local derivative $2z = 12$ — and this needs the cached $z = 6$ . So $\partial \mathcal{L} / \partial z = 0.25 \times 12 = 3$ .
Multiply node: local derivative with respect to $w$ is $x = 2$ . So $\partial \mathcal{L} / \partial w = 3 \times 2 = 6$ .

Verify independently: substituting gives $\mathcal{L} = (wx)^2/4$ , and at $x = 2$ that is $4w^2/4 = w^2$ , so $\partial \mathcal{L} / \partial w = 2w = 6$ . Agreed.

The line to stare at is the square node. Its local derivative was $2z$ , a number that only existed because the forward pass ran. Overwrite that cached 6 and the backward pass computes a wrong gradient with no error raised. That is the in-place bug, made concrete.

Code

The point is the cache: a three-node tape records forward values and replays them backward, using only the standard library.

1class Tape:2    """A three-node graph, recorded forward and replayed backward."""3
4    def forward(self, w: float, x: float) -> float:5        self.x, self.z = x, w * x  # cache z: the backward pass needs it6        self.a = self.z ** 27        return 0.25 * self.a8
9    def backward(self) -> float:10        d_a = 0.25               # dL/da11        d_z = d_a * 2 * self.z   # dL/dz <- uses the CACHED z12        return d_z * self.x      # dL/dw13
14
15tape = Tape()16print(tape.forward(3.0, 2.0), tape.backward())  # -> 9.0 6.017
18# Corrupt the cache the way an in-place op would. No error, wrong gradient.19tape.forward(3.0, 2.0)20tape.z = 999.021print(tape.backward())  # -> 999.0, silently wrong

Nine and six match the hand calculation. The last line is the whole warning: with $z$ corrupted to 999, the backward pass returns $0.25 \times 2 \times 999 \times 2 = 999.0$ , and nothing raises.

Watch Out For

Retaining the graph in an evaluation loop

Every forward pass builds a graph and holds activations for a backward pass that, during evaluation, never comes. Run validation without torch.no_grad() and memory climbs with every batch until the process dies.

It is easy to write, because the evaluation loop looks like the training loop minus .backward() — and the missing piece is invisible. Worse, a small validation set finishes before memory runs out, so the bug ships and appears later when the set grows or the batch size changes.

The related version is accumulating metrics wrongly. total_loss += loss keeps a reference to the loss tensor, which keeps its whole graph alive, so you retain one graph per batch for the entire epoch. total_loss += loss.item() extracts a plain float and lets the graph be freed.

Wrap evaluation in torch.no_grad() (or torch.inference_mode(), which is stricter and slightly faster), and detach anything you store.

Modifying a tensor the graph is still holding

In-place operations overwrite memory that a cached node may still need. When the backward pass reaches that node, the value it relies on is gone.

PyTorch catches many of these with a version counter and raises an error naming the modified variable. That error is a gift: it is loud, and it names the problem. The dangerous cases are the ones the counter misses, where you get a wrong gradient and a model that trains slightly worse for no visible reason.

The rule of thumb: in-place is safe on a tensor that is not needed for any backward computation. When in doubt, write the out-of-place version and let the framework's memory planner handle it. And if a gradient looks subtly wrong, run a numerical gradient check before suspecting anything else.

The Quick Version

A computational graph is a DAG: nodes are operations, edges carry values, the root is the scalar loss.
Forward pass traverses topological order and caches each node's output. Backward pass traverses the reverse.
Each node differentiates only itself. Gradients from multiple consumers are summed at a node.
Being acyclic is what guarantees a topological order exists, and therefore that this works at all.
Static graphs optimise better; dynamic graphs debug better. torch.compile traces a dynamic graph to get both.
With $w = 3$ and $x = 2$ , the tape gives $z = 6$ , $a = 36$ , loss 9, and gradient 6.
Corrupting the cached 6 to 999 silently moves the gradient to 999.0.
torch.no_grad() for evaluation, and .item() when accumulating metrics. Both prevent retained graphs.
In-place operations can overwrite a cached value the backward pass needs.