PyTorch 2.0 Compile Mode: A Practical Guide

by

PyTorch 2.0 introduced torch.compile(), a one-line speedup for PyTorch models. It compiles your model’s computation graph using TorchDynamo and TorchInductor, fusing operations and generating optimised kernels for your hardware. On NVIDIA GPUs, compile mode typically delivers 10–50% speedup with no changes to model architecture or training loop. This guide covers how to use it, what the compilation modes do, when it helps most, and the pitfalls to avoid.

Basic Usage

import torch
import torch.nn as nn

model = MyModel().cuda()

# One line — that's it
compiled_model = torch.compile(model)

# Training loop unchanged
optimizer = torch.optim.AdamW(compiled_model.parameters(), lr=1e-4)
for batch in dataloader:
    x, y = batch
    loss = criterion(compiled_model(x), y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

The first few forward passes are slow — this is the compilation step where TorchDynamo traces the computation graph and TorchInductor generates optimised CUDA kernels. Subsequent passes run at full speed. The warm-up overhead is typically 30–120 seconds depending on model size; for long training runs this is negligible.

Compilation Modes

torch.compile() accepts a mode parameter that controls the trade-off between compilation time and runtime speed:

  • mode="default" — balanced. Good speedup with reasonable compilation time. Start here.
  • mode="reduce-overhead" — reduces Python overhead from repeated small kernel launches. Best for models with many small operations or high CPU-GPU synchronisation overhead.
  • mode="max-autotune" — tries many kernel configurations and picks the fastest. Compilation takes much longer (minutes) but produces the fastest runtime. Use for production inference where you compile once and run many times.
# For training (compilation speed matters)
model = torch.compile(model, mode="default")

# For inference in production (runtime speed matters most)
model = torch.compile(model, mode="max-autotune")

# For models with many small ops
model = torch.compile(model, mode="reduce-overhead")

Checking if Compile Helped

import time, torch

model = MyModel().cuda()
x = torch.randn(32, 3, 224, 224).cuda()

# Warmup
with torch.no_grad():
    for _ in range(5):
        _ = model(x)

# Benchmark eager
t0 = time.perf_counter()
with torch.no_grad():
    for _ in range(100):
        _ = model(x)
torch.cuda.synchronize()
eager_ms = (time.perf_counter() - t0) * 1000 / 100

# Benchmark compiled
compiled = torch.compile(model)
with torch.no_grad():
    for _ in range(5):   # warmup for compilation
        _ = compiled(x)

t0 = time.perf_counter()
with torch.no_grad():
    for _ in range(100):
        _ = compiled(x)
torch.cuda.synchronize()
compiled_ms = (time.perf_counter() - t0) * 1000 / 100

print(f"Eager:    {eager_ms:.2f}ms")
print(f"Compiled: {compiled_ms:.2f}ms")
print(f"Speedup:  {eager_ms/compiled_ms:.2f}x")

When Compile Mode Helps Most

Compile mode provides the largest speedups in these situations. Transformer and attention-based models benefit most — the fused attention kernels from FlashAttention integration in TorchInductor are particularly effective. Large batch sizes give the GPU more parallelism to exploit with the optimised kernels. Models with repeated identical operations (like transformer encoder layers) compile especially efficiently because the same pattern is optimised once and reused. Inference-only deployments where you compile once and serve millions of requests amortise the compilation cost to near zero.

Compile mode helps less — or not at all — for small models (overhead from kernel launches is negligible), models with dynamic shapes where the graph changes every forward pass (though dynamic=True handles this), and models that spend most time on data loading or CPU operations rather than GPU computation.

Dynamic Shapes

By default, torch.compile() recompiles for different input shapes, which is slow if your batch sizes or sequence lengths vary. Enable dynamic shapes to handle variable inputs without recompilation:

# Handle variable input shapes without recompiling
compiled = torch.compile(model, dynamic=True)

# Or use fullgraph for models that can be fully traced
compiled = torch.compile(model, fullgraph=True, dynamic=True)

Using torch.compile with Mixed Precision

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()
compiled = torch.compile(model)

for batch in dataloader:
    x, y = batch[0].cuda(), batch[1].cuda()
    optimizer.zero_grad()

    with autocast():
        out = compiled(x)
        loss = criterion(out, y)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Compile mode and mixed precision (AMP) are fully compatible and their speedups stack — compile mode optimises the computation graph, AMP reduces memory and arithmetic precision. Together they can deliver 2–4x total speedup over eager FP32 training on modern NVIDIA GPUs.

Figure 1 — torch.compile() pipeline: from Python model to optimised CUDA kernels

PyTorch Model (eager) TorchDynamo Graph capture TorchInductor Kernel generation Optimised CUDA kernels

Debugging Compile Issues

When torch.compile() produces wrong results or errors, set TORCH_LOGS="+dynamo" to see what the tracer is doing. The most common issues: graph breaks where TorchDynamo cannot trace through a Python control flow and falls back to eager mode for that section (reduces speedup but does not break correctness), unsupported operations that cause silent fallback to eager, and recompilation due to changing input shapes (fix with dynamic=True).

import os, torch

os.environ["TORCH_LOGS"] = "+dynamo"

# See graph breaks
compiled = torch.compile(model, fullgraph=False)  # default — allows graph breaks

# Strict mode — fails if there are any graph breaks (useful for debugging)
compiled = torch.compile(model, fullgraph=True)

# Disable for debugging
torch._dynamo.config.suppress_errors = True

Hardware Requirements

Full performance from torch.compile() requires NVIDIA GPUs with CUDA, PyTorch 2.0+, and a Linux or Windows environment (macOS support is limited). The best speedups occur on Ampere architecture (RTX 3000 series, A100) and newer where TorchInductor can target Triton kernels specifically. On older GPUs (pre-Turing), speedups are smaller but still present for large models. CPU compilation is supported via the backend="inductor" option, though speedups are typically more modest than GPU.

Adding torch.compile() to an existing PyTorch training script takes one line and typically delivers 10–50% speedup for free. The only cost is compilation time on the first few batches. For any training run longer than an hour, that tradeoff is clearly worth it without further analysis — a 20% training speedup on a 10-hour run saves two hours of GPU time on the very first run. For inference serving, use mode="max-autotune", compile once at startup, and the latency reduction at high throughput will be significant. Start with mode="default", measure, and move to max-autotune or reduce-overhead if you need more.

Figure 1 — torch.compile() pipeline: from Python model to optimised CUDA kernels

PyTorch Model (eager) TorchDynamo Graph capture TorchInductor Kernel generation Optimised CUDA kernels

Expected Speedups by Model Type

Realistic speedups on NVIDIA Ampere GPUs (A100, RTX 3090): CNN image models typically see 15–30% improvement; transformer models (BERT, GPT-style) see 20–50%, with the largest gains on longer sequence lengths where attention dominates; RNNs and LSTMs see modest gains of 5–15%; simple MLP models often see negligible improvement. The speedup also scales with batch size — larger batches give the optimised kernels more work to parallelise. If your model trains in short bursts with small batches, compile mode may not help much. The benchmark in the “Checking if Compile Helped” section above takes 5 minutes and gives you the actual number for your specific model and hardware.

torch.compile with torch.no_grad and Inference

For inference-only workflows, combining torch.compile() with torch.no_grad() and torch.inference_mode() gives the best performance. Inference mode disables gradient tracking and some autograd bookkeeping, allowing the compiler to generate more aggressive optimisations:

import torch

model = MyModel().cuda().eval()
compiled = torch.compile(model, mode="max-autotune")

# Fastest inference pattern
with torch.inference_mode():
    output = compiled(input_tensor)

# For serving with batched requests
@torch.inference_mode()
def predict(batch):
    return compiled(batch.cuda()).cpu()

Compile Mode with Custom CUDA Extensions

If your model uses custom CUDA extensions or third-party operators, they may cause graph breaks — TorchDynamo cannot trace through them and falls back to eager for those sections. Use torch._dynamo.allow_in_graph to mark trusted custom functions as traceable, or wrap them in torch.compiler.disable() to explicitly exclude them while compiling everything else:

import torch

# Option 1: exclude a function from compilation
@torch.compiler.disable
def my_custom_op(x):
    return custom_cuda_extension(x)

# Option 2: mark as safe to include in graph
@torch._dynamo.allow_in_graph
def my_safe_op(x):
    return fast_custom_kernel(x)

compiled = torch.compile(model)  # custom_op excluded, safe_op included

Saving Compiled Models

Compiled models cannot be saved directly with torch.save() — save the original uncompiled model and recompile at load time. The compiled artifact is hardware-specific and not portable across different GPU types:

# Save the original model, not the compiled wrapper
torch.save(model.state_dict(), 'model.pt')

# Load and recompile
model = MyModel()
model.load_state_dict(torch.load('model.pt'))
model = model.cuda().eval()
compiled = torch.compile(model, mode="max-autotune")  # recompile after loading

For deployment, this means the compilation overhead occurs at service startup rather than at training time — acceptable for long-running inference services but worth accounting for in cold-start latency budgets.

Compile Mode on Multiple GPUs

Compile mode works seamlessly with torch.nn.DataParallel and DistributedDataParallel. Compile the model before wrapping it in DDP for the best results — compiling after DDP can cause graph breaks on the DDP communication hooks:

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# Compile first, then wrap in DDP
model = torch.compile(model)
ddp_model = DDP(model, device_ids=[local_rank])

Common Compile Errors and Fixes

The most frequent issues when first enabling compile mode. RecursionError during tracing — your model has a deeply recursive structure that TorchDynamo cannot unroll; use fullgraph=False (default) to allow graph breaks or refactor the recursive calls. Numerical differences — rare but possible due to different floating-point operation ordering; verify with torch.allclose(eager_out, compiled_out, atol=1e-5) and if they diverge, try mode="default" instead of max-autotune. OOM during compilation — TorchInductor generates and JIT-compiles kernels in GPU memory; reduce mode to "default" or set torch._inductor.config.triton.cudagraphs=False. Slow recompilations — changing input shapes between batches triggers recompilation; fix with dynamic=True or pad all inputs to a fixed shape.

These issues are uncommon for standard model architectures — the majority of PyTorch training loops compile cleanly with no modifications. The cases where compilation causes problems are typically models with heavy Python-level control flow, custom C++/CUDA extensions that are not registered with Dynamo, or models that rely on side effects in the forward pass. For everything else, compile mode is a safe, easy performance upgrade. The PyTorch team has invested heavily in making graph breaks non-fatal and in expanding operator coverage — each release reduces the list of unsupported patterns, so a model that did not compile cleanly in PyTorch 2.0 often compiles without issues in 2.1 or 2.2. Staying on a recent PyTorch version is itself a form of compile-mode optimisation, as the compiler backend improves continuously between releases.

Adding torch.compile() to an existing PyTorch training script takes one line and typically delivers 10–50% speedup for free. The only cost is compilation time on the first few batches. For any training run longer than an hour, that is worth it without further analysis. For inference serving, use mode="max-autotune", compile once at startup, and the latency reduction at high throughput is significant. Start with mode="default", measure, and move to max-autotune or reduce-overhead if you need more.

Leave a Comment