Barbara
UNE INTRODUCTION AUX TECHNOLOGIES DE COMPILATION EN APPRENTISSAGE PROFOND DANS PYTORCH POUR LA CAPTURE DE GRAPHES, LES REPRÉSENTATIONS INTERMÉDIAIRES, LA FUSION D’OPÉRATEURS ET LA GÉNÉRATION OPTIMISÉE DE CODE C++ ET GPU
La programmation informatique est magique. Nous écrivons du code dans des langages lisibles par des humains, et comme par magie, il se traduit en courants électriques à travers des transistors en silicium qui agissent comme des interrupteurs et leur permettent de mettre en œuvre une logique complexe – juste pour que nous puissions regarder des vidéos de chats sur Internet. Entre le langage de programmation et les processeurs matériels qui l’exécutent, se trouve une technologie importante – le compilateur. Le rôle d’un compilateur est de traduire et de simplifier notre code en langage lisible par l’humain en instructions compréhensibles par un processeur.
Dans l’apprentissage profond, les compilateurs jouent un rôle très important pour améliorer les performances de l’entraînement et de l’inférence, améliorer l’efficacité énergétique et cibler des matériels accélérateurs d’IA diversifié. Dans cet article de blog, je vais discuter des technologies de compilation d’apprentissage profond qui alimentent PyTorch 2.0. Je vais vous guider à travers les différentes phases du processus de compilation et discuter des différentes technologies sous-jacentes avec des exemples de code et des visualisations.
Un compilateur d’apprentissage profond traduit le code de haut niveau écrit dans des frameworks d’apprentissage profond en code optimisé de niveau inférieur spécifique au matériel pour accélérer l’entraînement et l’inférence. Il trouve des opportunités dans les modèles d’apprentissage profond pour optimiser les performances en réalisant une fusion de couche et d’opérateur, une meilleure planification de la mémoire et en générant des noyaux fusionnés optimisés spécifiques à la cible pour réduire les frais généraux d’appel de fonction.
Contrairement aux compilateurs logiciels traditionnels, les compilateurs d’apprentissage profond doivent travailler avec du code hautement parallélisable souvent accéléré sur des matériels accélérateurs d’IA spécialisés (GPUs, TPUs, AWS Trainium/Inferentia, Intel Habana Gaudi, etc.). Pour améliorer les performances, un compilateur d’apprentissage profond doit tirer parti des fonctionnalités spécifiques au matériel telles que la prise en charge de la précision mixte, les noyaux optimisés pour les performances et minimiser la communication entre l’hôte (CPU) et l’accélérateur d’IA.
Bien que les algorithmes d’apprentissage profond continuent d’avancer à un rythme rapide, les accélérateurs matériels d’IA évoluent également pour suivre les besoins de performance et d’efficacité des algorithmes d’apprentissage profond. Dans cet article de blog, je vais me concentrer sur le côté logiciel des choses, et plus particulièrement sur le sous-ensemble de logiciels plus proches du matériel – les compilateurs d’apprentissage profond.
PyTorch 2.0 comprend de nouvelles technologies de compilation pour améliorer les performances des modèles et l’efficacité d’exécution et cibler diverses arrière-plans matériels avec une API simple : torch.compile (). Bien que d’autres articles de blog et des articles aient discuté en détail des avantages de performance de PyTorch 2.0, je vais me concentrer sur ce qui se passe sous le capot lorsque vous invoquez le compilateur PyTorch 2.0. Si vous recherchez des avantages de performance quantifiés, vous pouvez trouver un tableau de bord de performance de différents modèles de huggingface, timm, et torchbench.
À un niveau élevé, les options par défaut pour le compilateur d’apprentissage profond de PyTorch 2.0 effectuent les tâches clés suivantes :
– Capture de graphes : Représentation graphique des calculs pour vos modèles et fonctions. Technologies PyTorch : TorchDynamo, Torch FX, FX IR
– Différenciation automatique : Suivi de graphique inverse à l’aide de la différenciation automatique et de la réduction aux opérateurs primitifs. Technologies PyTorch : AOTAutograd, Aten IR
– Optimisations : Optimisations de niveau de graphique avant et arrière et fusion d’opérateurs. Technologies PyTorch : TorchInductor (par défaut) ou autres compilateurs
– Génération de code : Génération de code C++/GPU spécifique au matériel. Technologies PyTorch : TorchInductor, OpenAI Triton (par défaut), autres compilateurs
À travers ces étapes, le compilateur transforme votre code et génère des représentations intermédiaires (RI) qui sont progressivement “abaissées”. Abaisser est un terme du lexique du compilateur qui fait référence à la mise en correspondance d’un large ensemble d’opérations (telles que supportées par l’API PyTorch) à un ensemble restreint d’opérations (telles que supportées par le matériel) par transformation automatique et réécriture effectuée par le compilateur. Le flux du compilateur PyTorch 2.0 :
Si vous êtes nouveau en ce qui concerne la terminologie du compilateur, ne soyez pas effrayé par tout cela pour l’instant. Je ne suis pas non plus un ingénieur de compilation. Continuez à lire et les choses deviendront claires, car je vais décomposer le processus à l’aide d’un exemple simple et de visualisations.
Pour simplifier, je vais définir une fonction très simple et la lancer à travers le processus de compilation PyTorch 2.0. Vous pouvez remplacer cette fonction par un modèle de réseau neuronal profond ou une sous-classe nn.Module, mais cet exemple devrait vous aider à comprendre ce qui se passe sous le capot beaucoup mieux qu’un modèle complexe de plusieurs millions de paramètres.
Le code PyTorch pour cette fonction est le suivant :
def f(x):
return torch.sin(x)**2 + torch.cos(x)**2
Maintenant, il est temps d’appeler torch.compile (). Tout d’abord, convainquons-nous que la compilation de cette fonction ne change pas sa sortie. Pour le même vecteur aléatoire 1×1000, l’erreur quadratique moyenne entre la sortie de notre fonction et un vecteur de 1 doit être zéro pour la fonction compilée et non compilée (sous une certaine tolérance d’erreur).
Tout ce que nous avons fait a été d’ajouter une seule ligne de code supplémentaire torch.compile() pour invoquer notre compilateur. Voyons maintenant ce qui se passe sous le capot à chaque étape.
La première étape pour le compilateur consiste à déterminer ce qu’il doit compiler. Et c’est là que TorchDynamo intervient. TorchDynamo intercepte l’exécution de votre code Python et le transforme en une représentation intermédiaire FX et le stocke dans une structure de données spéciale appelée FX Graph. À quoi cela ressemble-t-il, demandez-vous ? Heureusement que vous demandez. Ci-dessous, nous allons examiner le code que nous utilisons pour générer cela, mais voici la transformation et la sortie :
Il est important de noter que les graphiques Torch FX ne sont que des conteneurs pour l’RI et ne spécifient pas vraiment les opérateurs qu’ils doivent contenir. Dans la section suivante, nous verrons le conteneur FX graphique apparaître à nouveau avec un ensemble différent de RI. Si vous comparez le code de fonction et l’IR FX, il y a très peu de différence entre les deux. En fait, c’est le même code PyTorch que vous avez écrit, mais disposé dans un format que la structure de données FX graph attend. Ils fourniront tous deux le même résultat lorsqu’ils seront exécutés.
Si vous appelez torch.compile () sans arguments, il utilisera les paramètres par défaut qui exécutent toute la pile de compilation qui inclut le compilateur de matériel d’arrière-plan par défaut appelé TorchInductor. Mais nous irions un peu trop vite si nous discutions de TorchInductor maintenant, donc mettons ce sujet de côté et revenons-y quand nous serons prêts. Tout d’abord, nous devons discuter de la capture de graphes et nous pouvons le faire en interceptant les appels de torch.compile (). Voici comment nous allons procéder :
torch.compile () vous permet également de fournir votre propre compilateur, mais comme je ne suis pas un ingénieur de compilation et que je n’ai pas la moindre idée de la manière d’écrire un compilateur, je vais fournir une fausse fonction de compilateur pour capturer l’IR de FX graph que TorchDynamo génère.
Ci-dessous est notre fausse fonction de backend de compilateur appelée inspect_backend pour torch.compile () et dans cette fonction, je fais deux choses :
– Imprimer le code IR FX qui a été capturé par TorchDynamo
– Sauvegarder la visualisation du graphique FX
La sortie de cet extrait de code ci-dessus est le code IR FX et le diagramme de graphique montrant notre fonction sin^2(x) + cos^2(x).
Notez que notre fausse fonction de compilateur inspect_backend n’est invoquée que lorsque nous appelons la fonction compilée avec des données, c’est-
