Skip to content

torch.compile Deep Dive

Overview

torch.compile is the default path to speed: it wraps a module, captures the computation with Dynamo (a graph-tracing compiler front-end), and lowers it with Inductor (a Triton/C++ kernel generator). On modern GPUs it commonly delivers 1.5-2x training and 2-4x inference speedups — sometimes with zero code changes.

  • torch.compile(model) → cached compiled artifact; subsequent calls hit compiled kernels.
  • Dynamo captures frames (Python bytecode) into FX-style graphs; unsupported parts get "graph breaks".
  • Inductor (default backend) generates fused Triton kernels and CUDA graphs.
  • Modes: default, reduce-overhead (CUDA graphs), max-autotune, max-autotune-no-cudagraphs.
  • Options: fullgraph=True (no graph breaks), dynamic=True, backend=....

💡 The first call is slow (compilation + autotuning). Measure steady-state throughput, not the first iteration.

The Simplest Win

import torch, torch.nn as nn, time

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(*[nn.Linear(256, 256) for _ in range(8)])
    def forward(self, x):
        return self.layers(x).relu()

model = MLP().cuda()
compiled = torch.compile(model)      # returns a wrapped module

x = torch.randn(128, 256, device='cuda')

def bench(fn, n=20):
    for _ in range(3): fn()          # warm-up (compilation happens here)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n): fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / n * 1e3

print("eager  : %.2f ms" % bench(lambda: model(x)))
print("compile: %.2f ms" % bench(lambda: compiled(x)))

How It Works (the two-stage machine)

Python forward
   │  Dynamo: trace bytecode, build FX graph
   │  └─ graph breaks on unsupported control flow -> pieced graphs
   ▼
FX graph
   │  Inductor: generate Triton kernels, fuse, autotune, cudagraphs
   ▼
compiled kernels  (cached per (code, shapes, device, options))

TORCH_LOGS for transparency

TORCH_LOGS=graph_breaks python train.py     # see where graphs break
TORCH_LOGS=inductor python train.py         # see generated code
TORCH_COMPILE_DEBUG=1 python train.py       # dump artifacts

Modes & Backends

Mode What it adds When
default Inductor kernels most cases
reduce-overhead CUDA graph capture small models / launch-bound
max-autotune wider autotune + cudagraphs inference / fixed shapes
max-autotune-no-cudagraphs autotune w/o capture dynamic shapes 
compiled_fast = torch.compile(model, mode="reduce-overhead")
compiled_max  = torch.compile(model, mode="max-autotune", fullgraph=True)

⚠️ max-autotune compiles much longer; only worth it when warm shapes are stable.

What Compilation Actually Fixes

  1. Python overhead — hundreds of tiny scalar ops in forward disappear.
  2. Kernel launch overhead — fused regions = fewer launches.
  3. Memory bandwidth waste — fusion keeps intermediates in registers/SRAM.
  4. Redundant copies/layout conversions — Inductor picks layouts.

It does NOT fix: bad algorithms, unavoidable memory traffic, or CPU-bound data loaders.

Graph Breaks — The Silent Killer

def bad_forward(self, x):
    if x.sum() > 0:        # tensor data check -> graph break!
        return self.a(x)
    return self.b(x)

Every if tensor, for over Python range with data, len(tensor), print(tensor), or arbitrary python lib call can break the graph. Breaks leave fragments compiled separately → overhead partially returns.

Fixing the usual suspects

# 1. Prefer vectorized ops instead of data-dependent branching
mask = (x > 0).float()
out = mask * self.a(x) + (1 - mask) * self.b(x)

# 2. Use torch ops, not python control flow, where possible
# 3. Set fullgraph=True to *fail fast* and reveal breaks
torch.compile(model, fullgraph=True)

Dynamic Shapes

torch.compile(model, dynamic=True)          # allow varying batch sizes
torch.compile(model, dynamic=False)         # static: best for fixed shapes
  • Dynamic shapes trigger recompiles; guard with dynamic=True if batch varies.
  • torch._dynamo.config.cache_size_limit controls recompile budget.

When NOT to Use torch.compile

Situation Why
Tiny models (launch-bound) compile overhead > gain
Highly dynamic control flow breaks everywhere
One-off inference compilation cost per run
Ops with no Inductor kernels falls back to eager anyway
Debugging numerics compiled path hides eager semantics

Classic rule: measure. torch.compile is 90% free win, but it is not 100%.

Interop with the Rest of the Stack

  • Works with DDP/FSDP (torch.compile + distributed is the default path now).
  • Works with AMP (torch.amp.autocast) — compile after wrapping.
  • Export: torch.compile is for speed, torch.export for deployment (Ch 07).
  • Custom ops: fusable via torch.ops / inductor pattern matches.

Key Takeaways

  • torch.compile(model) = Dynamo capture + Inductor codegen; measure steady state.
  • Modes trade compile time for runtime: start default, escalate to reduce-overhead/max-autotune.
  • Kill graph breaks (data-dependent control flow) with vectorization or fullgraph=True checks.
  • Classic wins: python overhead, launches, bandwidth; not algorithmic changes.
  • Always benchmark on your shapes/devices before assuming.