Triton for MTIA: Bridging the Programming Model Gaps for Custom AI Accelerators
Cet article présente le premier déploiement à l'échelle de la production du langage de programmation Triton sur l'accélérateur personnalisé MTIA-2i de Meta, démontrant qu'un nouveau backend de compilateur et des extensions de langage minimales permettent un développement de noyaux performant et efficace qui comble l'écart entre les frameworks d'apprentissage automatique et le matériel personnalisé.
Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète
Imaginez que vous essayiez de construire un système de trains massifs et à grande vitesse pour transporter des passagers (des données) à travers un continent. Pendant des années, les seuls trains disponibles étaient les modèles électriques standards et de haute technologie (les GPU) que tout le monde savait conduire. Ils avaient des portes automatiques, un GPS intégré et une conduite fluide. Mais récemment, un nouveau type de moteur de train a été inventé (les accélérateurs IA personnalisés). Ces nouveaux moteurs sont incroyablement puissants et conçus spécifiquement pour le transport de lourdes cargaisons, mais ils sont bizarres. Ils n'ont pas de portes automatiques ; vous devez les ouvrir et les fermer manuellement. Ils n'ont pas de GPS ; vous devez dessiner la carte vous-même. Et ils ne roulent pas sur les mêmes rails que les anciens trains.
Le problème est que les ingénieurs qui construisent les wagons de train (les développeurs de logiciels) sont habitués aux anciens trains électriques, faciles à conduire. S'ils veulent utiliser les nouveaux moteurs, super rapides, ils doivent apprendre une toute nouvelle façon de conduire, très difficile, ou bien ils doivent construire un nouveau wagon de train de toutes pièces pour chaque nouvel itinéraire. Cela ralentit tout. La grande question est : pouvons-nous apprendre l'ancien style de conduite, facile d'utilisation, à fonctionner sur ces nouveaux moteurs bizarres sans perdre l'avantage de la vitesse ? Ce document explore exactement cela, en tentant de combler le fossé entre le monde familier des logiciels d'IA standards et le matériel étrange et personnalisé construit par des entreprises comme Meta pour faire fonctionner leurs modèles d'IA.
L'Histoire : Apprendre à un nouveau moteur à parler une ancienne langue
Meta, l'entreprise derrière Facebook et Instagram, construit ses propres puces informatiques spéciales appelées MTIA-2i pour faire fonctionner ses modèles d'IA plus rapidement et à moindre coût. Considérez le MTIA-2i comme un moteur de voiture de course conçu sur mesure. Il est incroyable dans ce qu'il fait, mais il est construit différemment des moteurs standards (les GPU) auxquels la plupart des programmeurs d'IA sont habitués.
Dans le monde de l'IA, il existe un langage de programmation populaire appelé Triton. Vous pouvez considérer Triton comme une « télécommande universelle » pour l'IA. Au lieu d'écrire du code complexe et désordonné pour dire à un GPU standard comment déplacer des données, les développeurs écrivent un code simple et propre en Triton, et un compilateur intelligent le traduit en langage machine que le GPU comprend. C'est comme parler anglais à un traductur qui sait comment parler à la machine.
Cependant, Triton a été conçu à l'origine uniquement pour les moteurs GPU standards. Lorsque Meta a essayé d'utiliser Triton sur son moteur de voiture de course personnalisé, le MTIA-2i, cela n'a pas fonctionné. La « télécommande » ne correspondait pas aux boutons du nouveau moteur. Le moteur MTIA-2i fonctionne de manière asynchrone (il reçoit des ordres et les exécute plus tard), gère sa propre mémoire d'une manière étrange via un « tampon circulaire », et nécessite des instructions très spécifiques pour que tout circule de manière fluide. Le code Triton standard ne pouvait tout simplement pas gérer ces particularités.
Ce que fait le document : Construire un nouvel adaptateur
L'équipe de Meta a décidé de régler ce problème. Ils ne se sont pas contentés de forcer l'ancienne télécommande à fonctionner ; ils ont construit un nouveau backend de compilateur spécifiquement pour le MTIA-2i. Ce nouveau système agit comme un adaptateur super intelligent. Il prend le code Triton simple et facile à lire que les développeurs écrivent et le traduit parfaitement en instructions de bas niveau complexes dont le moteur MTIA-2i a besoin.
Voici comment ils ont réussi, en utilisant quelques analogies amusantes :
- La ligne de buffet « FIFO » : Sur les GPU standards, l'ordinateur gère automatiquement la mémoire comme un buffet intelligent où les assiettes sont saisies et remplacées instantanément. Sur le MTIA-2i, c'est plutôt comme une ligne d'assemblage manuelle avec un tampon « Premier Entré, Premier Sorti » (FIFO). Le nouveau compilateur a appris à gérer cette ligne parfaitement, garantissant que les données arrivent exactement quand la machine en a besoin, sans que le programmeur ait à pousser manuellement chaque assiette.
- La cuisine à deux chefs : La puce MTIA-2i possède deux « cœurs » (comme deux chefs dans une cuisine) qui peuvent travailler en même temps. Le compilateur a trouvé comment répartir les tâches de cuisine entre eux afin qu'ils ne se gênent pas ou ne s'attendent pas les uns les autres. Il permet même aux programmeurs experts de dire précisément aux chefs qui doit hacher les oignons et qui doit remuer la soupe s'ils veulent faire les choses en grand.
- La route de livraison en « zigzag » : Imaginez que vous devez livrer des colis à 64 maisons. Si vous descendez simplement la rue en ligne droite (linéaire), les premières maisons sont servies rapidement, mais les dernières attendent éternellement. L'équipe a découvert qu'un itinéraire de livraison en « zigzag » (en faisant des va-et-vient) équilibrait bien mieux le travail. Ils ont ajouté une fonctionnalité à Triton pour que les programmeurs puissent facilement choisir ce chemin plus intelligent.
Les Résultats : Rapide, Flexible et Prêt pour le Monde Réel
L'équipe ne s'est pas contentée de construire cela en laboratoire ; elle l'a mis au travail dans le monde réel. Elle a déployé avec succès ces nouveaux noyaux (kernels) Triton en production sur 60 types différents de modèles d'IA.
Voici ce qu'ils ont découvert :
- Vitesse : Le code écrit dans le langage simple Triton s'est avéré aussi rapide que le code écrit par des experts en C++ de bas niveau complexe (le « langage natif » de la puce). En fait, pour la multiplication de matrices générales (GEMM) — une tâche mathématique lourde — le code Triton a atteint plus de 80 % de la vitesse maximale théorique de la puce.
- Adoption : En seulement un trimestre, le nombre de types de modèles utilisant Triton a triplé.
- Impact : Triton gère désormais 50 % des couches et 47 % du temps d'exécution non-GEMM de ces modèles. C'est presque la moitié du travail effectué par ce nouveau système facile à utiliser !
Pourquoi cela compte
Le document soutient que vous n'avez pas à choisir entre « facile à écrire » et « super rapide ». Avant cela, si vous vouliez utiliser une puce personnalisée comme le MTIA, vous deviez être un sorcier du codage de bas niveau et tout écrire à partir de zéro. Si vous vouliez utiliser un langage de haut niveau comme Triton, vous étiez coincé avec les GPU standards.
En comblant ce fossé, Meta a démontré que les langages de haut niveau peuvent être adaptés pour fonctionner sur presque n'importe quel matériel personnalisé. Cela signifie qu'à l'avenir, lorsque les entreprises construiront des puces d'IA encore plus étranges et spécialisées, les développeurs n'auront pas besoin d'apprendre un nouveau langage difficile pour chacune d'elles. Ils pourront continuer à utiliser les outils familiers et puissants qu'ils connaissent déjà, tandis que le compilateur gérera les détails complexes du nouveau matériel.
Le document note également qu'ils ont testé auparavant un autre langage de très bas niveau appelé KNYFE, mais qu'ils l'ont abandonné car il était trop difficile à apprendre pour la plupart des gens. Ils ont prouvé qu'il vaut mieux s'en tenir à un langage flexible et de haut niveau comme Triton, même avec quelques petites modifications personnalisées, car c'est la stratégie gagnante pour maintenir le développement de l'IA rapide et efficace.
En résumé, le document démontre qu'avec le bon « tour de magie » du compilateur, nous pouvons faire fonctionner la « télécommande universelle » sur presque n'importe quelle « télévision », même celles qui sont étranges et fabriquées sur mesure, sans sacrifier la qualité de l'image.
Noyé(e) sous les articles dans votre domaine ?
Recevez des digests quotidiens des articles les plus récents correspondant à vos mots-clés de recherche — avec des résumés techniques, dans votre langue.