Exploiter la puissance des GPU AMD pour l’entraînement de modèles de langage avancés
Cet article détaille les efforts de l’équipe de LightOn pour adapter l’architecture de modèle de langage Mamba afin de l’entraîner sur les GPU AMD Instinct MI250/MI250X. Malgré les défis techniques liés à l’optimisation des opérations clés du modèle Mamba sur l’architecture AMD, les résultats obtenus en termes de performances d’entraînement et de convergence du modèle sont très encourageants. Cela démontre la possibilité de rendre l’entraînement de modèles de langage de pointe plus indépendant de la plateforme matérielle, ouvrant la voie à une plus grande accessibilité et flexibilité pour les praticiens de l’IA.
Points clés
- Les GPU AMD Instinct MI250/MI250X offrent des performances de pointe pour les charges de travail d’IA, avec jusqu’à 362 TFLOPs en FP16/BF16
- L’équipe de LightOn a développé des noyaux personnalisés pour adapter l’architecture Mamba aux GPU AMD, en se concentrant sur les opérations clés comme la convolution causale 1D et le “selective scan”
- Les tests de performance montrent des résultats compétitifs par rapport aux GPU NVIDIA, avec un débit allant jusqu’à 44 000 tokens/seconde sur 8 GPU MI250
- La convergence du modèle Mamba entraîné sur AMD est équivalente à celle obtenue sur GPU NVIDIA, avec des résultats similaires sur les benchmarks d’évaluation
- Le partenariat avec Nscale, un fournisseur de cloud européen spécialisé dans les solutions HPC et IA durables, a permis d’accéder aux ressources AMD nécessaires à ces travaux
À retenir
Ces résultats démontrent que l’entraînement de modèles de langage avancés comme Mamba peut désormais être réalisé de manière performante sur une variété de plateformes matérielles, y compris les GPU AMD. Cela ouvre la voie à une plus grande accessibilité et flexibilité pour les praticiens de l’IA, qui pourront choisir leur matériel en fonction des coûts d’utilisation plutôt que de se limiter à une seule architecture. Bien que des optimisations supplémentaires soient encore possibles, cette prouesse technique est une étape importante vers l’indépendance des modèles de langage vis-à-vis de la plateforme d’entraînement.
Sources
