Il y a de cela quelques semaines, j’avais abordé une fonctionnalité très intéressante de PyTorch 2.0 : la possibilité de compiler le graphe de calcul de rétropropagation dynamique directement en CUDA. Cette avancée remarquable qui permet d’accélérer l’entraînement des réseaux de neurones jusqu’à deux fois plus rapidement!
Aujourd’hui, je souhaite vous présenter une nouvelle fonctionnalité introduite avec PyTorch 2.1 (uniquement disponible sur Linux au moment d’écrire ces lignes). Il s’agit de la possibilité de convertir du code Numpy en CUDA sans nécessiter de modifications! De plus, comme PyTorch utilise tous les cœurs disponibles par défaut, on a du «multi-processing» transparent d’inclus! Le tout se traduit par des accélérations incroyables. Dans un article récent, l’équipe de PyTorch rapporte des exécutions jusqu’à 57 fois plus rapides sur un ordinateur doté de 32 cœurs.
Il est également possible d’inspecter le code C++ généré. À la place du code CUDA illisible, PyTorch génère du code Triton, offrant ainsi une meilleure lisibilité.
L’article fournit une multitude d’exemples illustrant l’utilisation de cette nouvelle fonctionnalité. Il explique notamment comment intégrer nos fonctions Numpy dans un contexte CUDA. Cependant, il convient de noter que quelques ajustements et refactorisations pour faire fonctionner le code en CUDA pourraient prendre plus de temps que prévu. Voici une liste des détails à vérifier en double et en triple.
1. Habituellement, pour convertir un tenseur de PyTorch en Numpy, une série d’opérations est requise : tensor.detach().cpu().numpy(). Cependant, avec l’intégration du code Numpy en CUDA, seule la méthode tensor.numpy() est désormais nécessaire, puisque le tenseur restera en GPU.
2. Dans l’article, il est aussi mentionné de faire attention à l’utilisation de masques booléens sur-place lors de la manipulation de tenseurs (je l’avoue, je suis le premier à le faire). Ce type de manipulation devrait être évité dans un contexte où le code est compilé dans un graphe CUDA, car cela peut entraîner des ruptures dans celui-ci.
3. Bien entendu, pour optimiser les performances de notre code, il est préférable d’éviter d’écrire des boucles for directement en Python. Cela est d’autant plus crucial si l’on prévoit compiler notre code en CUDA en utilisant torch.compile(). Cependant, dans les cas où l’utilisation d’une boucle for est inévitable, la meilleure option semble être de compiler uniquement l’intérieur de la boucle.
Enfin, l’article évoque la possible introduction de fonctionnalités similaires pour les librairies Scipy et Scikit-Learn. Si cela se concrétise, il deviendra envisageable de créer des solutions en intelligence artificielle (IA), compilées en CUDA, sans quitter l’environnement Python. Cela marquerait une petite révolution dans le domaine de l’IA en temps réel et, je n’en doute pas, éviterait beaucoup de futurs maux de tête chez les chercheurs, moi y compris.