Vue d'ensemble

PyTorch 2.6, publie le 30 janvier 2025, ameliore torch.compile et les performances d'entrainement.

Fonctionnalites principales

torch.compile ameliore

torch.compile supporte davantage d'operations et reduit le temps de compilation initiale. Le mode reduce-overhead est plus stable.

python 
import torch

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

# Compilation pour acceleration
compiled = torch.compile(model, mode='reduce-overhead')
x = torch.randn(32, 784)
out = compiled(x)
print(f'Output shape: {out.shape}')

Performances ameliorees

Les operations CUDA sont optimisees avec un meilleur support de FlashAttention et des noyaux fuses pour les Transformers.

python 
import torch
import torch.nn.functional as F

# Scaled dot-product attention (FlashAttention)
q = torch.randn(2, 8, 128, 64)
k = torch.randn(2, 8, 128, 64)
v = torch.randn(2, 8, 128, 64)

out = F.scaled_dot_product_attention(q, k, v)
print(f'Attention shape: {out.shape}')

Sources