🤖 machine learning

Step-TP: A Grounded, Step-Level Dataset with Chain-of-Thought Reasoning for LLM-Guided Tensor Program Optimization

Cet article présente Step-TP, un jeu de données post-entraînement novateur qui améliore l'optimisation des programmes tensoriels guidée par les LLM en fournissant une supervision ancrée, atomique et au niveau des étapes, dotée d'un raisonnement structuré de type chaîne de pensée et d'une représentation intermédiaire économe en tokens, afin de permettre une prise de décision fiable en plusieurs étapes.

Auteurs originaux : Mengfan Liu, Da Zheng, Junwei Su, Chuan Wu

Publié 2026-05-26

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Mengfan Liu, Da Zheng, Junwei Su, Chuan Wu

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

La Vue d'Ensemble : Enseigner à un Chef à Cuisiner Plus Vite

Imaginez que vous avez un chef brillant (le Grand Modèle de Langage ou LLM) qui est incroyable pour lire des recettes et comprendre les saveurs. Cependant, ce chef n'a jamais réellement cuisiné dans une cuisine industrielle à grande vitesse. Lorsque vous lui demandez d'optimiser une recette pour qu'elle soit cuisinée plus vite, il fait souvent des suppositions basées sur l'apparence finale du plat, plutôt que de comprendre les étapes spécifiques nécessaires pour y parvenir. Il pourrait dire : « Ajoutez simplement plus de chaleur ! » sans réaliser que cela brûlerait la nourriture.

Dans le monde de l'informatique, « cuisiner » signifie exécuter des programmes mathématiques complexes (appelés programmes tensoriels) sur des cartes graphiques puissantes (GPU). Rendre ces programmes plus rapides est incroyablement difficile car il existe des millions de façons minuscules de réorganiser le code, et un seul faux mouvement peut tout casser.

Ce papier introduit Step-TP, un nouveau « livre de recettes » conçu pour enseigner au chef exactement comment effectuer ces modifications, étape par étape, avec une explication claire de pourquoi chaque étape fonctionne.

Le Problème : Le Livre de Recettes « Boîte Noire »

Avant ce papier, les jeux de données utilisés pour entraîner ces chefs IA présentaient trois problèmes majeurs :

Ils Montraient Seulement le Résultat : La plupart des anciens livres de recettes montraient simplement la recette « Avant » et la recette « Après ». Ils ne montraient pas les étapes intermédiaires. C'était comme montrer un poulet cru et un poulet rôti, mais en cachant l'assaisonnement, la température du four et le temps de cuisson. L'IA mémorisait simplement l'apparence du plat final au lieu d'apprendre la technique culinaire.
Trop de Désordre : Les recettes étaient écrites dans un langage très désordonné et technique (comme du code informatique brut) rempli de détails inutiles. C'était comme essayer de lire une recette où chaque instruction était écrite dans une police différente, avec des notes de bas de page sur la marque du four. Cela rendait difficile pour l'IA de se concentrer sur la logique réelle de la cuisson.
Pas de « Pourquoi » : L'IA n'apprenait pas le raisonnement. Si le chef voyait un nouveau type de légume, il ne pouvait pas comprendre comment le cuisiner car il avait seulement mémorisé des plats spécifiques, et non les règles générales de la cuisine.

La Solution : Step-TP

Les auteurs ont créé un nouveau jeu de données appelé Step-TP qui résout ces problèmes. Voici comment ils l'ont fait, en utilisant trois ingrédients principaux :

1. Un Langage Plus Clair (LEIR)

Les auteurs ont inventé une nouvelle façon d'écrire les « recettes » appelée LEIR (Représentation Intermédiaire d'Équations de Boucle).

L'Analogie : Imaginez traduire une note manuscrite désordonnée, pleine de gribouillis et de taches de café, en une liste propre et dactylographiée avec des puces.
Ce que cela fait : LEIR élimine tout le « code de base » informatique (les choses ennuyeuses et répétitives) et ne laisse que la logique essentielle : les boucles (les étapes répétitives) et les équations (les mathématiques). Cela rend la recette beaucoup plus courte et plus facile à lire et à comprendre pour l'IA.

2. Entraînement Étape par Étape (Étapes Atomiques)

Au lieu de montrer à l'IA toute la transformation du début à la fin, Step-TP la décompose en étapes uniques et minuscules.

L'Analogie : Au lieu de dire « Transformez le poulet cru en poulet rôti », le jeu de données dit :
- Étape 1 : Assaisonnez le poulet.
- Étape 2 : Préchauffez le four.
- Étape 3 : Placez le poulet dans le four.
Pourquoi c'est important : Cela apprend à l'IA à prendre une petite décision sûre à la fois. Elle apprend que « l'Étape 1 » mène à un état spécifique, ce qui permet ensuite « l'Étape 2 ». Cela empêche l'IA de faire des bonds gigantesques et risqués qui pourraient casser le programme.

3. Raisonnement en Chaîne de Pensée (Le « Pourquoi »)

Chaque étape du jeu de données est accompagnée d'une explication de « Chaîne de Pensée » (CoT).

L'Analogie : C'est comme une émission de cuisine où le chef ne fait pas seulement l'action, mais explique : « Je retourne le poulet maintenant parce que le dessous est doré. »
Ce que cela fait : Le jeu de données indique explicitement à l'IA pourquoi un changement spécifique a été effectué (par exemple : « Nous réorganisons ces boucles pour mieux utiliser la mémoire »). Cela aide l'IA à apprendre la logique sous-jacente afin qu'elle puisse appliquer ces règles à de nouveaux problèmes jamais vus auparavant.

Les Résultats : Un Chef Maître

Le papier a testé ce nouveau jeu de données sur différentes tailles de modèles d'IA. Voici ce qu'ils ont découvert :

Efficacité : Parce que les « recettes » (LEIR) étaient si propres, l'IA pouvait les traiter beaucoup plus vite, en utilisant beaucoup moins de « tokens » (mots) qu'auparavant. C'était comme passer de la lecture d'un manuel de 100 pages à la lecture d'une fiche triche de 10 pages.
Précision : L'IA est devenue beaucoup meilleure pour effectuer des modifications qui fonctionnaient réellement. Lors des tests, les modèles pouvaient transformer avec succès des programmes et les maintenir en fonctionnement correct plus de 90 % du temps, même pour des tâches très complexes.
Longs Parcours : L'IA pouvait gérer de longues chaînes d'optimisations. Au lieu de faire seulement un petit changement, elle pouvait planifier une séquence de 10 ou 15 étapes pour rendre un programme des centaines de fois plus rapide. C'était comme le chef apprenant à planifier un menu complet de banquet plutôt que de simplement cuisiner un seul accompagnement.

Résumé

Step-TP est un jeu de données d'entraînement spécialisé qui enseigne aux modèles d'IA comment optimiser les programmes informatiques en :

Utilisant un langage propre et simplifié (LEIR) pour éliminer le désordre.
Décomposant les tâches complexes en petites étapes gérables.
Fournissant des explications pour chaque étape afin que l'IA comprenne la logique, et non seulement le motif.

Le résultat est une IA capable d'agir comme un ingénieur expert fiable, apportant des améliorations précises et multi-étapes aux performances des logiciels sans rien casser.

Résumé Technique : Step-TP

Énoncé du Problème

L'optimisation de l'efficacité d'exécution des programmes tensoriels sur les GPU constitue un défi critique pour l'apprentissage profond, particulièrement à mesure que les architectures de modèles gagnent en complexité. Bien que les grands modèles de langage (LLM) aient démontré de solides capacités de raisonnement, leur application à l'optimisation des programmes tensoriels a été entravée par le manque de données d'entraînement adaptées. Les jeux de données existants fournissent généralement uniquement des paires de programmes optimisés de bout en bout, utilisant des représentations verbeuses et inefficaces en termes de tokens (telles que le CUDA brut ou le TVM TIR). Ces jeux de données manquent de supervision vérifiable au niveau des étapes, contraignant les LLM à s'appuyer sur l'imitation des résultats ou sur un appariement de motifs superficiel plutôt que sur l'apprentissage du processus décisionnel sous-jacent. Par conséquent, les LLM peinent à prendre des décisions fiables à l'étape unique au sein des vastes espaces d'optimisation combinatoire requis pour un réglage efficace des programmes tensoriels.

Méthodologie

Pour remédier à ces limitations, les auteurs introduisent Step-TP, un jeu de données de post-entraînement conçu pour fournir une supervision ancrée, atomique et au niveau des étapes, accompagnée d'un raisonnement structuré de type Chaîne de Pensée (CoT). La méthodologie repose sur quatre principes directeurs fondamentaux :

1. Représentation Intermédiaire Équation-Boucle (LEIR)

Les auteurs proposent LEIR, une représentation intermédiaire efficace en tokens et vérifiable, adaptée à l'optimisation par LLM. Contrairement aux IR existants (CUDA, TVM TIR) encombrés par la boilerplate des compilateurs et les détails spécifiques au matériel, LEIR distille les programmes tensoriels en deux composantes sémantiques irréductibles :

Structure de Boucle : Définit l'espace d'itération à l'aide de symboles compacts pour les boucles sérielles ( $L$ ), parallèles ( $P$ ), vectorisées ( $V$ ), déroulées ( $U$ ) et liées aux threads/blocs ( $B$ ).
Structure d'Équation : Spécifie le calcul algébrique à chaque itération en utilisant des opérateurs élément par élément, des variables tensorielles (avec des types de données explicites et une hiérarchie mémoire) et des délimiteurs.
LEIR utilise une initialisation implicite pour les motifs de réduction courants et emploie une syntaxe basée sur LaTeX pour exploiter les connaissances préalables des LLM, réduisant ainsi considérablement la consommation de tokens par rapport au CUDA et au TIR.

2. Construction du Jeu de Données au Niveau des Étapes

Le jeu de données est construit via un pipeline multi-étapes :

Génération de Sources : Des programmes PyTorch sont générés à partir d'opérateurs uniques et d'architectures composites (par exemple, modules d'attention), couvrant des formes et des types de données divers.
Transformation Pilotée par Stratégie en Une Étape : Au lieu d'une mise en correspondance de bout en bout, le système décompose les trajectoires d'optimisation complexes en stratégies atomiques et composables. Pour chaque LEIR d'origine, les stratégies applicables sont identifiées sur la base de préconditions, appliquées individuellement, et le LEIR transformé résultant est généré.
Supervision CoT Structurée : Chaque transformation est accompagnée d'une trace de raisonnement expliquant la justification de la stratégie et documentant les changements structurels spécifiques (transitions d'état IR-à-IR).
Vérification et Filtrage : Les programmes transformés sont abaissés vers le TVM TIR et exécutés pour vérifier l'équivalence numérique. Un mécanisme de filtrage des stratégies est appliqué pour équilibrer la distribution du jeu de données, empêchant les LLM de surajuster aux stratégies simples en sous-pondérant les transformations « faciles » et en conservant les « difficiles » (par exemple, softmax en ligne, variantes FlashAttention).

3. Taxonomie des Stratégies

Le jeu de données couvre 43 stratégies d'optimisation distinctes, catégorisées en quatre niveaux :

Niveau Graphique : par exemple, réordonnancement d'expressions, fusion d'opérateurs.
Niveau Opérateur : par exemple, réordonnancement de boucles, tuilage de boucles, division de boucles.
Niveau Mémoire : par exemple, lecture/écriture en cache, définition de la portée de stockage.
Niveau Mathématique : par exemple, softmax en ligne, factorisation, annulation.

Contributions Clés

Conception de LEIR : Une IR haute densité et vérifiable permettant une représentation compacte des états intermédiaires des programmes et l'application transparente de stratégies d'optimisation atomiques, réduisant l'empreinte en tokens d'environ 60 % par rapport au TIR et de 83 % par rapport au CUDA.
Construction de Step-TP : Le premier jeu de données de post-entraînement pour l'optimisation tensorielle par LLM fournissant une supervision ancrée, atomique et au niveau des étapes avec un CoT structuré. Il transforme l'espace de recherche d'optimisation en décisions interprétables à l'étape unique.
Validation Empirique : Des expériences extensives démontrant que Step-TP permet aux modèles de générer des transformations exécutables et ancrées à travers diverses stratégies et de naviguer dans des trajectoires d'optimisation à long horizon.

Résultats

Les auteurs ont évalué des modèles (série Qwen3) entraînés sur Step-TP à travers des tâches d'optimisation à étape unique et multi-étapes :

Efficacité des Tokens : LEIR a systématiquement requis moins de tokens que le TIR et le CUDA, permettant aux processus d'optimisation complexes de s'adapter dans des fenêtres de contexte pratiques.
Transformations à Étape Unique : Les modèles entraînés sur Step-TP ont atteint une haute fidélité dans la génération de transformations exécutables et numériquement équivalentes. Par exemple, le modèle Qwen3-32B a atteint un taux de « Pass Égal » (sorties identiques au niveau des bits) de 92 % sur des cas de test tenus en réserve. Les modèles ont couvert avec succès une gamme diversifiée de catégories de stratégies proportionnellement à la distribution du jeu de données.
Optimisation Multi-Étapes : En utilisant des algorithmes de recherche (Glouton, Beam, BFS, DFS, MCTS, Recherche en Chaîne), les modèles ont obtenu des accélérations de temps d'exécution substantielles. La recherche Beam, par exemple, a atteint une accélération moyenne de 42,90× et un pic de 561,82×. Les modèles ont composité avec succès des séquences de 2 à 10 types de stratégies distincts pour optimiser des trajectoires à long horizon.
Généralisation : Les connaissances d'optimisation capturées par Step-TP sont restées efficaces sur différentes architectures GPU (testées sur NVIDIA H20-3e et A100).
Études d'Ablation : Les résultats ont confirmé que la combinaison de LEIR, de CoT structuré et de filtrage des stratégies est critique. La suppression du CoT a considérablement réduit l'efficacité de la recherche multi-étapes (de 0,87 à 0,10), et la suppression du filtrage des stratégies a conduit à un biais vers des transformations simples.

Signification

L'article affirme que Step-TP comble une lacune critique dans le domaine en transformant les LLM de générateurs d'heuristiques grossières en agents d'optimisation fiables capables de raisonner sur les programmes tensoriels à grande échelle. En fournissant une supervision vérifiable au niveau des étapes, le jeu de données permet un apprentissage efficace en données et un raisonnement stable, permettant aux modèles de généraliser à travers les charges de travail et les backends matériels sans dépendre d'un apprentissage par renforcement coûteux ou d'un ajustement fin basé sur la recherche. Ce travail établit une fondation pour des cadres d'optimisation tensorielle plus flexibles et efficaces en données, capables de raisonner systématiquement sur le comportement d'exécution des GPU.