torch.compile() — Graph Mode Without Giving Up Eager

The biggest single feature in PyTorch 2.x

torch.compile(model) is a one-line transformation that JIT-compiles your model into optimized fused kernels. Typical speedups: 1.5–3x on CUDA, often more for transformer-shaped models. The amazing part: you keep eager-mode debuggability everywhere outside the compiled region.

What's actually happening

TorchDynamo intercepts the Python bytecode of your forward, captures the operations into an FX graph.
AOTAutograd rewrites the graph to include backward operations.
TorchInductor lowers the graph to optimized Triton/CUDA kernels (or C++ for CPU).

The graph is built lazily on the first call. If your code does something Dynamo can't capture (data-dependent control flow, custom Python objects in forward), it gracefully falls back to eager for that section — a "graph break". The compile won't fail outright; it'll just lose some speedup.

Three modes

default — good balance of compile time and speedup.
mode="reduce-overhead" — minimizes Python overhead between kernel launches. Best for small models or small batch sizes.
mode="max-autotune" — exhaustively searches kernel variants for max speed. Slow to compile (sometimes minutes), fastest to run.

What breaks the graph

Data-dependent control flow on tensor values (if x.sum() > 0) — usually OK if it's on a Python int / config.
Calls into untraced libraries (some OpenCV / PIL ops in forward).
Mutating tensor values via Python attributes.
Some custom autograd Functions (improving in PyTorch 2.x).

Code

Compile a model — one line·python

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(1024, 4096), nn.GELU(),
            nn.Linear(4096, 4096), nn.GELU(),
            nn.Linear(4096, 1024),
        )
    def forward(self, x):
        return self.layers(x)

model = MLP().cuda()
model = torch.compile(model)               # done.

x = torch.randn(64, 1024, device='cuda')
y = model(x)                                # first call: slow (compiles)
y = model(x)                                # subsequent: fast

Picking the mode·python

import torch

# Default — best for most workloads
model = torch.compile(model)

# Reduce overhead — when batch is small / model is light
model = torch.compile(model, mode="reduce-overhead")

# Max autotune — exhaustive search; slow compile, fastest runtime
model = torch.compile(model, mode="max-autotune")

Detecting graph breaks — the diagnostic env var·python

import os
import torch

# Set BEFORE importing torch (or just before compiling)
os.environ['TORCH_LOGS'] = 'graph_breaks'

@torch.compile
def forward(x, mask):
    # This .item() call CAUSES a graph break — the value goes to Python
    if mask.sum().item() > 0:
        return x * 2
    return x * -1

x = torch.randn(8)
mask = torch.tensor([1, 0, 1, 0, 1, 0, 1, 0])
forward(x, mask)
# Logs will show: 'Graph break: tensor.item()'

# Fix: use torch.where for tensor-valued conditions
@torch.compile
def forward_fixed(x, mask):
    return torch.where(mask.bool(), x * 2, x * -1)

Compile inside training — drop in·python

import torch, torch.nn as nn
from torch.amp import autocast

model = MyModel().cuda()
model = torch.compile(model, mode="default")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for x, y in loader:
        x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
        optimizer.zero_grad(set_to_none=True)
        with autocast(device_type='cuda', dtype=torch.bfloat16):
            out = model(x)
            loss = criterion(out, y)
        loss.backward()
        optimizer.step()

# torch.compile and bf16 autocast compose cleanly. So does DDP/FSDP.

torch.compile() — Graph Mode Without Giving Up Eager

The biggest single feature in PyTorch 2.x

What's actually happening

Three modes

What breaks the graph

Code

External links

Exercise

Progress

Comments 0