Vue d'ensemble

PyTorch 2.3, publié le 14 juin 2024, améliore torch.export pour l'exportation de modèles et intègre les noyaux Triton personnalisés dans torch.compile.

Fonctionnalités principales

torch.export amélioré

torch.export produit un graphe IR propre pour le déploiement, sans dépendance au runtime Python.

python 
import torch

class MonModele(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

modele = MonModele()
exemple = (torch.randn(1, 10),)
exported = torch.export.export(modele, exemple)
print(exported.module()(torch.randn(1, 10)))

Noyaux Triton

Les noyaux Triton personnalisés sont désormais intégrables dans le pipeline torch.compile, permettant d'optimiser les opérations GPU spécifiques.

python 
import torch

# torch.compile intègre les kernels Triton
@torch.compile
def ma_fonction(x, y):
    return torch.matmul(x, y) + torch.relu(x.sum(dim=-1, keepdim=True))

a = torch.randn(32, 64, device='cuda')
b = torch.randn(64, 128, device='cuda')
result = ma_fonction(a, b)
print(result.shape)

Sources