torch.compile has been the headline feature of PyTorch since 2.0 shipped in March 2023, and by 2026 it is mature enough that not using it on a production model is a decision you should be able to justify. This tutorial measures what it actually buys you on two common model shapes, shows how to find out why it is not helping when it is not, and gives you a reproducible script to run on your own hardware.
Everything below was written against PyTorch 2.x on an NVIDIA GPU. Numbers are deliberately not quoted as headline figures: they vary by GPU generation, model, batch size and driver, and the point of the script is that you produce your own.
What torch.compile does, in one paragraph
torch.compile wraps a module or function. On the first call TorchDynamo traces the Python bytecode into an FX graph of tensor operations; TorchInductor then generates fused kernels for that graph (Triton kernels on GPU, C++ on CPU). Subsequent calls with compatible inputs run the compiled artifact instead of the Python. The win comes from kernel fusion (fewer memory round trips) and from removing Python overhead between ops. That is the whole mental model you need; the official docs cover the internals if you want them.
The benchmark harness
We measure eager versus compiled on a ResNet-50 and a small transformer encoder. Install torch and torchvision (a current 2.x release) and save this as bench.py.
import time
import torch
import torch.nn as nn
import torchvision
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
device = "cuda"
def small_transformer():
layer = nn.TransformerEncoderLayer(
d_model=512, nhead=8, dim_feedforward=2048, batch_first=True
)
return nn.TransformerEncoder(layer, num_layers=6)
def timeit(fn, *args, warmup=10, iters=50):
for _ in range(warmup):
fn(*args)
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(iters):
fn(*args)
torch.cuda.synchronize()
return (time.perf_counter() - start) / iters * 1000 # ms
@torch.no_grad()
def run(name, model, example, dtype=torch.bfloat16):
model = model.to(device, dtype).eval()
example = example.to(device, dtype)
eager = timeit(model, example)
compiled = torch.compile(model, mode="max-autotune-no-cudagraphs")
comp = timeit(compiled, example)
print(f"{name:20s} eager {eager:7.2f} ms compiled {comp:7.2f} ms "
f"speedup {eager / comp:4.2f}x")
if __name__ == "__main__":
for bs in (1, 32):
run(f"resnet50 bs={bs}", torchvision.models.resnet50(),
torch.randn(bs, 3, 224, 224))
run(f"transformer bs={bs}", small_transformer(),
torch.randn(bs, 128, 512))
Three details matter. The warmup loop absorbs compilation time, which can be tens of seconds on the first call; never include it in a latency number. torch.cuda.synchronize() is required because kernel launches are asynchronous. And BF16 is the realistic production dtype on any GPU from Ampere onward; measuring FP32 flatters eager mode because both paths are memory-bound.
What you will typically see
Run python bench.py and you will get a table like the one below. The shape of the result is consistent across hardware even though the magnitudes are not.
| Model | Batch | Eager | Compiled | Typical speedup |
|---|---|---|---|---|
| ResNet-50 | 1 | Python-overhead bound | launch overhead removed | 1.3-2x |
| ResNet-50 | 32 | GPU-bound | fused convs/BN/ReLU | 1.1-1.4x |
| Transformer | 1 | Python-overhead bound | fused attention/MLP | 1.5-2.5x |
| Transformer | 32 | GPU-bound | fused MLP, SDPA | 1.2-1.6x |
Small batches gain the most because eager mode is spending its time in Python and kernel launches, both of which compilation removes. Large batches gain less because the GPU was already the bottleneck; the remaining gain is from fusion.
For training add mode="max-autotune" only once the shapes are stable, and compile the model rather than the whole training step unless your optimizer step is also a hot spot. Compiling the loss and backward is automatic: AOTAutograd traces both.
Finding graph breaks
When the speedup is disappointing the usual cause is a graph break: a point where Dynamo could not trace and fell back to eager, splitting your model into several small compiled regions with Python between them. Ask the compiler to explain itself:
explanation = torch._dynamo.explain(model)(example)
print(explanation.graph_break_count)
for reason in explanation.break_reasons:
print(reason)
Common causes and fixes:
- Data-dependent Python control flow (
if tensor.item() > 0). Replace withtorch.whereor restructure; or accept the break if it is outside the hot path. - Printing, logging, or
.numpy()insideforward. Move it out or guard it behind a flag that is false in production. - Unsupported third-party ops. Wrap them with
torch.compiler.disableso the break is explicit and small, or register a custom op. - Python generators and unusual containers. Convert to lists or tensors before the forward.
Set TORCH_LOGS="graph_breaks" in the environment to print breaks as they happen, and TORCH_LOGS="recompiles" for the next problem.
Dynamic shapes and recompilation
Dynamo guards on input shapes. If your batch size or sequence length changes, the first new shape triggers a recompile; after a few, Dynamo marks the dimension dynamic automatically. If you know a dimension varies, say so up front to avoid the stall:
compiled = torch.compile(model, dynamic=True)
# or per-tensor:
torch._dynamo.mark_dynamic(example, 0) # batch dimension
Dynamic kernels are slightly slower than specialized ones, so for serving with a fixed set of shapes it is often better to pad to a few buckets and let each bucket compile once. Watch the recompile log in staging; a model that recompiles in production is slower than eager.
When compile does not pay
Be honest about the cases where you should not bother:
- Tiny models at small batch where the whole forward takes under a millisecond; compile overhead per call can exceed the work.
- Heavy Python in the forward that cannot be traced and produces dozens of graph breaks.
- Highly dynamic architectures (mixture-of-experts routing with Python loops, variable graph structure per sample).
- Short-lived processes. A batch job that runs the model once will never amortize the compile time. Use the compile cache (
TORCHINDUCTOR_CACHE_DIR) to share artifacts across runs if the process restarts frequently. - CPU inference on older toolchains where Inductor's C++ backend may not have a suitable compiler available.
Putting it in production
Compile once at process start, behind a warmup request that exercises the shapes you will serve. Keep mode="max-autotune-no-cudagraphs" unless you have verified CUDA graphs are safe for your input handling, since CUDA graphs require static memory addresses. Pin the PyTorch version; compiled artifacts are not portable across releases. And keep the eager path one environment variable away so you can compare when something looks wrong.
If you want these numbers on your own model, with the graph breaks found and fixed, that is exactly the work of our performance optimization practice. Contact us to talk it through.