torch.compile

What it is

torch.compile (introduced in PyTorch 2.0) is the production graph-capture-and-compile path. It wraps a function or module, transparently captures graphs at the bytecode level via Dynamo, traces the joint forward + backward via AOT Autograd, and lowers to fused GPU/CPU kernels via Inductor.

Usage:

import torch

@torch.compile
def loss(model, x, y):
    return ((model(x) - y) ** 2).mean()

# or
model = torch.compile(model, mode="max-autotune", fullgraph=True)

torch.compile is non-invasive: the wrapped callable still accepts arbitrary Python, falls back to eager on unsupported constructs (with a warning), and re-uses the compiled artefact on subsequent calls when guards pass.

The full pipeline

graph LR
    Fn[Python fn / Module] -->|@torch.compile| Hook[Frame eval hook]
    Hook -->|Dynamo| FX1[FX graph]
    FX1 -->|AOT Autograd<br/>trace joint fwd+bwd| FX2[Joint FX graph]
    FX2 -->|partition| F[Fwd FX]
    FX2 -->|partition| B[Bwd FX]
    F -->|Inductor<br/>lowering + scheduling| Kf[Compiled fwd<br/>(Triton + C++)]
    B -->|Inductor| Kb[Compiled bwd]
    Kf & Kb -->|wrapped in autograd.Function| Compiled[Compiled callable]
    Compiled -->|on call, check guards| Hook

Each stage is documented separately:

• Dynamo — bytecode capture and guards.
• AOT Autograd — joint graph + partitioning.
• Inductor — IR, scheduling, codegen.
• FX — the IR shared by all of the above.

Modes

mode= controls the compile budget vs. perf trade-off:

Mode	What it does
"default"	Tuned for short compile time + good throughput.
"reduce-overhead"	Adds CUDA Graphs to amortize launch overhead. Best for small-batch inference.
"max-autotune"	Aggressive autotuning of matmul/conv/attention templates. Long compile, high perf.
"max-autotune-no-cudagraphs"	Same minus CUDA Graphs.

fullgraph=True causes Dynamo to raise on a graph break instead of falling back; useful for ensuring you actually got fully-compiled code.

dynamic=True/False/None controls dynamic-shape behaviour:

• False — specialize on every shape (recompile if shape changes).
• True — assume every dim is dynamic (recompile only on rank/dtype/device changes).
• None (default) — auto-detect: specialize first, mark dimensions dynamic on the second compilation when shapes differ.

Backends

Dynamo dispatches the captured graph to a backend. The default is inductor. Other built-ins (most for debugging):

• eager — run the FX graph eagerly.
• aot_eager — run through AOT autograd then eager.
• aot_eager_decomp_partition — like aot_eager, with decompositions.
• cudagraphs — wrap eager in CUDA Graphs.
• onnxrt — execute via ONNX Runtime.
• tvm, openxla, tensorrt, tensorrt_static — third-party.

Pick via torch.compile(..., backend="aot_eager").

Caching

Compiled artefacts are cached at multiple levels:

• In-process cache keyed by the guard set; this is the default.
• FX graph cache (torch._dynamo.config.fx_graph_cache) — disk-backed.
• Triton autotune cache — caches the best tile config per input shape.
• AOTAutograd cache — caches the partitioned forward+backward per traced graph.

Cache locations and keys are controlled via TORCHINDUCTOR_CACHE_DIR, TORCH_COMPILE_DEBUG, and torch._dynamo.config.

Graph breaks and recompilation

When Dynamo can't trace something, it graph-breaks and the unsupported region runs in eager. Common causes: side effects, exceptions, certain stdlib calls, opaque C extensions. Set TORCH_LOGS=graph_breaks to see every break with a Python traceback.

When a guard fails, Dynamo recompiles. Common causes: shape change (with dynamic=False), nn.Module structure change, global state mutation. Set TORCH_LOGS=recompiles to see every recompilation.

What you get

• Memory wins from rematerialization in AOT autograd's min-cut partitioner — typical 2–4x activation savings on transformers.
• Latency wins from fused Triton kernels that combine pointwise prologues, reductions, and pointwise epilogues into one kernel.
• CUDA Graphs in reduce-overhead mode for small-batch inference.
• Specialized templates for matmul / attention / convolution that can outperform stock cuBLAS on common shapes.

Debugging

A short toolkit:

Tool	Purpose
TORCH_LOGS=dynamo,recompiles,graph_breaks	Dynamo behaviour
TORCH_LOGS=aot_graphs,aot_joint_graph	AOT autograd graphs
TORCH_LOGS=inductor,output_code	Inductor scheduling and final kernels
TORCH_COMPILE_DEBUG=1	Dump everything to torch_compile_debug/
torch._dynamo.explain(fn)(...)	Concise summary of graph breaks
torch._dynamo.repro.after_dynamo/after_aot	Auto-minified repros for compile bugs

Where to read next

• The systems pages (linked above) for the implementation.
• Features / Distributed training for compiling distributed code.
• Features / torch.export and AOTInductor for the export-and-deploy path that shares most of the same machinery.