Overview

PyTorch 2.2, released on January 31, 2024, integrates FlashAttention-2 and improves torch.compile for better performance.

Main Features

FlashAttention-2

FlashAttention-2 is integrated into torch.nn.functional.scaled_dot_product_attention, accelerating transformer training.

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 automatically selected
out = F.scaled_dot_product_attention(q, k, v)
print(out.shape)  # torch.Size([2, 8, 128, 64])

Improved torch.compile

torch.compile supports more operations and reduces compilation times through better caching.

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)  # first run compiles
out2 = compiled(x)  # subsequent runs fast
print(out.shape)  # torch.Size([32, 128])

Sources