Ajout en cours de la prise en charge de CUDA dans Apple MLX

• MLX est un framework de tableaux destiné aux chercheurs en machine learning sur Apple Silicon, offrant une ergonomie proche de NumPy et de PyTorch
• Des travaux sont en cours pour ajouter un backend CUDA à MLX
• L’objectif principal est la prise en charge de la mémoire unifiée (unified memory) et une exploitation étendue du matériel NVIDIA
• Pour l’instant, seuls les exemples des tutoriels fonctionnent, et la compilation ainsi que les tests ont été vérifiés dans un environnement Ubuntu 22.04 + CUDA 11.6
• Après une phase initiale de résolution de problèmes de performances et de goulots d’étranglement, les travaux d’optimisation et de refactorisation se poursuivent de manière répétée
• Le backend CUDA de MLX est développé avec le soutien d’Apple, avec pour objectif d’offrir une expérience cohérente entre le développement sur Mac et le déploiement dans de grands environnements de calcul grâce à la mémoire unifiée et à la prise en charge du matériel NVIDIA

---

Présentation du framework MLX

• MLX est un framework de tableaux pour le machine learning développé par l’équipe de recherche en apprentissage automatique d’Apple
• Il a été conçu pour permettre aux ingénieurs et aux chercheurs d’expérimenter efficacement et de valider rapidement leurs idées
• Sa spécificité est la prise en charge d’Apple Silicon, mais l’ajout d’un backend CUDA vise désormais à étendre ses possibilités d’exécution aux environnements équipés de GPU NVIDIA

Caractéristiques principales

• API familière : l’API Python est proche de NumPy, tandis que les packages de plus haut niveau mlx.nn et mlx.optimizers proposent une interface similaire à PyTorch
  • Des API C++, C et Swift sont également fournies, avec une ergonomie cohérente avec l’API Python
• Transformations de fonctions composables : inclut des fonctions comme la différenciation automatique, la vectorisation automatique et l’optimisation du graphe de calcul
• Évaluation paresseuse : plusieurs opérations peuvent être regroupées et calculées uniquement quand le résultat est réellement nécessaire, pour une utilisation plus efficace des ressources
• Création dynamique de graphes : même si la shape des données d’entrée change, l’exécution reste immédiate sans compilation lente, ce qui facilite le développement et le débogage
• Modèle multi-appareils et de mémoire unifiée :
  • Auparavant, seuls le CPU et le GPU Apple étaient pris en charge, mais le backend CUDA vise maintenant la prise en charge des environnements à GPU NVIDIA
  • Grâce à une architecture de mémoire unifiée, il devient possible de manipuler le même objet sur tous les appareils pris en charge sans copie mémoire
• Conception concise et extensible : les chercheurs peuvent étendre et améliorer le framework facilement

Différences avec les autres frameworks et inspirations

• Son design s’inspire de NumPy, PyTorch, Jax, ArrayFire, entre autres
• Il met particulièrement en avant le modèle de mémoire unifiée, une interface concise et un environnement d’exécution immédiate/de débogage

Exemples représentatifs et domaines d’utilisation

• Le dépôt MLX Examples contient divers exemples pratiques
  • Entraînement de modèles de langage Transformer
  • Génération de texte à grande échelle basée sur LLaMA et fine-tuning avec LoRA
  • Génération d’images avec Stable Diffusion
  • Implémentations de modèles récents, dont la reconnaissance vocale basée sur OpenAI Whisper

Signification de MLX et prise en charge de CUDA

• MLX, qui était un framework centré sur Apple Silicon, étend désormais son potentiel d’usage général dans les environnements à GPU NVIDIA avec l’introduction d’un backend CUDA
• Grâce à la prise en charge récente de CUDA, il pourrait s’imposer comme un framework adapté à la fois à la recherche de pointe et aux usages industriels, sur le matériel Apple comme NVIDIA

[gosarinamu]

Il suffirait juste de permettre l’utilisation des GPU Nvidia sur macOS… haha.

[yangeok]

C’est une bonne nouvelle, haha. J’espère qu’ils ajouteront vite la prise en charge de CUDA pour qu’on puisse aussi faire de l’entraînement à haute vitesse sur Mac~ !

[GN⁺]

Commentaires sur Hacker News

• Je veux vérifier si j’ai bien compris : 1, les programmes compilés avec MLX peuvent fonctionner sur des puces compatibles CUDA. Mais 2, il est impossible de faire fonctionner des programmes CUDA sur Apple Silicon. Je me demande si c’est bien parce que le point 2 violerait le droit d’auteur, en particulier le fameux fossé défensif de Nvidia
  • Le point 1 est correct. Cela permet aux développeurs de travailler sur des appareils Apple relativement moins puissants, UMA comprise, puis de déployer le code final sur des systèmes Nvidia relativement plus performants. C’est utile pour plusieurs raisons
  • Le point 2 n’est pas une violation du droit d’auteur. Il est possible de réimplémenter une API
  • Concernant le point 2, cela ne me semble vraiment pas exact. L’AMD HIP fait aussi essentiellement la même chose, après avoir abandonné OpenCL vers 2017–2018
  • Moi, je veux un point 3. J’aimerais pouvoir connecter un GPU NVIDIA à de l’Apple Silicon et utiliser CUDA. J’imagine Apple Silicon, la mémoire unifiée, le GPU et CUDA ensemble avec PyTorch, JAX ou TensorFlow, même si je n’ai pas encore vraiment testé MLX
  • Si le point 2 n’existe pas, c’est surtout parce que c’est beaucoup plus difficile
• Si vous vous êtes dit : « Attends, la plateforme Apple n’a pas de support CUDA officiel ? », il faut aussi noter que cette série de patchs ajoute le support des « plateformes Linux avec CUDA 12 et SM 7.0 (Volta) ou plus »
  Guide d’installation
• Ce projet est mené par zcbenz, connu notamment pour avoir créé Electron et d’autres projets
  Présentation de zcbenz
• L’une des principales caractéristiques de MLX est l’utilisation d’une architecture à mémoire unifiée (UMA), et je me demande comment cela fonctionne
  Voir la puce correspondante dans le readme du dépôt : MLX GitHub
  Je pense que ce serait intéressant de l’étendre à tous les APU UMA, quel que soit le constructeur, mais je me demande si cela ne devrait pas fonctionner différemment avec un GPU discret
  En lisant les commentaires de la PR, il semble que CUDA prenne aussi directement en charge une API UMA et gère les copies de manière transparente quand c’est nécessaire
  • D’après mon expérience, l’absence de fonctionnalités de prefetch provoque de gros goulots d’étranglement mémoire à cause de l’attente des copies. Si tout le dataset tient en VRAM, le prefetch manuel fonctionne, mais dans mon application, l’entraînement ML, la perte de performance était trop importante, donc je suis passé directement à un chargement en streaming
• C’est un peu un autre sujet, mais il semble qu’une bonne partie des contributeurs à MLX ne soient pas officiellement employés par Apple. Voir par exemple le Twitter de prince_canuma
  Je me demande si Apple soutient cela en coulisses sans l’afficher officiellement. Il y avait aussi des rumeurs de transfert d’équipe
• Dans la description de la PR, il est indiqué qu’il s’agit d’un « ongoing effort » pour ajouter un backend CUDA
  Autrement dit, cela signifie que le code MLX peut être compilé et exécuté sur du matériel x86 + GeForce, et non l’inverse
• Comme le dit l’idée selon laquelle « écrire et tester le code sur Mac puis le déployer sur un supercalculateur offrirait une bonne expérience de développement », je me demande maintenant si l’on peut utiliser MLX sur Linux
  J’ai testé moi-même, et seule la version Python 3.12 est publiée sur PyPI
  Page PyPI de MLX-CUDA
• Je me demande si cela est lié à l’effet Strix Halo. J’avais du budget restant pour mon poste de travail, donc j’ai acheté un EVO-X2, et c’est comparable, voire meilleur, que des machines EPYC znver5 de milieu de gamme. C’est largement devant la plupart des instances EC2 ou GCE, même côté NVMe. Ce matériel coûte 1800
  On voit aussi apparaître récemment des reviews de youtubeurs affirmant que le DGX Spark est moins intéressant en rapport qualité-prix que Strix Halo. L’inconvénient, c’est que ROCm reste encore insuffisant pour l’instant, même si cela semble devoir s’améliorer bientôt. Si CUDA fonctionnait parfaitement sur les appareils Apple, ce serait une option à envisager sérieusement, même si Strix est bon marché et très performant
  • Strix Halo correspond en fait à l’AMD Ryzen AI Max+ 395. On le trouve dans le Framework Desktop et dans plusieurs mini-PC. La bande passante mémoire de cette puce est de 200 GB/s, ce qui est très bon pour une plateforme x86 classique, mais reste très loin d’un GPU Nvidia, par exemple une 5090 à 1792 GB/s, ou d’un Apple M3 Ultra à 800 GB/s. Le rapport performance/prix est excellent, mais pour des charges très gourmandes en mémoire comme les LLM, j’ai l’impression qu’on vient seulement d’atteindre un minimum acceptable de performance
  • La description de la PR montre clairement que la cible est l’entraînement sur des clusters cloud
  • Cela fait 10 ans qu’on nous dit que les calculs de réseaux neuronaux sur GPU AMD vont bientôt concurrencer Nvidia, et cela ne s’est toujours pas matérialisé
  • Je me demande comment cela se compare à un Mac mini M4
• Je suis très impressionné par les modèles MLX. On peut ouvrir des modèles locaux à toute la famille sans craindre un risque d’incendie comme avec un ordinateur Nvidia. J’aimerais qu’Apple Silicon devienne un concurrent sérieux du camp des puces Nvidia. Je me demande aussi si ce support CUDA pourrait relever d’une stratégie de type embrace, extend, extinguish (EEE)
• Apple prévoit à l’avenir de construire des datacenters basés sur les puces de la série M afin de les utiliser pour le développement d’applications, les tests et l’hébergement de services externes