Compute graphs: how deep learning frameworks track operations
A compute graph is a directed acyclic graph where nodes are values and edges are the operations that produced them.
In training, the forward pass builds the graph, and the backward pass walks it in reverse topological order, applying the chain rule at each node to compute gradients.
That's what loss.backward() does in PyTorch, and it's why training needs far more memory than inference: the graph's intermediate values are kept until the backward pass.
Forward and backward
Every operation records its inputs and how to compute its local derivative. After the forward pass, the loss is the last node. The backward pass starts at the loss with gradient 1, and each node passes gradient back to its inputs, multiplied by its local derivative. Inputs used in several places accumulate gradient from each use.
A tiny autograd engine
import math
class Value:
def __init__(self, data: float, parents: tuple = ()):
self.data = data
self.grad = 0.0
self.parents = parents
self.backward_fn = lambda: None
def __add__(self, other: "Value") -> "Value":
out = Value(self.data + other.data, (self, other))
def backward_fn():
self.grad += out.grad
other.grad += out.grad
out.backward_fn = backward_fn
return out
def __mul__(self, other: "Value") -> "Value":
out = Value(self.data * other.data, (self, other))
def backward_fn():
self.grad += other.data * out.grad
other.grad += self.data * out.grad
out.backward_fn = backward_fn
return out
def tanh(self) -> "Value":
t = math.tanh(self.data)
out = Value(t, (self,))
def backward_fn():
self.grad += (1 - t * t) * out.grad
out.backward_fn = backward_fn
return out
def backward(self) -> None:
order, seen = [], set()
def visit(node: "Value") -> None:
if id(node) not in seen:
seen.add(id(node))
for parent in node.parents:
visit(parent)
order.append(node)
visit(self)
self.grad = 1.0
for node in reversed(order):
node.backward_fn()
x, w, b = Value(2.0), Value(-0.5), Value(0.3)
y = (x * w + b).tanh()
y.backward()
print(f"y = {y.data:.4f}")
print(f"dy/dw = {w.grad:.4f}, dy/dx = {x.grad:.4f}, dy/db = {b.grad:.4f}")
The graph here is one neuron: multiply, add, tanh. A transformer is the same idea with billions of values grouped into tensors.
Topological order
The backward pass needs each node's gradient to be complete before it passes gradient on, so nodes are processed in reverse topological order. That's a depth-first search over a DAG, the same algorithm you'd use to order build steps or task dependencies, and it shows up again in graphs for agents.
Why training needs more memory than inference
Inference keeps only what the next layer needs. Training keeps intermediate activations for the backward pass, plus gradients, plus optimizer state (Adam keeps two extra values per parameter). A rough rule for full fine-tuning with Adam in mixed precision is about 16 bytes per parameter before activations, which is why methods like LoRA train a small number of new parameters and freeze the rest.
Static and dynamic graphs
PyTorch builds the graph as your Python code runs, so ordinary loops and conditions work.
Compilers like torch.compile and XLA capture the graph ahead of time so they can fuse operations and plan memory, which is how attention kernels avoid writing intermediate matrices to memory.
Check your understanding
Compute dy/dw for the example by hand with the chain rule and compare with the output.
From the floor
In the AI data structures lecture, compute graphs were the last stage in the sequence I laid out, after the KV cache, because once you understand that every operation on tensors is recorded as a graph, training, fine-tuning and the memory they need stop being a black box.