HyperAI
il y a 12 jours

CUDA-L1 : Amélioration de l'optimisation CUDA grâce à l'apprentissage par renforcement contrastif

Xiaoya Li, Xiaofei Sun, Albert Wang, Jiwei Li, Chris Shum
CUDA-L1 : Amélioration de l'optimisation CUDA grâce à l'apprentissage par renforcement contrastif
Résumé

La croissance exponentielle de la demande en ressources de calcul GPU a créé un besoin urgent de stratégies automatisées d’optimisation CUDA. Bien que les avancées récentes dans les modèles de langage à très grande échelle (LLM) montrent un potentiel prometteur pour la génération de code, les modèles actuels de pointe (SOTA) obtiennent des taux de réussite faibles lorsqu’il s’agit d’améliorer la vitesse CUDA. Dans cet article, nous introduisons CUDA-L1, un cadre d’apprentissage par renforcement automatisé pour l’optimisation CUDA, basé sur un nouvel algorithme de renforcement contrastif. CUDA-L1 obtient des améliorations significatives de performance sur la tâche d’optimisation CUDA : entraîné sur une NVIDIA A100, il permet une accélération moyenne de x3,12, avec une médiane de x1,42, sur l’ensemble des 250 noyaux CUDA du benchmark KernelBench, les accélérations maximales atteignant x120. En outre, le modèle démontre une portabilité remarquable sur différentes architectures GPU, obtenant des accélérations moyennes de x3,12 sur L40, x2,50 sur RTX 3090, x2,39 sur H100 et x2,37 sur H20, malgré une optimisation spécifique pour l’A100. Les capacités de CUDA-L1 montrent qu’il est possible de transformer, par simple signal de récompense basé sur l’accélération, un modèle de langage initialement peu performant en un optimiseur CUDA efficace, sans recourir à l’expertise humaine ni à des connaissances spécifiques au domaine. Ce paradigme ouvre la voie à une optimisation automatisée des opérations CUDA, et promet de contribuer de manière significative à l’amélioration de l’efficacité des GPU, tout en atténuant la pression croissante sur les ressources de calcul GPU. Nous identifions également des défis importants liés à l’entraînement de modèles de renforcement pour des tâches telles que le développement CUDA, où le renforcement apprend souvent à exploiter des failles dans les fonctions de récompense plutôt que de résoudre correctement les problèmes d’optimisation visés. En identifiant ces modes d’échec et en analysant leurs causes profondes, nous proposons des méthodes concrètes pour concevoir des procédures d’entraînement plus robustes, capables de prévenir les comportements de « hacking de récompense » (reward hacking).