One call, three graphs

A compiler backend under torch.compile is just a function that takes an FX graph and gives back something callable. That makes it the easiest instrument in the whole stack: write a backend that prints its argument and returns it unchanged, and you get to see exactly what dynamo produced, with nothing between you and the object.

Do that on a two-layer step and the graph that arrives is not what chapter 7's closing line would have you expect. aot_module_simplified is the function that turns it into ATen, and it runs inside the backend, after dynamo is already done. So there is a moment where the captured graph exists and is not ATen at all.

The script below stacks three printers: one for the graph dynamo hands over, one for the forward that comes out of aot_autograd, and one for the backward. All three fire from a single compiled call plus a single .backward().

import torch
import torch.nn as nn
from torch._functorch.aot_autograd import aot_module_simplified
from functorch.compile import make_boxed_func

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(3, 4)
    def forward(self, x):
        return torch.relu(self.fc(x)).sum()

def show(tag):
    def compiler(gm, inputs):
        print(f"--- {tag}")
        print(gm.code.strip())
        return make_boxed_func(gm.forward)
    return compiler

def backend(gm, example_inputs):
    print("--- what dynamo handed the backend")
    print(gm.code.strip())
    return aot_module_simplified(gm, example_inputs,
                                 fw_compiler=show("aot forward"),
                                 bw_compiler=show("aot backward"))

torch.manual_seed(0)
loss = torch.compile(MLP(), backend=backend)(torch.randn(2, 3))
loss.backward()

Read the first block and count the targets. self.L__self___fc(l_x_) is a module call, still a whole nn.Linear sitting unopened in the graph. torch.relu is the Python-level function, not aten.relu.default. relu.sum() is a method call on a tensor. Nothing in that graph is an operator overload yet.

The second block is the same computation with every one of those opened up. The Linear became a transpose and an addmm; the method call became aten.sum.default. Dynamo never did that rewriting. It happened when aot_autograd re-traced the graph under a dispatcher that records what each Python-level call actually dispatches to.

The third block is the piece chapter 7 promises exists and never shows. It takes primals_3 and relu, which are two of the three values the forward returned, plus tangents_1, which is the incoming gradient. threshold_backward is the relu derivative; the two transposes and the mm are the weight gradient; sum.dim_IntList over dimension 0 is the bias gradient.

The forward graph returns more than the forward result. The extras are the backward’s inputs.

--- what dynamo handed the backend
def forward(self, L_x_ : torch.Tensor):
    l_x_ = L_x_
    l__self___fc = self.L__self___fc(l_x_);  l_x_ = None
    relu = torch.relu(l__self___fc);  l__self___fc = None
    sum_1 = relu.sum();  relu = None
    return (sum_1,)
--- aot forward
def forward(self, primals_1, primals_2, primals_3):
    t = torch.ops.aten.t.default(primals_1);  primals_1 = None
    addmm = torch.ops.aten.addmm.default(primals_2, primals_3, t);  primals_2 = t = None
    relu = torch.ops.aten.relu.default(addmm);  addmm = None
    sum_1 = torch.ops.aten.sum.default(relu)
    return [sum_1, primals_3, relu]
--- aot backward
def forward(self, primals_3, relu, tangents_1):
    detach = torch.ops.aten.detach.default(relu);  relu = None
    expand = torch.ops.aten.expand.default(tangents_1, [2, 4]);  tangents_1 = None
    detach_1 = torch.ops.aten.detach.default(detach);  detach = None
    threshold_backward = torch.ops.aten.threshold_backward.default(expand, detach_1, 0);  expand = detach_1 = None
    t_1 = torch.ops.aten.t.default(threshold_backward)
    mm = torch.ops.aten.mm.default(t_1, primals_3);  t_1 = primals_3 = None
    t_2 = torch.ops.aten.t.default(mm);  mm = None
    sum_2 = torch.ops.aten.sum.dim_IntList(threshold_backward, [0], True);  threshold_backward = None
    view = torch.ops.aten.view.default(sum_2, [4]);  sum_2 = None
    t_3 = torch.ops.aten.t.default(t_2);  t_2 = None
    return [t_3, view, None]

Two separate graphs come out, but nothing traced them separately. aot_autograd builds one graph containing both halves and then hands it to a partitioner, and the docstring on create_aot_dispatcher_function says so in order: trace to a joint graph, pass it through partition_fn to isolate the forward and backward portions, compile each with its own compiler.

You can watch the cut happen by writing a partition_fn that prints its argument and then delegates. The joint graph below is from a smaller program than the MLP, (x @ w).sin().sum(), because the joint form is easier to read when the whole thing fits on a screen.

Two naming conventions are worth learning here, because they show up in every aot printout you will ever read. primals are the forward inputs. tangents are the incoming gradients. In the joint graph they arrive as arguments to the same function, and the return carries the forward output and the input gradients together, with None in the slot for the input that did not require grad.

def forward(self, primals, tangents):
    primals_1, primals_2, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
    mm = torch.ops.aten.mm.default(primals_1, primals_2);  primals_2 = None
    sin = torch.ops.aten.sin.default(mm)
    sum_1 = torch.ops.aten.sum.default(sin);  sin = None
    expand = torch.ops.aten.expand.default(tangents_1, [2, 4]);  tangents_1 = None
    cos = torch.ops.aten.cos.default(mm);  mm = None
    mul = torch.ops.aten.mul.Tensor(expand, cos);  expand = cos = None
    t = torch.ops.aten.t.default(primals_1);  primals_1 = None
    mm_1 = torch.ops.aten.mm.default(t, mul);  t = mul = None
    return pytree.tree_unflatten([sum_1, None, mm_1], self._out_spec)

import torch
from functorch.compile import aot_function, make_boxed_func, default_partition

def nop(gm, inputs):
    return make_boxed_func(gm.forward)

def peek(joint_gm, joint_inputs, **kwargs):
    print(joint_gm.code.strip())
    return default_partition(joint_gm, joint_inputs, **kwargs)

def f(x, w):
    return (x @ w).sin().sum()

x = torch.randn(2, 3)
w = torch.randn(3, 4, requires_grad=True)
aot_function(f, fw_compiler=nop, bw_compiler=nop, partition_fn=peek)(x, w).backward()

# def forward(self, primals, tangents):
#     primals_1, primals_2, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
#     mm = torch.ops.aten.mm.default(primals_1, primals_2);  primals_2 = None
#     sin = torch.ops.aten.sin.default(mm)
#     sum_1 = torch.ops.aten.sum.default(sin);  sin = None
#     expand = torch.ops.aten.expand.default(tangents_1, [2, 4]);  tangents_1 = None
#     cos = torch.ops.aten.cos.default(mm);  mm = None
#     mul = torch.ops.aten.mul.Tensor(expand, cos);  expand = cos = None
#     t = torch.ops.aten.t.default(primals_1);  primals_1 = None
#     mm_1 = torch.ops.aten.mm.default(t, mul);  t = mul = None
#     return pytree.tree_unflatten([sum_1, None, mm_1], self._out_spec)

Where the cut falls is a choice, and PyTorch ships two answers to it. default_partition collects the operators between the forward inputs and the forward outputs, and its docstring names the consequence directly: the stashed tensors become the output of the generated forward graph. min_cut_rematerialization_partition opens with a different sentence, that the backward recomputes the forward, trading memory bandwidth against computation.

Run the same three-op chain through both and the difference is countable. Under the default cut, the forward returns four values and the backward takes four arguments. Under the min-cut, both numbers drop to two, and sin and cos reappear as the first two lines of the backward, recomputed from the input rather than carried across the boundary.

Which one you get depends on how you entered. aot_function defaults to default_partition. Inductor does not use that default: compile_fx.py wraps min_cut_rematerialization_partition in its own partition_fn and passes that down, so an ordinary torch.compile on the inductor backend is already recomputing rather than saving wherever the heuristic says it should.

This is activation checkpointing, decided per tensor by a solver instead of per block by you. It also explains a shape of profile that otherwise looks like a bug: a backward that runs more ops than the forward did, on a model you never wrote a checkpoint wrapper for.

default_partition: forward returns 4, backward takes 4
def forward(self, primals_1):
    sin = torch.ops.aten.sin.default(primals_1)
    cos = torch.ops.aten.cos.default(sin)
    sin_1 = torch.ops.aten.sin.default(cos)
    sum_1 = torch.ops.aten.sum.default(sin_1);  sin_1 = None
    return [sum_1, primals_1, sin, cos]
min_cut_rematerialization_partition: forward returns 2, backward takes 2
def forward(self, primals_1):
    sin = torch.ops.aten.sin.default(primals_1)
    cos = torch.ops.aten.cos.default(sin);  sin = None
    sin_1 = torch.ops.aten.sin.default(cos);  cos = None
    sum_1 = torch.ops.aten.sum.default(sin_1);  sin_1 = None
    return [sum_1, primals_1]

One more thing the printers tell you, and it is about time rather than structure. Swap the graph printing for a one-line marker in each compiler and the ordering is unambiguous: the forward compiler runs during the compiled call, the backward compiler does not run until .backward() is invoked.

So a first training step pays for two compilations at two different moments, and a benchmark that times only the forward will miss the second one entirely. Chapter 9's harness exists for exactly this class of measurement error, and this is one more thing its warmup has to cover.

The lazy backward also means a model you compile and only ever run under torch.no_grad() never builds a backward graph at all. Nothing warns you either way. The compiler simply never gets called.

before the forward call
[forward compiler ran]
forward call returned
[backward compiler ran]
backward returned