Benchmark Results

PyTorch 2.0 with torch.compile significantly accelerates training and inference. On HuggingFace, TorchBench, and TIMM benchmarks, average gains reach 43% in training and 21% in inference on NVIDIA A100 GPUs.

Measuring Gains

python 
import torch
import time

model = torch.nn.Linear(4096, 4096).cuda()
x = torch.randn(256, 4096, device='cuda')

# Without compilation
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(1000):
    _ = model(x)
torch.cuda.synchronize()
eager_time = time.perf_counter() - t0

# With compilation
compiled = torch.compile(model)
_ = compiled(x)  # warm-up
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(1000):
    _ = compiled(x)
torch.cuda.synchronize()
compile_time = time.perf_counter() - t0

print(f'Eager: {eager_time:.3f}s')
print(f'Compiled: {compile_time:.3f}s')
print(f'Speedup: {eager_time / compile_time:.1f}x')

Sources