Résumé technique : GPU Forecasters : Les modèles de langage comme substituts sélectifs pour l'optimisation du temps d'exécution des noyaux

Énoncé du problème

Les noyaux (kernels) GPU sont les moteurs fondamentaux du deep learning moderne, où des gains d'efficacité mineurs se traduisent par des réductions significatives des coûts d'entraînement et d'inférence. L'optimisation de ces noyaux implique généralement un cycle génération-évaluation dans lequel des modèles de langage étendus (LLM) proposent des noyaux candidats, qui sont ensuite compilés et exécutés sur du matériel physique pour mesurer la performance. Cependant, ce processus d'évaluation constitue un goulot d'étranglement sévère. Chaque évaluation nécessite une compilation (prenant 60 à 120 secondes pour CUDA) et une exécution sur l'appareil, ce qui rend le débit de l'évaluation des ordres de grandeur plus lent que la génération. À mesure que les budgets de recherche augmentent, le coût des mesures physiques répétées limite la capacité à explorer le vaste espace des implémentations potentielles de noyaux.

Alors que les travaux antérieurs se sont concentrés sur l'utilisation des LLM pour générer de meilleurs noyaux, cet article traite du problème complémentaire de la réduction du coût d'évaluation. Les auteurs proposent d'utiliser les LLM non seulement comme générateurs, mais comme des substituts sélectifs (selective surrogates) pour le GPU physique. L'objectif est de prévoir la performance relative d'un noyau candidat par rapport à une implémentation de référence et d'identifier les cas où le modèle est incertain, reportant ainsi uniquement les cas incertains vers le GPU physique coûteux.

Méthodologie

Formulation de la tâche : Prévision sélective

Les auteurs formulent le problème comme la prévision de l'accélération relative ($S_i = T_{ref} / T_i$) d'un noyau candidat par rapport à un noyau de référence, étant donné les spécifications matérielles. Prédire le temps d'exécution absolu en secondes est jugé intraitable en raison de la complexité de la modélisation du comportement du compilateur et des détails micro-architecturaux. Au lieu de cela, la cible est discrétisée en huit classes ordinales d'accélération (allant de "Ralentissement sévère" à "Accélération extrême").

Le substitut (un LLM) prend en entrée le noyau de référence, le noyau candidat et les spécifications matérielles. Il génère une chaîne de raisonnement et produit une distribution de probabilité sur les huit classes. La confiance déclarée du modèle est définie comme la probabilité maximale assignée à une classe unique.

Stratégie de prédiction sélective

Le système fonctionne sous un cadre de prédiction sélective. Le substitut est utilisé pour classer les candidats. Si la confiance du substitut dépasse un certain seuil, la prévision est acceptée et le candidat n'est pas mesuré sur le GPU. Si la confiance est faible (indiquant une incertitude), le candidat est reporté vers le GPU physique pour une mesure de vérité terrain. Cela permet au processus de recherche de considérer beaucoup plus de candidats dans un budget de mesure GPU fixe.

Apprentissage par renforcement (RL) pour le calibrage

Pour améliorer l'exactitude du substitut et, surtout, son calibrage (l'alignement entre la confiance déclarée et l'exactitude réelle), les auteurs appliquent l'apprentissage par renforcement à un modèle aux poids ouverts (GPT-OSS-20B). Ils définissent des récompenses basées sur :

1. La correction : Si la classe prédite correspond à la classe mesurée.
2. Le calibrage catégoriel (Score de Brier) : Pénaliser la distribution en fonction de la distance par rapport à la cible "one-hot" réelle.
3. Le calibrage sensible à la distance (CRPS) : Pénaliser la distribution cumulative, en tenant compte de la nature ordinale des classes d'accélération (par exemple, prédire une accélération de 2× quand la vérité est de 1,5× est moins grave que de prédire 4×).

Les données d'entraînement sont construites comme un sous-produit de la recherche de noyaux elle-même, en utilisant les candidats réussis des exécutions de recherche précédentes comme exemples d'entraînement.

Principales contributions

1. Les LLM comme substituts GPU : L'article démontre que les LLM peuvent atteindre une précision non triviale pour prévoir comment un GPU physique répond à des noyaux candidats, récupérant une fraction substantielle de l'accélération disponible en classant les candidats sur la base de prévisions plutôt que de mesures.
2. Cadre de prédiction sélective : Les auteurs présentent la prévision GPU comme un problème de prédiction sélective, montrant que les substituts peuvent rapporter une incertitude calibrée. Cela permet une stratégie de recherche qui fait confiance aux prévisions de haute confiance et réserve les mesures physiques pour les cas incertitudes.
3. Impact du RL sur le comportement du substitut : À travers des expériences, l'article caractérise comment l'entraînement par RL modifie les propriétés du substitut. Spécifiquement, l'ajout de récompenses de calibrage (Brier ou CRPS) améliore la fiabilité des estimations de confiance du modèle, tandis qu'un entraînement basé uniquement sur la correction peut dégrader la relation entre confiance et exactitude.
4. Amélioration de la recherche de bout en bout : L'étude montre que l'intégration du substitut dans une boucle de recherche de noyaux (spécifiquement une recherche basée sur PUCT) permet d'explorer significativement plus de candidats par mesure GPU. Sur la plupart des tâches, cela permet de trouver des noyaux plus rapides qu'une ligne de base qui mesure chaque proposition sous le même budget GPU.
5. Publication du jeu de données : Les auteurs publient une collection réutilisable de 12 388 noyaux GPU générés par LLM avec des temps d'exécution mesurés, couvrant 118 problèmes, plusieurs types de GPU (A100, L40S, H100) et diverses méthodes de génération. Ce jeu de données soutient la recherche plus large dans la prédiction de performance GPU et la recherche de noyaux.

Résultats

• Précision de la prévision : Les LLM prêts à l'emploi (ex. GPT-OSS-20B) récupèrent environ 83 % de la meilleure accélération disponible lorsque seulement le top 1 % des candidats est mesuré sur le GPU. Cette performance s'élève à 94 % lorsqu'on mesure le top 50 %.
• Calibrage et RL : L'entraînement par RL modifie considérablement le comportement du substitut.
  • L'entraînement axé uniquement sur la correction améliore l'exactitude de la classe exacte mais dégrade la corrélation entre la confiance et l'erreur (le modèle devient trop confiant sur les mauvaises prédictions).
  • L'entraînement de calibrage (Brier/CRPS) améliore la monotonie de l'erreur de prévision par rapport à la confiance. Notamment, l'entraînement calibré par Brier améliore la métrique "accélération récupérée" d'environ 2,1 points de pourcentage par rapport au modèle de base, suggérant que des distributions mieux calibrées mènent à de meilleurs signaux de classement pour l'algorithme de recherche.
• Efficacité de la recherche : Dans une recherche de noyaux de bout en bout, la recherche assistée par le substitut trouve des noyaux plus rapides que la ligne de base sur quatre des six tâches, atteignant souvent la meilleure accélération de la ligne de base avec nettement moins de mesures GPU. Sur les tâches où la ligne de base gagne, la recherche semble saturée précocement, indiquant que le substitut est plus bénéfique lorsque l'espace de recherche est vaste et inexploré.
• Moments de découverte : Le substitut peut identifier les "moments de découverte" (mutations où un noyau enfant est significativement plus rapide que son parent) mieux que le hasard, bien que la précision reste un défi, nécessitant une confirmation par GPU pour les décisions à enjeux élevés.

Signification et affirmations

L'article affirme que les LLM peuvent jouer un rôle plus large dans l'optimisation des noyaux au-delà de la simple génération de code. En agissant comme des modèles virtuels d'un GPU, ils peuvent servir de filtres à haut débit et à faible coût qui gèrent la rareté des ressources matérielles physiques.

Les auteurs soulignent que la valeur primaire de cette approche n'est pas seulement la vitesse, mais la capacité d'effectuer une optimisation basée sur un modèle sans interroger le système d'origine à chaque étape. Le substitut offre un débit plus élevé et un coût par requête plus faible qu'un GPU, en échange d'un bruit de prévision qui peut être géré via des estimations d'incertitude calibrées. Ce travail relie l'optimisation des noyaux au domaine plus large de la construction de simulateurs neuronaux pour des systèmes physiques complexes, suggérant une voie vers des optimiseurs qui traitent la mesure matérielle comme une ressource rare à utiliser de manière sélective.