AOT Autograd

Active contributors: bdhirsh, aorenste

Purpose

AOT Autograd is the middle layer of the torch.compile stack. It takes a forward FX graph (produced by Dynamo or by a manual tracer), runs it through Functionalize + Autograd + FakeTensor to pre-trace the joint forward + backward graph, then partitions that joint graph into a forward and a backward, and hands those to a downstream compiler (Inductor by default).

Without AOT Autograd, torch.compile would only be able to compile inference; with it, the entire training step (including the backward pass) compiles into a single graph.

The package lives at torch/_functorch/_aot_autograd/ (and its top-level entry torch/_functorch/aot_autograd.py). It grew out of the same functorch transforms that implement vmap/grad.

Directory layout

Path	Contents
torch/_functorch/aot_autograd.py	Top-level entry: aot_module, aot_function
torch/_functorch/_aot_autograd/	Implementation
torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py	Build joint graph and dispatch to compiler
torch/_functorch/_aot_autograd/runtime_wrappers.py	Compiled artefact wrapper that re-injects autograd
torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py	Caching & inlining of compiled artefacts
torch/_functorch/_aot_autograd/input_output_analysis.py	Detection of view, mutation, and aliasing patterns
torch/_functorch/_aot_autograd/utils.py	Helpers
torch/_functorch/partitioners.py	Forward/backward graph splitting policies
torch/_subclasses/fake_tensor.py	FakeTensor (used heavily here)
torch/_subclasses/functional_tensor.py	Python functional-tensor wrapper
aten/src/ATen/FunctionalTensorWrapper.cpp	C++ functionalize key

Key abstractions

Type	File	Purpose
aot_function / aot_module	torch/_functorch/aot_autograd.py	User-facing API
AOTConfig	torch/_functorch/_aot_autograd/utils.py	Compile config
FakeTensorMode	torch/_subclasses/fake_tensor.py	Dispatch mode that replaces real ops with shape-only ops
FunctionalTensorMode	torch/_subclasses/functional_tensor.py	Dispatch mode that functionalizes view/mutation
min_cut_rematerialization_partition	torch/_functorch/partitioners.py	Default fwd/bwd graph partitioning
ViewMutationMetadata	torch/_functorch/_aot_autograd/input_output_analysis.py	Captures view/aliasing/mutation patterns
CompiledRuntimeWrapper	torch/_functorch/_aot_autograd/runtime_wrappers.py	Wraps the compiled fwd/bwd into an autograd.Function

How it works

The trace

Given a forward function f(args) and example inputs:

1. Wrap inputs in FakeTensors — every tensor becomes a FakeTensor carrying shape/dtype/device but no real data. This is essential: we want to symbolically execute through Python without doing actual compute, and FakeTensor provides shape semantics for every ATen op via meta kernels.
2. Push a FunctionalTensorMode — every view and in-place op is intercepted and replaced with a pure functional equivalent. The Functionalize dispatch key (in aten/src/ATen/FunctionalTensorWrapper.cpp) does the heavy lifting; the Python mode wraps that for tracing.
3. Push a ProxyTorchDispatchMode — every dispatcher call is recorded as an FX node.
4. Run the function with autograd enabled — the Autograd* keys are still active, so they record their own backward Nodes onto the tape.
5. Call .backward() on the outputs — this drives the backward dispatch of every recorded autograd op, which re-enters the same proxy mode and gets recorded into the same FX graph.

The result is a joint FX graph where forward and backward live side by side. Tangent arguments (the gradient inputs to backward) are explicit graph inputs.

The partition

torch/_functorch/partitioners.py has two main strategies:

• default_partition — cuts the joint graph at the outputs of the forward; saved tensors are everything that crosses the cut.
• min_cut_rematerialization_partition — solves a min-cut problem to minimize saved-tensor memory by rematerializing cheap ops in backward instead of saving their outputs. This is the default for torch.compile and a major contributor to its memory wins.

The partitioner emits a forward FX graph (returns outputs + saved values) and a backward FX graph (takes saved values + tangents → input grads).

Compile + runtime wrap

Each of the two graphs is handed to the configured backend (Inductor by default). The compiled artefacts are wrapped in a synthetic torch.autograd.Function whose forward calls the compiled forward and stores the saved values in ctx, and whose backward calls the compiled backward. From the autograd engine's perspective, it just sees an opaque Function — so AOT-compiled regions compose naturally with eager autograd outside the region.

graph LR
    DynGraph[Dynamo FX graph] --> Trace[trace_joint<br/>Fake + Functionalize + Proxy]
    Trace -->|joint FX graph| Part[Partitioner<br/>fwd + bwd]
    Part -->|fwd graph| BackendF[Backend compile<br/>(Inductor)]
    Part -->|bwd graph| BackendB[Backend compile<br/>(Inductor)]
    BackendF & BackendB --> Wrap[CompiledRuntimeWrapper<br/>autograd.Function]
    Wrap --> User[User calls compiled fn]

Mutation, views, and aliasing

A surprising amount of AOT Autograd's complexity is in handling aliasing and mutation correctly when the user passes in or mutates real tensors. input_output_analysis.py classifies each input/output into categories: input-mutated, output-as-view-of-input, output-as-view-of-intermediate, etc. The runtime wrappers then re-apply these effects on the real inputs after the compiled (functional) graph runs.

This is the bookkeeping that lets torch.compile accept code that does x.add_(1) or return x[:, 1:] without the user noticing anything different.

Caching

jit_compile_runtime_wrappers.py implements an in-process cache keyed by guard set. PyTorch 2.x also includes optional persistent caches (FX graph cache, Triton autotune cache, AOTAutograd cache) controlled by env vars / torch._dynamo.config.

Integration points

• Dynamo sends in the forward FX graph + sample inputs.
• FakeTensor / FunctionalTensor are required infrastructure; meta kernels in torch/_meta_registrations.py and torch/_decomp/ keep the symbolic execution working as the op set evolves.
• Inductor (or any backend with the right interface) consumes the partitioned graphs.
• Higher-order ops are passed through with their subgraphs intact; HOP-specific tracing logic lives in torch/_higher_order_ops/.

Entry points for modification

• A new partitioning strategy → add to partitioners.py and select via aot_config.partition_fn.
• New decomposition rules (replace an ATen op with a sub-FX graph) → register in torch/_decomp/decompositions.py.
• New runtime mutation/view support → runtime_wrappers.py and input_output_analysis.py.
• For debugging, set TORCH_LOGS=aot_graphs,aot_joint_graph and TORCH_COMPILE_DEBUG=1 to dump the joint and partitioned graphs.

Key source files

File	Purpose
torch/_functorch/aot_autograd.py	Top-level entry
torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py	Joint graph build + dispatch
torch/_functorch/_aot_autograd/runtime_wrappers.py	Wrap compiled artefacts as autograd.Function
torch/_functorch/partitioners.py	fwd/bwd partitioning
torch/_subclasses/fake_tensor.py	FakeTensor
torch/_subclasses/functional_tensor.py	FunctionalTensor
torch/_decomp/decompositions.py	ATen op decompositions
torch/_meta_registrations.py	Meta kernels for shape inference