Résultats de benchmark

PyTorch 2.0 avec torch.compile accélère significativement l'entraînement et l'inférence. Sur les benchmarks HuggingFace, TorchBench et TIMM, les gains moyens atteignent 43 % en entraînement et 21 % en inférence sur GPU NVIDIA A100.

Mesurer les gains

python 
import torch
import time

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

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

# Avec 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'Compilé : {compile_time:.3f}s')
print(f'Gain : {eager_time / compile_time:.1f}x')

Sources