torch.compile: When One Line Speeds Up PyTorch — and When It Quietly Doesn't

If your PyTorch training loop feels slower than the hardware should allow, the usual advice is to rewrite hot paths, fuse operations by hand, or reach for a different framework. Since PyTorch 2.0 there's a lower-effort option: wrap your model in torch.compile() and let the framework capture your eager code into an optimized graph. The thesis of this post is simple: torch.compile is often worth trying first because it costs one line, but you have to measure it correctly and know the cases where it silently falls back to eager speed.

What torch.compile actually does

torch.compile(model) returns a wrapped version of your model. On the first call, a component called TorchDynamo traces your Python code and captures it as a graph of tensor operations. A backend compiler — Inductor by default — then fuses operations and generates optimized kernels for that graph. Fusion matters because eager PyTorch launches many small GPU kernels, each with overhead; a fused graph launches fewer, larger kernels.

Two consequences follow. First, compilation is lazy: the first execution pays the full tracing and codegen cost, so iteration one is much slower, not faster. Second, the compiled artifact is specialized to what it saw — including input shapes. Feed it a new shape and it recompiles (or falls back), which we will come back to.

A worked example you can run

Here is a minimal benchmark on a CUDA GPU. Run it in a normal Python environment with PyTorch 2.x installed (import torch; print(torch.__version__) — anything 1.x simply does not have this API). No special permissions are needed.

import time
import torch

print(torch.__version__)  # requires 2.x

model = torch.nn.Sequential(
    torch.nn.Linear(1024, 4096),
    torch.nn.GELU(),
    torch.nn.Linear(4096, 1024),
).cuda()

compiled = torch.compile(model)
x = torch.randn(256, 1024, device="cuda")

def bench(fn, iters=200, warmup=20):
    for _ in range(warmup):   # warmup absorbs compilation cost
        fn(x)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn(x)
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / iters

eager_ms = bench(model) * 1000
compiled_ms = bench(compiled) * 1000
print(f"eager: {eager_ms:.3f} ms  compiled: {compiled_ms:.3f} ms")

# Numerical check: fusion can change reduction order slightly
out_eager = model(x)
out_compiled = compiled(x)
print("outputs match:", torch.allclose(out_eager, out_compiled, atol=1e-5))

The critical details: the warmup loop exists because the first compiled call triggers tracing and codegen — benchmarking without it would make compilation look like a catastrophic slowdown. The torch.cuda.synchronize() calls matter too, because CUDA kernels are asynchronous and naive timing measures launch time, not execution time. Finally, torch.allclose with a tolerance confirms correctness: fused kernels can reorder floating-point reductions, so bitwise equality is not guaranteed and tiny differences are expected.

What speedup should you expect? It depends heavily on the model, batch size, GPU, and PyTorch version — small kernel-bound models like this MLP tend to benefit more than models already dominated by one huge matmul. Treat any multiplier you read online (including from your own run here) as specific to that setup, not a promise.

Choosing a mode

The mode argument trades compile time against runtime speed:

• "default" — a balanced choice; compile in seconds, decent speedup. Start here.
• "reduce-overhead" — uses CUDA graphs to cut per-call CPU launch overhead. Most useful for small batches where the GPU is starved by launch latency, at the cost of extra memory and some constraints on tensor reuse.
• "max-autotune" — spends a long time (sometimes many minutes) benchmarking alternative kernel implementations to pick the fastest. Reasonable for a fixed production inference shape; painful during iterative development.

A practical workflow: develop with default, then try reduce-overhead for latency-sensitive small-batch serving, and reserve max-autotune for shapes that will run millions of times.

The quiet failure mode: graph breaks

TorchDynamo cannot capture everything. Data-dependent control flow (e.g., if x.sum() > 0:), certain unsupported ops, and some custom autograd functions cause a graph break: Dynamo compiles what it can and runs the rest in eager mode. Your model still produces correct results — it just gets less speedup, and nothing warns you loudly.

Check before you celebrate. Run torch._dynamo.explain(model)(x) to see how many graphs were captured and where breaks occurred. If you see many breaks, the fix is usually to restructure the offending code (move data-dependent branching out of the forward pass) rather than abandon compilation.

The related trap is shape variability. If your batches have varying shapes — common in NLP with variable sequence lengths — each new shape can trigger a recompilation, and compile time can swamp your gains. Passing dynamic=True tells Dynamo to compile for symbolic shapes, accepting slightly less optimization in exchange for not recompiling per shape. Test this explicitly: run two batches of different shapes and watch whether wall time spikes on the second.

Trade-offs and limitations

Compilation interacts poorly with some distributed setups and exotic custom ops, so production adoption deserves a real test pass, not a one-line change on Friday. Backend maturity also differs: CUDA on NVIDIA GPUs is the most exercised path, while CPU and other accelerators work but with less consistent wins. And because outputs can differ slightly from eager mode, any pipeline that asserts exact reproducibility against a stored baseline needs tolerance-based comparisons instead.

What to do next

Take one real model from your codebase, wrap it with torch.compile, and run the benchmark pattern above: warmup excluded, CUDA synchronized, outputs checked with torch.allclose, and torch._dynamo.explain consulted for graph breaks. If the compiled run is meaningfully faster and numerically equivalent, keep it; if graph breaks eat the gains, you have learned that in an afternoon rather than after a rewrite. That cheap experiment is the real value of the feature.