La plateforme de machine learning MLX d’Apple, initialement conçue pour les puces Apple Silicon, fait un pas significatif avec l’ajout d’un backend CUDA. Cette évolution permet aux développeurs d’exécuter des modèles MLX directement sur les GPU NVIDIA, une avancée qui mérite qu’on s’y attarde.
Ce projet est piloté par le développeur @zcbenz sur GitHub, qui a commencé à expérimenter avec le support CUDA il y a quelques mois. Au fil du temps, il a décomposé le projet en plusieurs segments, les intégrant progressivement à la branche principale de MLX.
Bien que le backend soit encore en développement, plusieurs opérations fondamentales comme la multiplication de matrices, le softmax, la réduction, le tri et l’indexation sont déjà prises en charge et testées.
C’est quoi CUDA ?
Cuda, ou Compute Unified Device Architecture, est à NVIDIA ce que Metal est à Apple : une plateforme de calcul créée pour exploiter pleinement les capacités de ses propres GPU lors de tâches de calcul parallèle à haute performance.
Pour beaucoup, CUDA représente la norme pour exécuter des charges de travail de machine learning sur les GPU NVIDIA. Ce système est omniprésent dans l’écosystème ML, allant de la recherche académique aux déploiements commerciaux. Des frameworks comme PyTorch et TensorFlow, de plus en plus connus même en dehors des cercles spécialisés, reposent sur CUDA pour bénéficier de l’accélération GPU.
Pourquoi MLX d’Apple supporte-t-il désormais CUDA ?
MLX étant étroitement intégré à Metal sur les plateformes Apple, il n’était pas à l’origine conçu pour fonctionner en dehors de macOS. L’ajout d’un backend CUDA modifie cette donne.
Cela permet aux chercheurs et ingénieurs de prototyper localement sur un Mac, en utilisant Metal et les puces Apple Silicon, puis de déployer le même code sur de vastes clusters GPU NVIDIA, qui dominent encore les charges de travail de formation en machine learning.
Cependant, des limitations subsistent. Actuellement, tous les opérateurs MLX ne sont pas encore implémentés, et le support des GPU AMD est encore à venir. Néanmoins, l’opportunité d’exécuter du code MLX sur des GPU NVIDIA, sans avoir à le réécrire, ouvre la voie à des tests, expérimentations et cas d’utilisation en recherche beaucoup plus rapides.
Si vous souhaitez explorer par vous-même, les détails sont disponibles sur GitHub.