torch.compile Deep Dive

PyTorch 2.0 introduces torch.compile(), which transforms an eager model into an optimized graph via TorchDynamo (capture) and TorchInductor (code generation). A single call yields 30-200% speedups without modifying model code.

Usage

python 
import torch

model = torch.nn.TransformerEncoderLayer(
    d_model=512, nhead=8, batch_first=True
)
model = model.cuda()

# Compilation: one line
compiled = torch.compile(model, mode='reduce-overhead')

# Available modes:
# 'default'         : good trade-off
# 'reduce-overhead' : minimize CPU overhead
# 'max-autotune'    : try more CUDA configs

x = torch.randn(32, 128, 512, device='cuda')
out = compiled(x)  # first call: compilation
out = compiled(x)  # subsequent calls: fast
print(out.shape)  # torch.Size([32, 128, 512])

Sources