Computation graphs and autograd

Problem

Backpropagation needs the whole chain of operations that produced the loss, in order, along with the intermediate values each local derivative uses. You could write that chain down by hand for a fixed network, but real models branch on data, loop a variable number of times, and change shape from batch to batch. Something has to record the chain while the forward pass runs, decide which tensors need gradients, keep only the values backward will need, and throw it all away afterward so memory does not grow with every step.

In PyTorch that something is autograd. Most of the time it is invisible: you call loss.backward() and read .grad. But the moment you freeze part of a model, evaluate without gradients, call backward twice, clip a gradient in flight, or write an operation PyTorch does not know how to differentiate, you are working with the graph directly, and its rules decide what happens.

Intuition

Think of the forward pass as leaving a trail of receipts. Each operation writes one: "I multiplied these two things; to send a gradient back, multiply it by the other input, and here is a copy of that input". The receipt points back to the receipts of whatever produced its inputs. The loss holds the last receipt, so following pointers from it reaches every operation that mattered, and nothing else.

Backward reads the receipts in reverse. It starts at the loss with a gradient of 1, asks each receipt to turn the gradient of its output into gradients of its inputs, and passes those along the pointers. Where the trail starts, at the parameters you created yourself, a special receipt adds the arriving gradient into .grad. Once a receipt has been used, its saved copies are shredded unless you asked to keep them.

The graph is dynamic: it is built fresh by running ordinary Python each step. An if that takes the other branch next step simply records different nodes. There is no separate "compile the graph" phase, which is why PyTorch code reads like normal code and debugs with print statements.

Mechanics

Nodes, grad_fn and next_functions

When an operation runs with gradient recording on and at least one input has requires_grad=True, PyTorch creates a backward node and stores it on the output as output.grad_fn. The node's class name says which derivative it computes: MulBackward0, AddBackward0, ClampBackward1 (the digit tells overloads apart). Every node is a torch.autograd.graph.Node with a name() and a next_functions tuple.

next_functions holds one (node, index) pair per input, in argument order. The node is the input's own grad_fn; it is None for an input that does not require gradients; and the integer says which output of that node the edge connects to (0 for single-output operations). Following next_functions from loss.grad_fn visits the whole graph.

Each node also saves what its derivative needs. MulBackward0 saves the other operand, ExpBackward0 saves its own output, AddBackward0 saves nothing. These saved tensors are most of the memory a training step uses beyond the weights.

Leaf tensors and AccumulateGrad

A tensor is a leaf when grad_fn is None: you created it (a parameter, an input batch), or it came out of operations on tensors that do not require gradients. A tensor that requires gradients and was produced by a recorded operation is not a leaf.

The edge to a leaf that requires gradients points at an AccumulateGrad node, one per leaf. When backward reaches it, it adds the incoming gradient into the leaf's .grad: it stores it if .grad is empty and adds to it otherwise. That is why gradients accumulate across backward() calls and you call optimizer.zero_grad() each step. Non-leaf tensors do not get a .grad (it stays None) unless you call retain_grad() on them.

Because loss = y·y uses y twice, MulBackward0.next_functions lists AddBackward0 twice. Edges point from a node to the nodes that produced its inputs, which is the direction gradients flow.

How backward walks the graph

loss.backward() calls torch.autograd.backward, which hands the roots to the C++ engine. The engine works in three steps:

1. Count dependencies. Starting from the root, it visits every reachable node once and counts, for each node, how many edges point into it. In the example AddBackward0 gets 2.
2. Run ready nodes. A node is ready when all of its incoming gradients have arrived (its count reaches zero). Their gradients are summed in an input buffer first, so a tensor used twice gets the sum of both contributions, which is the multivariable chain rule. Among ready nodes on one device, the one with the highest sequence_nr (created most recently in the forward pass) runs first. AccumulateGrad nodes have the largest possible sequence number, so leaf gradients are written as soon as they are complete.
3. Release. After a node runs, unless the graph is being kept, the engine calls its release_variables(), freeing the saved tensors.

retain_graph and create_graph

Because saved tensors are freed as backward goes, a second backward() through the same nodes fails with "Trying to backward through the graph a second time". Pass retain_graph=True to keep them, for example when two losses share a trunk and you call backward on each. Its default is None, which means "the same as create_graph". create_graph=True records the backward pass itself as a graph, so the gradients are differentiable and you can take second derivatives (Hessian-vector products, gradient penalties). Calling backward on the sum of the losses usually beats retaining the graph.

no_grad, inference_mode and detach

• torch.no_grad() switches recording off for its block on the current thread. Every result has requires_grad=False and no grad_fn, so no nodes or saved tensors are kept. Use it for evaluation and for in-place parameter updates. Factory functions that take a requires_grad argument (like creating an nn.Parameter) are not affected.
• torch.inference_mode() goes further: it also skips view tracking and version counter bumps, so it is faster, but tensors made inside it can never be used later in a recorded computation.
• t.detach() returns a new tensor that shares storage with t but has no grad_fn and requires_grad=False: a leaf cut off from the graph. Gradients stop at it. It is how you use a value without differentiating through it: a target network in reinforcement learning, a stop-gradient in a contrastive loss, or a value you want to log.
• requires_grad_(False) on parameters freezes them: operations that touch only frozen tensors record nothing.

Hooks

Hooks run your function in the middle of backward:

• tensor.register_hook(fn) calls fn(grad) when the gradient with respect to that tensor has been computed. If fn returns a tensor, that replaces the gradient for everything further back. It raises if the tensor does not require gradients. Use it to inspect, clip or rescale a gradient in flight.
• tensor.register_post_accumulate_grad_hook(fn), leaves only, runs after .grad has been updated and receives the tensor itself. Optimizers can step a parameter the moment its gradient is ready and free the gradient right away.
• node.register_hook(fn) and node.register_prehook(fn) on a grad_fn see the gradients going into and out of one node.
• torch.autograd.graph.saved_tensors_hooks(pack, unpack) intercepts every tensor saved for backward, for example to move it to CPU (save_on_cpu) or compress it.

Module-level hooks (register_full_backward_hook) are built from these.

Custom autograd.Function

To define your own operation's derivative, subclass torch.autograd.Function with two static methods. forward(ctx, *inputs) computes the output and stashes what backward needs with ctx.save_for_backward(...) (tensors) or as plain attributes on ctx (everything else). backward(ctx, *grad_outputs) returns one gradient per forward input, None for inputs that need none. Call it with MyFn.apply(x), never MyFn()(x); its node in the graph is named after the class, MyFnBackward.

The usual reasons are a derivative that is cheaper or more stable than the one autograd would compose, an operation implemented outside PyTorch, or a deliberately different gradient such as the straight-through estimator. Helpers on ctx: mark_dirty for in-place outputs, mark_non_differentiable for outputs like indices, and set_materialize_grads(False) to receive None instead of zeros. Check a hand-written backward with torch.autograd.gradcheck.

Worked example

Steps 3 and 4 were both ready after step 2; AccumulateGrad's maximal sequence number puts it first. Without retain_graph=True each node frees its saved values right after it runs.

Implementation

The worked example in PyTorch:

import torch

x = torch.tensor(2.0)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
y = w * x + b
loss = y * y

print(loss.grad_fn)                    # <MulBackward0 ...>
print(loss.grad_fn.next_functions)     # ((<AddBackward0>, 0), (<AddBackward0>, 0))
print(y.grad_fn.next_functions)        # ((<MulBackward0>, 0), (<AccumulateGrad>, 0))
print(w.is_leaf, y.is_leaf)            # True False

loss.backward()
print(w.grad, b.grad)                  # tensor(28.) tensor(14.)
# loss.backward()                      # RuntimeError: Trying to backward through the graph a second time ...

with torch.no_grad():
    z = w * x
print(z.requires_grad, z.grad_fn)      # False None


class StraightThroughClamp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, t, low, high):
        return t.clamp(low, high)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None, None   # identity for t; low and high get no gradient


w.grad = b.grad = None
y = w * x + b
c = StraightThroughClamp.apply(y, 0.0, 5.0)
print(c.grad_fn.name())                # StraightThroughClampBackward
(c * c).backward()
print(w.grad, b.grad)                  # tensor(20.) tensor(10.)

A minimal engine that follows the same rules, in plain Python (the tested helper behind this page's numbers is a slightly longer version):

class Node:
    count = 0

    def __init__(self, name, inputs, backward):
        self.name, self.backward = name, backward
        self.next = [t.grad_fn or t.acc if t.requires_grad else None for t in inputs]
        self.seq = Node.count
        Node.count += 1


def backward(root):
    deps, stack, seen = {}, [root], {root}
    while stack:                                   # 1. count incoming edges
        for n in filter(None, stack.pop().next):
            deps[n] = deps.get(n, 0) + 1
            if n not in seen:
                seen.add(n)
                stack.append(n)
    buf, ready = {root: 1.0}, [root]
    while ready:                                   # 2. run newest ready node
        ready.sort(key=lambda n: n.seq)
        node = ready.pop()
        grads = node.backward(buf.pop(node))       # AccumulateGrad: adds into .grad, returns ()
        for n, g in zip(node.next, grads):
            if n is None:
                continue
            buf[n] = buf.get(n, 0.0) + g           # sum contributions
            deps[n] -= 1
            if deps[n] == 0:
                ready.append(n)

Tradeoffs

• Dynamic graphs. Rebuilding the graph every step costs Python overhead per operation and hides the whole program from the compiler, but it makes control flow, debugging and variable shapes trivial. torch.compile recovers much of the speed by tracing the same Python into a fused graph.
• Memory. Saved tensors scale with depth, width and batch size. no_grad during evaluation, activation checkpointing (recompute instead of save) and saved_tensors_hooks to offload are the levers.
• retain_graph=True keeps every saved tensor alive until the graph is dropped. Summing losses and calling backward once is almost always better.
• detach versus no_grad. detach cuts one tensor out and keeps recording everything else; no_grad stops recording for a block. Detaching shares storage, so an in-place change to the detached tensor changes the original and can break a later backward.
• Hooks are powerful and easy to get wrong: a tensor hook that returns a new gradient changes training silently, and hooks registered inside a loop pile up unless you keep the handle and call remove().
• Custom Functions buy control over the derivative at the cost of correctness being your problem. Run gradcheck in double precision on small inputs, unless the gradient is deliberately not the true derivative.

Check yourself

Connections

References

• Paszke et al. Automatic Differentiation in PyTorch. NIPS Autodiff Workshop, 2017.
• Paszke et al. PyTorch: An Imperative Style, High-Performance Deep Learning Library. NeurIPS, 2019 (arXiv 1912.01703).
• Baydin, Pearlmutter, Radul and Siskind. Automatic Differentiation in Machine Learning: a Survey. JMLR, 2018 (arXiv 1502.05767).
• Bengio, Léonard and Courville. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. 2013 (arXiv 1308.3432; the straight-through estimator).
• PyTorch v2.15.0: torch/autograd/graph.py, torch/autograd/function.py, torch/_tensor.py, once_differentiable, mark_non_differentiable, set_materialize_grads, and the engine in engine.cpp.