Topolearntorch.compile often gives 1.3-2x speedups on training and inference with no model rewrite.
torch.compile(model) returns an optimised wrapper around a module or function. It works in layers:
- TorchDynamo hooks into Python bytecode execution and captures tensor operations into an FX graph, guarded by checks on shapes, dtypes and Python values.
- AOTAutograd traces the backward pass ahead of time.
- TorchInductor (the default backend) generates fused kernels: Triton on GPUs, C++/OpenMP on CPUs.
The first call is slow because it compiles; later calls reuse the cached code while the guards still hold.
Things to know:
- Graph breaks: code Dynamo cannot trace (some third-party calls, data-dependent Python control flow,
printof tensor values) splits the graph and falls back to Python, costing performance.fullgraph=Trueturns breaks into errors so you can find them. - Recompilation: a change in a guarded property, such as a new input shape, triggers a recompile. After a shape changes, Dynamo marks that dimension dynamic; you can also pass
dynamic=Trueor calltorch._dynamo.mark_dynamic. - Modes:
"reduce-overhead"uses CUDA graphs for small batches,"max-autotune"benchmarks kernel variants (slow compile, fastest run).
Debug with TORCH_LOGS="graph_breaks,recompiles". The compiled module shares parameters with the original, so the optimizer and state_dict work as before (keys gain an _orig_mod. prefix).
Example
import torch
from torch import nn
model = nn.Sequential(nn.Linear(256, 256), nn.GELU(), nn.Linear(256, 10))
compiled = torch.compile(model) # same parameters, optimised execution
x = torch.randn(32, 256)
out = compiled(x) # first call: trace + compile (slow)
out = compiled(x) # later calls: cached kernels
print(torch.allclose(out, model(x), atol=1e-5)) # True
@torch.compile(fullgraph=True) # error instead of a silent graph break
def scaled_gelu(x, s):
return nn.functional.gelu(x) * s
print(scaled_gelu(torch.randn(4), 2.0).shape) # torch.Size([4])
# Find graph breaks and recompiles from the shell:
# TORCH_LOGS="graph_breaks,recompiles" python train.pyRemember
- Dynamo captures graphs, Inductor generates fused kernels
- First call compiles; guards decide when to recompile
- Use
fullgraph=Trueand TORCH_LOGS to hunt graph breaks
Watch out for
- Benchmarking the first call, which includes compilation time.
- Feeding constantly changing shapes, causing repeated recompiles.