Vue d'ensemble

PyTorch 2.2, publie le 31 janvier 2024, integre FlashAttention-2 et ameliore torch.compile pour de meilleures performances.

Fonctionnalites principales

FlashAttention-2

FlashAttention-2 est integre dans torch.nn.functional.scaled_dot_product_attention, accelerant l'entrainement des transformers.

python 
import torch
import torch.nn.functional as F

q = torch.randn(2, 8, 128, 64, device='cuda')
k = torch.randn(2, 8, 128, 64, device='cuda')
v = torch.randn(2, 8, 128, 64, device='cuda')

# FlashAttention-2 selectionne automatiquement
out = F.scaled_dot_product_attention(q, k, v)
print(out.shape)  # torch.Size([2, 8, 128, 64])

torch.compile ameliore

torch.compile supporte davantage d'operations et reduit les temps de compilation grace a un meilleur cache.

python 
import torch

model = torch.nn.Linear(256, 128).cuda()
compiled = torch.compile(model, mode='reduce-overhead')

x = torch.randn(32, 256, device='cuda')
out = compiled(x)  # premiere execution compile
out2 = compiled(x)  # executions suivantes rapides
print(out.shape)  # torch.Size([32, 128])

Sources