Learning Bug Context for PyTorch-to-JAX Translation with LLMs

Résumé Technique : Contexte de l'apprentissage des bugs pour la traduction PyTorch vers JAX avec les LLM

Énoncé du problème
Bien que les modèles de langage de grande taille (LLM) fassent preuve de solides performances lors de la traduction de code entre des langages de programmation généraux, leur fiabilité diminue considérablement lorsqu'ils traitent du code spécifique à un domaine, comme les frameworks de deep learning. Plus précisément, la traduction de code de PyTorch vers JAX présente des défis uniques dus aux différences fondamentales de sémantique d'exécution, notamment l'exécution de graphes dynamiques, la compilation juste-à-temps (JIT), les approches de différenciation automatique et les flux de vectorisation. Des travaux antérieurs, tels que UniTrans, tentent d'améliorer la traduction via une correction itérative utilisant des entrées de test ou des erreurs de compilation. Cependant, ces méthodes échouent souvent à rectifier efficacement les traductions, particulièrement lorsqu'elles reposent sur des LLM plus petits ou « faibles » (par exemple, ceux ayant moins de 10 milliards de paramètres), ce qui entraîne un code non fonctionnel ou des changements comportementaux subtils qui s'écartent de l'intention originale.

Méthodologie : Le framework T2J
Les auteurs introduisent T2J, un framework d'augmentation de prompt conçu pour améliorer la traduction PyTorch-vers-JAX basée sur les LLM en exploitant des ensembles de données curatés de bugs et de leurs corrections. Le framework fonctionne via plusieurs modules clés :

1. Construction de données (Dataset de bugs corrigés) :

   • Les auteurs partent de 20 extraits de code PyTorch issus du dataset TorchLeet (code de résolution de problèmes).
   • Ces extraits sont traduits en JAX à l'aide d'un LLM « faible » (gpt-4o-mini).
   • Des développateurs de logiciels professionnels sont engagés pour déboguer et réparer les implémentations JAX générées. Ce processus implique plusieurs cycles de vérification manuelle pour garantir que le code JAX est compilable, exécutable et produit des sorties équivalentes au code PyTorch original.
   • Cela aboutit à un dataset de 163 paires bug-solution curatées, documentant les erreurs de traduction prévalentes (ex: décalages d'API, écarts de gradient) et leurs corrections correspondantes.

2. Apprentissage en contexte via l'augmentation de prompt :

   • Au lieu de procéder au fine-tuning du modèle, T2J utilise l'apprentissage en contexte (in-context learning).
   • Le prompt de traduction est augmenté d'un guidage structuré dérivé du dataset de bugs corrigés. Spécifiquement, le prompt inclut une liste formatée en JSON des schémas d'erreurs courants et de leurs solutions comme contexte.
   • Le LLM « faible » (gpt-4o-mini) utilise ce contexte augmenté pour générer le code JAX corrigé sans nécessiter de réentraînement coûteux.

3. Framework d'évaluation :

   • Datasets : L'évaluation est menée sur deux datasets : un ensemble Intrinsèque (20 extraits TorchLeet avec une vérité terrain vérifiée par des humains) et un ensemble Extrinsèque (100 extraits PyTorch provenant de GitHub, où la vérité terrain est approximée par un LLM « fort », gpt-4o, suivi d'une vérification humaine).
   • Métriques : Le papier propose trois nouvelles métriques aux côtés de la traditionnelle CodeBLEU :
     • Score T2J CodeTrans : Une métrique de type « LLM-as-a-judge » évaluant l'Utilité et la Correction Fonctionnelle (avec et sans code de référence), inspirée de l'ICE-score.
     • Score T2J FixCost : Quantifie l'effort requis pour corriger le code en comptant le nombre d'étapes de correction humaine.
     • Score T2J Comparison : Un jugement binaire par un LLM comparant deux sorties de traduction pour déterminer laquelle est supérieure.

Contributions clés

1. Dataset T2J : La création du premier dataset de bugs corrigés spécifiquement pour les traductions PyTorch-vers-JAX, contenant des annotations détaillées des schémas d'erreurs et des corrections pour faciliter l'amélioration de la fiabilité de la traduction de code par LLM.
2. Framework T2J : Une nouvelle technique d'augmentation de prompt qui comble les lacunes spécifiques au domaine en intégrant un contexte structuré de bug-correction, offrant une approche scalable pour les migrations entre écosystèmes sans le coût computationnel du fine-tuning.
3. Validation empirique : Une évaluation rigoureuse démontrant que le framework proposé améliore significativement la qualité de la traduction pour les LLM faibles.

Résultats

• Évaluation intrinsèque : En utilisant T2J, le LLM faible (gpt-4o-mini) a obtenu une amélioration de 20 % du score proposé T2J CodeTrans (spécifiquement en correction fonctionnelle et utilité) par rapport au prompt standard de base. Le Score T2J Comparison a indiqué que 100 % des traductions générées par T2J étaient jugées supérieures à la version de base.
• Efficacité : Le framework a réduit l'effort humain requis pour le débogage. Le Score T2J FixCost est passé de 163 étapes (baseline) à 87 étapes, représentant une réduction d'environ 50 % de l'effort de correction.
• Performance : Le code corrigé généré par T2J s'est exécuté environ 2,5 fois plus vite que la sortie de base lors des tests de temps d'exécution.
• Évaluation extrinsèque : Sur le dataset GitHub, T2J a montré des améliorations de la correction fonctionnelle et de l'utilité (jusqu'à 1,2 point), bien que les scores CodeBLEU soient légèrement inférieurs à la baseline, soulignant les limites des métriques traditionnelles pour ce domaine spécifique.
• Corrélation : L'étude a constaté que les métriques T2J CodeTrans proposées présentaient la plus forte corrélation avec les coûts de correction humaine par rapport aux métriques traditionnelles comme CodeBLEU.

Signification et affirmations
Le papier affirme que T2J fournit une solution rentable et scalable pour améliorer la traduction de code entre frameworks. En curatant un dataset de « contexte d'apprentissage des bugs » et en utilisant l'apprentissage en contexte, les auteurs démontent que même des LLM plus faibles et à bas coût peuvent atteindre des traductions de haute qualité qui approchent les performances de modèles plus puissants, à condition d'être guidés par des exemples structurés d'erreurs communes et de leurs résolutions. Les auteurs positionnent ce travail comme une étape vers la modernisation des systèmes de deep learning hérités et la facilitation de la migration entre frameworks, tout en reconnaissant des limites telles que l'accent actuel mis sur les datasets de résolution de problèmes et l'absence d'application aux LLM open-source en raison de contraintes budgétaires.