torch.compile and TorchInductor

PyTorch 2.0 shipped a single function that, on average, made 165 open-source models run 20% faster at float32 and 36% faster under AMP - with no changes to the model code. That function is torch.compile. Understanding why it works requires following the graph from Python bytecode down to machine code.

The compilation pipeline in four layers

torch.compile is not one tool but a stack of four cooperating components:

Python source
    │
    ▼
TorchDynamo      (trace and capture the graph)
    │
    ▼
AOT Autograd     (capture the backward pass ahead-of-time)
    │
    ▼
Compiler backend (default: TorchInductor)
    │
    ▼
Triton / C++     (generated kernel code)

TorchDynamo sits at the CPython level. It hooks into the Frame Evaluation API (PEP 523), intercepts bytecode at the point of function entry, and symbolically executes operations to record a torch.fx graph. The key design choice: Dynamo does not require the user to write tracing-friendly code. If it encounters something it cannot capture (a C extension call, data-dependent branching, a print), it falls back gracefully and inserts a "graph break" - execution resumes in eager mode for that fragment.

AOT Autograd then captures both the forward and backward passes ahead of time, before any data flows through. This means the backward graph is also available to the compiler, enabling cross-pass fusion that is impossible in classic eager mode.

TorchInductor is the default backend. It lowers the torch.fx graph to a loop-level intermediate representation (loop IR), then emits either Triton kernels for CUDA/ROCm/Intel GPUs or vectorised C++ for CPU. The critical optimisation at this stage is operator fusion: instead of writing activations back to DRAM after every elementwise op, Inductor recognises fuseable chains and merges them into a single Triton kernel.

Why fusion matters so much

A na??ve PyTorch forward pass through a transformer block calls dozens of separate kernels. Each one: 1. Reads its input tensors from DRAM. 2. Computes. 3. Writes its output back to DRAM.

For elementwise chains (LayerNorm -> dropout -> residual add), the compute is trivially cheap relative to the memory roundtrips. The roofline model makes this concrete: a modern A100 has ~312 TFLOP/s of compute but only ~2 TB/s of memory bandwidth. An unfused four-operation chain that each move, say, 1 GB through DRAM consumes 4 x (1 GB / 2 TB/s) = 2 ms in memory time alone, even if the arithmetic finishes in microseconds.

TorchInductor recognises the fuseable subgraph and emits a single Triton kernel that keeps intermediate values in registers or shared memory. The same workload now does one DRAM read and one DRAM write. Latency drops by roughly 4x for that subgraph.

How TorchDynamo handles dynamic shapes

By default, Dynamo compiles a specialised version of the graph for the exact input shapes it first sees, using those shapes as guards. If the shapes change (a different batch size, a variable sequence length), the guards fail, Dynamo recompiles, and the new compiled version is cached.

To avoid recompilation storms in production, pass dynamic=True:

model = torch.compile(model, dynamic=True)

This tells Dynamo to track shapes symbolically where possible, producing a single compiled artifact valid across a range of inputs. The tradeoff: symbolic shape reasoning is harder, and Inductor may produce slightly less optimised kernels when shapes are symbolic because certain tiling decisions become conservative.

A middle path is torch.compiler.mark_dynamic(tensor, dim) to tag only the dimensions you know will vary, letting all others remain specialised.

Using torch.compile in practice

import torch

model = MyTransformer().cuda()

# Minimal: wraps the entire model
compiled = torch.compile(model)

# With options
compiled = torch.compile(
    model,
    backend="inductor",   # default; alternatives: "eager", "aot_eager", "cudagraphs"
    mode="reduce-overhead",  # trades compile time for lower kernel launch overhead
    fullgraph=False,         # allow graph breaks (default); set True to error on them
)

# From PyTorch 2.0+: works with torch.autocast and torch.no_grad transparently
with torch.autocast("cuda"):
    loss = compiled(x).sum()
    loss.backward()

The mode parameter controls the optimisation/compile-time tradeoff:

• "default": balanced; good starting point.
• "reduce-overhead": enables CUDA Graphs under the hood, cutting kernel launch latency. Best for small batch sizes where launch overhead dominates.
• "max-autotune": runs Inductor's autotuning loop to pick the best tile sizes. Compile time can be minutes; inference time is minimal.

When it falls down

Graph breaks kill the benefit. Every graph break means the captured subgraph is smaller, and fusion opportunities across the break boundary are lost. The most common causes are: calling .item() on a tensor (forces a CPU synchronisation), data-dependent control flow (if tensor > 0:), and using Python objects that Dynamo cannot trace through (certain custom nn.Module subclasses, tqdm wrappers, print inside forward). Run torch._dynamo.explain(model)(x) to enumerate breaks before committing to compilation.

First-call latency. Compilation is deferred to the first forward pass. A model with complex shapes or many ops can take tens of seconds or more to compile. In serving, this manifests as a severe cold-start spike. Mitigate with torch.export + torch._inductor.aot_compile to produce a pre-compiled artifact that loads instantly.

Distributed training friction. DDP wraps the model in communication collectives that Dynamo cannot always trace cleanly. The recommended pattern is to compile before wrapping with DDP: model = torch.compile(model); model = DDP(model). Even so, the communication kernels themselves are not fused or optimised by Inductor - only the compute subgraphs between AllReduces benefit.

Numerical divergence. Inductor's fused Triton kernels can produce results that differ from eager mode at the ULP level. This is usually benign, but training with very tight convergence criteria (certain reinforcement learning setups, for example) can diverge. Use torch.compile(model, backend="aot_eager") to isolate whether a divergence is a graph-capture bug or a codegen numerics issue.

Dynamic shapes compile cost. Enabling dynamic=True shifts some shape decisions to runtime. Triton's autotuner may run multiple kernel configurations on each new shape, causing latency spikes on the first call per shape. For latency-sensitive inference, profile with the exact shape distribution you expect before enabling dynamic shapes.

Memory pressure. Compiled models, especially with max-autotune, cache multiple compiled versions plus Triton-compiled PTX. On memory-constrained GPUs, this can cause OOMs that do not appear in eager mode.

Further reading

• torch.compiler API documentation (PyTorch 2.7)
• PyTorch 2.0 release post with benchmark details
• Triton language and compiler documentation
• torch.compile FAQ: graph breaks, debugging, known issues