Il y a eu beaucoup de discussions suite à la présentation du noyau MLA Jalapeño au HotChips et aux commentaires ultérieurs de SemiAnalysis. En tant qu'équipe matérielle d'OpenAI, nous avons à peine effleuré cette petite pépite : le fait que l'IA écrit nos noyaux, et que, lorsqu'elle le fait, nous n'avons pas vraiment besoin de comprendre ce que fait le noyau ligne par ligne. Nous avons manifestement omis : comment une telle chose est-elle possible ? Quelle est la bonne façon de penser cela, par rapport à une méthode plus traditionnelle de génération de code ? Le noyau optimisé est-il aussi fiable que le noyau non optimisé ?
Pour vous situer, j'ai travaillé sur des compilateurs pour accélérateurs pendant plus d'une décennie. J'ai créé XLA, qui est une excellente infrastructure de compilation avec une équipe et un effort intersociétés remarquables. Depuis plus de 2 ans chez OpenAI, j'essaie de reconceptualiser le fonctionnement des compilateurs à l'ère de l'IA. De nouvelles formulations de compilateurs s'appuieront sur les forces existantes, mais il est impossible de nier qu'il existe un nouvel outil puissant à exploiter dans la boîte à outils.
Ce sera un petit voyage, mais j'espère éclairer comment l'IA est utilisée pour l'automatisation de l'amélioration des programmes informatiques ; c'est-à-dire la compilation optimisante. Je pense qu'avec l'IA, nous pourrions vivre ce que nous considérons comme les « compilateurs 2.0 ». L'IA est fondamentalement moins contrainte dans ce qu'elle peut proposer, et ce qu'elle propose est le résultat de l'entraînement et du contexte du modèle, ce qui m'amène à la classer comme un « optimiseur stochastique » – cela peut poser des défis mais, comme nous le verrons, c'est aussi une source de grandes forces…
Une grande partie de la recherche académique et des applications industrielles s'oriente déjà dans cette direction, et découvre rapidement le potentiel de l'implication de l'IA dans le domaine des compilateurs optimisants, mais nous sommes à un point où cela mérite une explication générale.
Contexte
Les compilateurs prennent des programmes et produisent des versions traduites ou améliorées de ces programmes.
Les programmes, à la fois en entrée et en sortie, ont une sémantique qui nous dit ce que les programmes signifient, ce qu'ils pourraient faire, et comment raisonner sur ces choses qu'ils pourraient faire.
Ceux d'entre nous qui travaillent sur les compilateurs les considèrent comme des fonctions pures – ils prennent une structure de données et produisent une structure de données qui doit avoir une sémantique correspondante.
Parfois, nos compilateurs se concentrent sur « l'abaissement » ou la « traduction ». Par exemple, ils peuvent prendre du C et produire de l'assembleur x86-64, ce que nous considérons souvent comme « de plus bas niveau ». Mais souvent, ils font plus que de la simple traduction dans le cadre de ce processus…
Nos compilateurs, en pratique, se concentrent sur l'« optimisation ». Ils peuvent prendre une structure de données qui représente le programme – dans notre jargon une « Représentation Intermédiaire » (RI) – et ils essaient de produire une meilleure version de ce programme. Parfois, « meilleur » signifie qu'il nécessite moins de cycles pour s'exécuter, parfois qu'il aura moins de code inutile, parfois qu'il se spécialise pour des choses que nous pouvons prouver « doivent être vraies » à propos du programme (évaluation partielle).
Maintenant, considérons brièvement que les LLM ont été créés à l'origine pour traduire un texte humain d'une langue à une autre. La traduction est clairement dans leur domaine de compétence. Et nous pouvons voir, à travers notre utilisation quotidienne des LLM, qu'ils peuvent également écrire de nouvelles solutions et améliorer des solutions existantes. Beaucoup d'entre nous, codeurs, avons aussi fait l'expérience de demander à un LLM « optimise cet extrait de code » et ils y parviennent remarquablement. (Cependant, nous devons savoir qu'ils ont optimisé le code correctement, ce que nous aborderons !) Cela vise simplement à souligner que les LLM ont les capacités que nous recherchons dans un compilateur optimisant.
Optimisation et Optimalité
Les compilateurs optimisants essaient, sans surprise, d'augmenter l'optimalité du programme sur lequel ils travaillent, selon un objectif donné (généralement le temps d'exécution). C'est tellement difficile à faire dans le cas général, pour un programme arbitraire, qu'il existe un théorème appelé le théorème du plein emploi pour les ingénieurs compilateurs. (Je ne l'ai découvert qu'après avoir choisi d'être ingénieur compilateur, mais cela m'a quand même réconforté !)
Les « superoptimiseurs » sont un sous-domaine fascinant des compilateurs optimisants. Imaginez qu'il existe un programme donné, et que nous puissions dire ce qu'il fait via sa sémantique. Quel est le programme le plus optimal qui a cette même sémantique ? C'est ce que les superoptimiseurs tentent de résoudre, et c'est effectivement un problème de recherche…
Imaginez que j'essaie de trouver le programme le plus court qui a cette même sémantique, et que j'aie un moyen de demander si un programme candidat a la même sémantique. Je pourrais, hypothétiquement, énumérer chaque programme dans un ordre objectif, et choisir le plus petit qui a la même sémantique.
Cependant, énumérer chaque programme dans un ordre objectif semble assez insoluble. L'un de mes articles académiques préférés, publié en 2013 et intitulé « STOKE » (Superoptimisation Stochastique), demandait : « et si nous modifiions simplement les programmes aléatoirement encore et encore, finirions-nous par observer le meilleur programme ? » Ils ont proposé que via une marche aléatoire (et avec notre ami OG de l'apprentissage automatique, Monte Carlo par chaîne de Markov / Metropolis-Hastings), vous finiriez par voir ce programme optimal.
La modification Monte Carlo est généralement stupide (vous choisissez une modification au hasard), mais rapide. Les LLM sont très intelligents (beaucoup de jetons de raisonnement), mais comparativement lents.
Et si, au lieu de la modification Monte Carlo stupide/rapide, nous laissions les LLM déterminer les directions dans lesquelles prendre les programmes ? Nous aurions un optimiseur stochastique très intelligent, promenant notre programme dans l'espace des programmes optimisés.
Intuitions pour l'Optimisation
Prenons du recul. Pensez à la personne que vous connaissez qui incarne le mieux « optimise à fond des extraits de code ». Appelons-la « l'optimiseuse Ollie ». Ollie a probablement une intuition viscérale du type de modifications de code qui pourraient porter leurs fruits. Ollie essaie probablement certaines choses pour voir si elles fonctionnent, et si elles ne fonctionnent pas, elle les annule et essaie autre chose. Mais elle a une intuition de ce qui est possible, et de comment elle pourrait battre le compilateur.
Ces intuitions qu'a Ollie dépassent souvent ce que font les compilateurs. Bien que les compilateurs optimisants modernes aient des résultats assez impressionnants, ils sont basés sur des règles et des heuristiques assez simples. Dans le jargon technique, ils sont basés sur l'idée d'une transformation locale de flux de données exécutée jusqu'à un point fixe. Nous ordonnons également les phases de considération ; c'est-à-dire que nous construisons des pipelines de compilation pour considérer A puis B, mais pas le problème composite AB. Les ordonnanceurs et les allocateurs de registres en sont un exemple notoire, de nombreux doctorats ont été tentés sur l'ordonnanceur-allocateur de registres composite (pour bénéficier de la suppression de l'ordre des phases), mais ils ont été difficiles à faire fonctionner en pratique.
C'est pourquoi l'expertise d'Ollie est précieuse. Souvent, Ollie sait comment équilibrer plusieurs problèmes NP-complets avec des heuristiques adaptées à la situation. Il y a donc une conscience et une sensibilité contextuelles plus personnalisées. Ollie est également capable d'employer des techniques que les compilateurs optimisants peuvent ne pas appliquer de manière rentable, surtout en combinaison, comme l'extraction de code, la création d'ABI personnalisés ou de transformations pour permettre la vectorisation, ou la multitude d'autres choses qui nous font grommeler « j'aimerais que le compilateur ait un moyen de faire ça… »
Considérez maintenant que l'IA, grâce à ses capacités de raisonnement, peut agir comme une mini Ollie. Elle n'a peut-être pas l'intuition correspondante en termes de ce qui portera ses fruits, mais elle a une idée de ce qui peut être rentable, et elle peut prendre de nombreuses, nombreuses tentatives.
Avec cette approche, contrairement à l'article STOKE, nous ne pouvons pas garantir qu'avec un temps infini nous verrons le programme optimal, mais parce que l'IA a des capacités de raisonnement « plus humaines », elle peut réellement obtenir une traction significative de type humain par unité de temps.
Faire le Lien : le Noyau MLA
Permettez-moi de commencer par dire : je ne sais pas quel code bas niveau l'IA a produit pour notre noyau MLA Jalapeño, mais je sais comment taper le numpy pour MLA.
Dans le compilateur XLA sur lequel j'ai travaillé auparavant, nous fusionnions ces opérations numpy en blocs, puis nous utilisions un métaprogramme appelé « émetteur » pour les abaisser en boucles, instructions et primitives de bas niveau.
Lorsque le compilateur XLA / programme émetteur faisait cela, je n'avais pas besoin de me soucier de l'assembleur qui sortait. Pour notre optimiseur stochastique, l'IA prend conceptuellement la place du métaprogramme émetteur – elle abaisse et optimise à la fois, et nous pouvons lui demander d'optimiser de plus en plus vers la ligne de toit.
J'espère que cela clarifie où se situe l'IA et comment elle est analogue à un composant d'un système de compilation optimisante existant. Il est également utile de penser : la couche que nous considérons comme « code assembleur » est maintenant en train de monter. Lorsque vous tapez du C++ normal et le compilez à -O3 (le niveau d'optimisation typique le plus élevé), vous ne vous attendez pas à comprendre l'assembleur qui en sort, même si vous compreniez le C++ que vous aviez tapé. Nous faisons la chose analogue ici, mais avec une spécification d'entrée de plus haut niveau et plus mathématique.
Maintenant, une question clé est comment nous vérifions que le programme que nous obtenons de l'IA est effectivement équivalent à la description de haut niveau / numpy. Ce mécanisme de vérification établit la fiabilité du processus d'optimisation stochastique par l'IA. Je m'attends à ce qu'un futur article de blog entre plus dans les détails à ce sujet, mais pour l'instant, qu'il suffise de dire que tester l'équivalence sémantique est possible et que nous le faisons. Les programmes pour accélérateurs sont particulièrement adaptés à des contrats forts et complets que nous pouvons vérifier « sont exactement ce que fait le programme optimisé par l'IA », car ils sont assez mathématiques et orientés flux de données dans leur contexte large.
Notez que de nombreuses techniques pertinentes dans ce domaine ont été pionnières par des efforts dans le sous-domaine de la synthèse de programmes. Alors que les compilateurs optimisants disent : « voici un programme avec une sémantique, améliorez-le mais avec une sémantique équivalente ! », la synthèse de programmes dit : « je crois qu'il existe un programme avec cette sémantique, veuillez essayer de trouver le meilleur possible ». La synthèse de programmes est un problème plus difficile que la compilation optimisante, mais elle est aussi fondamentalement moins contrainte. C'est effectivement ce que font les humains comme Ollie lorsqu'ils font mieux que le compilateur optimisant, et c'est quelque chose que l'IA peut maintenant nous aider à automatiser. L'IA peut puiser « l'inspiration » du programme original, mais elle n'a pas besoin de se contenter d'effectuer des transformations locales mineures sur celui-ci. Les compilateurs optimisants classiques ne verront pas « oh, vous avez écrit un tri à bulles » et, en comprenant le contrat, le remplacer par un tri rapide, mais Ollie et l'IA sont capables de le faire. C'est ce qui nous place davantage dans le régime de la synthèse de programmes avec optimisation stochastique que dans le régime classique des compilateurs optimisants.
Tout cela converge vers le fait que vous pouvez commencer avec quelque chose qui « n'est pas très loin du numpy », attendre 48 heures, et obtenir un noyau optimisé avec la même sémantique, comme nous l'avons montré dans notre présentation HotChips :
Comme le note également la diapositive, sur notre machine, nous sommes souvent capables d'observer l'IA dépasser les performances de nos experts humains, même sur les noyaux que nous pensions assez bien réglés. Souvent, un pourcentage réalisable décent reste encore à explorer en raison des nombreuses variétés de combinaisons / permutations qui peuvent être nécessaires. Celles-ci sont souvent d'une complexité fastidieuse pour un ingénieur en performance humain.
Récapitulatif et Conclusion
Un compilateur, en fin de compte, n'est qu'une fonction. Nous donnons notre programme à cette fonction, et nous obtenons une meilleure version de notre programme. Nous attendons du programme que nous obtenons et du programme que nous mettons qu'ils aient la même sémantique.
Les compilateurs optimisants traditionnels améliorent les programmes via des règles de flux de données et des heuristiques. Ceux-ci sont entièrement compréhensibles dans leur provenance, mais peuvent aussi être plus limités dans les transformations qu'ils peuvent effectuer.
En revanche, l'IA, en tant qu'optimiseur stochastique, doit simplement « réfléchir intensément » et produire quelque chose. Ses transformations ne sont pas fondamentalement limitées, ce qui la rend plus analogue à notre expert humain optimiseur. Nous avons besoin de moyens pour vérifier que les programmes qu'elle produit sont fiables et implémentent la même sémantique que celle que nous avons mise, et nous les avons en place. Et ce type d'optimisation par l'IA est particulièrement adapté aux opérations mathématiques qui ont des contrats très forts. Les contrats évitent le besoin de comprendre ce que fait le noyau ligne par ligne.
C'est ainsi que nous avons obtenu le noyau MLA généré par l'IA !
