Overview

PyTorch 2.0, released on March 16, 2023, is a major release introducing torch.compile() for automatic graph optimization, delivering 30-200% speedups with no code changes required.

Main Features

torch.compile()

torch.compile() automatically captures and optimizes the PyTorch computation graph. It analyzes Python code and generates optimized GPU code without changing the existing API.

python 
import torch

model = torch.nn.Sequential(
    torch.nn.Linear(784, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10),
)

# Automatic model compilation
compiled_model = torch.compile(model)

# Same usage, improved performance
x = torch.randn(32, 784)
y = compiled_model(x)  # 30-200% faster

Graph optimization

The new TorchDynamo backend captures the computation graph at the Python level and passes it to TorchInductor for optimized code generation (Triton for GPU, C++ for CPU).

python 
import torch

@torch.compile(mode='reduce-overhead')
def train_step(model, x, y, optimizer, loss_fn):
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return loss

# Available modes:
# 'default'         - good trade-off
# 'reduce-overhead' - minimize latency
# 'max-autotune'    - maximize optimization

Performance gains

Benchmarks show significant speedups across various models: 43% average on HuggingFace models, 46% on TorchBench, and 26% on TIMM models, all with zero code changes.

python 
import torch
import time

model = torch.nn.Transformer(
    d_model=512, nhead=8, num_encoder_layers=6
)
x = torch.randn(10, 32, 512)

# Without compilation
start = time.time()
for _ in range(100):
    model(x, x)
print(f'Without compile: {time.time() - start:.2f}s')

# With compilation
model_c = torch.compile(model)
model_c(x, x)  # warm-up
start = time.time()
for _ in range(100):
    model_c(x, x)
print(f'With compile: {time.time() - start:.2f}s')

Sources