Premiers pas avec MLX pour les puces Apple

Bonjour, je m’appelle Awni. Aujourd’hui, je vais vous présenter MLX. MLX est un framework open source optimisé pour les puces Apple. C’est un outil flexible qui permet aussi bien de faire des calculs simples que de faire tourner des modèles d’apprentissage automatique très puissants sur les appareils Apple. Si vous voulez créer du texte, des images, du son ou même de la vidéo avec les derniers modèles d’IA, MLX est fait pour vous. Vous pouvez l’utiliser pour entraîner, ajuster ou personnaliser vos modèles d’apprentissage automatique sur votre Mac. Je vais d’abord vous expliquer ce qu’est MLX et où il est utile. Je vais aussi couvrir les bases de MLX en Python, comme l’installation et les opérations sur les tableaux. Ensuite, je vous présenterai les principales fonctionnalités qui distinguent MLX des autres frameworks. Je vais aussi présenter les outils de MLX qui permettent d’accélérer l’apprentissage automatique sur la puce Apple.

Dans la dernière partie, je vous présenterai MLX Swift et comment commencer à l’utiliser. Passons à une introduction de MLX. MLX a été conçu dès le départ pour offrir de hautes performances sur la puce Apple, que ce soit sur le CPU ou le GPU. MLX s’adapte à des usages variés, des calculs simples à l’apprentissage automatique à grande échelle. Il est conçu pour être simple et flexible, sans sacrifier la vitesse ni l’efficacité.

MLX offre une API de base proche de NumPy. Il peut souvent remplacer NumPy pour accélérer la plupart des calculs. MLX intègre tout pour l’apprentissage automatique, comme la différenciation automatique et des bibliothèques avancées. Ces API avancées sont similaires à PyTorch et JAX. Si vous avez déjà utilisé l’un de ces frameworks, MLX vous semblera familier et encore plus facile à utiliser. MLX exécute de l’apprentissage automatique avancé directement sur l’appareil. Par exemple, MLX est utilisé par LM Studio, une appli populaire qui génère du texte avec des modèles de langage directement sur Mac.

Le module MLX LM, basé sur MLX, permet de générer du texte et d’ajuster des modèles de langage comptant jusqu’à des centaines de milliards de paramètres. Pour en savoir plus, consultez la séance « Explore large language models on Apple silicon with MLX ». MLX dispose d’une API Python complète, utile pour le prototypage rapide. Il propose aussi une API en Swift, avec tous les outils nécessaires pour créer et optimiser des réseaux neuronaux. Il a aussi des API en C++ et en C. MLX permet d’exécuter les derniers modèles d’apprentissage automatique sur une puce Apple, y compris Mac, iPhone, iPad et Vision Pro. Tout le logiciel MLX est open source sous licence MIT permissive. Le cœur du logiciel est disponible sur GitHub, avec plusieurs exemples et packages créés avec les API Python et Swift. MLX a aussi une communauté active de créateurs de modèles sur Hugging Face. De nombreux modèles récents sont disponibles dans l’organisation MLX sur Hugging Face et de nouveaux y sont ajoutés chaque jour. Le plus simple pour débuter avec MLX en Python est de l’installer depuis PyPi. Il suffit d’exécuter une seule ligne dans le terminal pip3 install mlx. MLX est facile à utiliser. Pour lancer des calculs sur des tableaux, ouvrez un fichier Python et importez MLX. Puis, vous pouvez créer des tableaux et faire des opérations simples. Par exemple, nous allons ajouter ici deux tableaux entiers.

Vous pouvez aussi voir facilement la forme et le type d’un tableau. Comme indiqué, l’API Python MLX est similaire à NumPy. Les opérations ont généralement les mêmes noms, signatures et comportements. Si vous avez déjà utilisé NumPy ou un framework similaire, MLX vous semblera familier et facile à utiliser. Maintenant que vous connaissez MLX et ses bases, voyons ses fonctionnalités clés. Elles incluent mémoire unifiée, évaluation paresseuse, transformations de fonctions et packages avancés pour créer et optimiser des réseaux neuronaux. MLX est conçu pour exploiter au mieux les capacités de la puce Apple. Cela inclut un nouveau modèle adapté à la mémoire unifiée.

La plupart des systèmes d’apprentissage automatique ont un GPU séparé avec sa propre mémoire. La puce Apple, quant à elle, possède une architecture mémoire unifiée. Cela signifie que le CPU et le GPU partagent la même mémoire physique.

Pour utiliser la mémoire unifiée, MLX fonctionne différemment des frameworks classiques. Dans les frameworks classiques, le calcul suit les données. Si le tableau est en mémoire CPU, le calcul se fait sur le CPU. Si le tableau est en mémoire GPU, le calcul se fait sur le GPU.

MLX alloue les tableaux en mémoire unifiée. Pas besoin de les copier pour les utiliser sur les différents appareils pris en charge.

Pour exécuter une opération, il suffit de spécifier l’appareil directement à l’opération. Par exemple, ici, nous ajoutons a et b sur le GPU et multiplions a et b sur le CPU. Ces opérations peuvent s’exécuter en parallèle, et MLX gère automatiquement les dépendances si nécessaire. Autre point fort de MLX : le calcul paresseux. Pour gagner en efficacité, surtout avec de très gros calculs, MLX utilise un moteur d’exécution paresseuse.

Aucun calcul n’est effectué tout de suite lorsqu’on additionne deux tableaux. À la place, un graphe de calcul est construit, comme celui-ci.

Le tableau C n’a pas encore été calculé. Il n’est calculé que si vous en avez réellement besoin. Par exemple, afficher C ou le convertir en liste Python à partir de MLX déclenchera son calcul. Vous pouvez forcer explicitement l’évaluation du graphe en utilisant mx.eval. Le calcul paresseux a plusieurs atouts. Le fait de séparer la création et l’exécution du graphe permet à MLX d’optimiser les calculs avant leur exécution. Et avec le calcul paresseux, vous ne payez que ce que vous utilisez. Autre point fort de MLX : les transformations de fonctions. Grâce à elles, MLX devient un outil performant pour l’entraînement et l’optimisation des modèles d’apprentissage automatique.

Ce sont des transformations qui reçoivent une fonction et produisent une nouvelle fonction en retour.

MLX a plusieurs transformations de fonction. On les classe en général en 2 catégories. Les transformations pour la différentiation automatique servent à optimiser le graphe de calcul. Par exemple, MLX permet, grâce à une transformation de fonction, de calculer automatiquement le gradient de toute fonction.

Supposons une fonction simple qui calcule le sinus de son entrée. Pour calculer le gradient de cette fonction, utilisez la transformation de fonction mx.grad. La fonction mx.grad renvoie une nouvelle fonction qui, appliquée à un tableau, fournit la dérivée. Les transformations de fonctions sont librement composables. Vous pouvez calculer la dérivée seconde en appliquant simplement mx.grad au résultat de mx.grad.

Le résultat est une autre fonction qui, appliquée à un tableau, renvoie la dérivée seconde du sinus.

MLX propose également 2 packages avancés pour créer et optimiser des réseaux neuronaux.

Le 1er, mlx.nn, une bibliothèque modulaire pour construire des réseaux neuronaux. Le second, mlx.optimizers, une bibliothèque d’algorithmes d’optimisation courants. Les deux packages peuvent être utilisés séparément, mais s’intègrent également de manière fluide.

Le package mlx.nn offre tout le nécessaire pour créer des réseaux neuronaux. nn.module est la classe de base principale dont héritent toutes les couches et tous les conteneurs. Il propose des méthodes pratiques pour accéder, charger, enregistrer les paramètres, etc.

La bibliothèque nn propose un ensemble de couches standard prêtes à l’emploi, comme nn.Linear, mais vous pouvez créer les vôtres en héritant de nn.Module.

Les fonctions de perte courantes et les méthodes d’initialisation sont également disponibles dans les sous-modules nn.losses et nn.init.

Voyons comment créer un réseau neuronal multicouche simple avec mlx.nn.

La première étape consiste à créer une classe personnalisée qui hérite de nn.Module. Dans ce cas, nous utiliserons un simple réseau neuronal à une couche cachée. On crée les couches linéaires dans la méthode d’initialisation du module. On implémente ensuite la fonction call, qui calcule la sortie du module pour une entrée donnée. La fonction call applique la première couche linéaire, la fonction d’activation relu, puis la seconde couche linéaire. Bien que MLX soit conçu et optimisé pour l’architecture à mémoire unifiée de la puce Apple, sa bibliothèque de réseaux neuronaux de haut niveau est également pensée pour être familière, avec une structure similaire à celle de frameworks courants comme PyTorch.

Comparons l’implémentation du modèle avec MLX à celle réalisée avec PyTorch. Elles sont presque identiques, à deux petites différences près dans la fonction qui calcule la sortie. Si vous avez déjà construit des modèles dans PyTorch, le passage à MLX devrait être très simple. Maintenant que vous avez vu la plupart des fonctionnalités de base de MLX voyons comment l’utiliser pour optimiser vos tâches d’apprentissage automatique. Commençons par voir comment compiler des fonctions pour les accélérer. Nous verrons ensuite mx.fast, qui propose des implémentations rapides prêtes à l’emploi des opérations courantes en apprentissage automatique et une API pour ajouter des noyaux Metal personnalisés. Puis, nous verrons comment utiliser la quantification pour réduire l’utilisation de la mémoire et accélérer l’exécution du modèle. Enfin, nous verrons comment utiliser MLX pour répartir un calcul sur plusieurs machines. Presque tous les calculs réalistes avec MLX sont composés de fonctions effectuant plusieurs opérations sur des tableaux. Une façon simple d’accélérer ces fonctions est d’utiliser la transformation mx.compile. Supposons que vous ayez une fonction qui effectue quelques opérations par élément, comme la fonction d’activation GELU illustrée ici. Le graphe de calcul de cette fonction contient plusieurs nœuds. Chacun de ces nœuds correspond en coulisses à un lancement distinct de noyau GPU.

La compilation du graphe fusionne tous ces noyaux distincts en un seul noyau optimisé. Cela permet d’économiser de la mémoire, de réduire la surcharge d’exécution du graphe et d’obtenir un calcul plus efficace.

Utiliser mx.compile, c’est aussi simple que d’ajouter le décorateur mx.compile à votre fonction. La compilation est efficace, mais pour des opérations plus complexes, y compris certaines opérations courantes dans l’apprentissage automatique, mx.fast est souvent plus performant. Par exemple, de nombreux composants clés d’un modèle Transformer utilisent des opérations de mx.fast. Cela inclut les encodages positionnels, les couches de normalisation et l’attention à produit scalaire. Les opérations de mx.fast sont plus spécialisées, mais optimisées pour offrir des performances maximales en entraînement comme en inférence. Elles sont aussi très configurables, ce qui permet d’accélérer de nombreuses variantes d’un même calcul. Par exemple, l’opération d’attention à produit scalaire peut accepter un masque en entrée. Le masque peut être additif, booléen ou une chaîne de caractères indiquant le type de masque. Examinons de plus près la norme RMS, qui est l’une des opérations de mx.fast. Elle est utilisée dans presque tous les grands modèles de langage modernes basés sur des transformateurs. Une implémentation simple avec les opérations MLX génère un graphe de calcul volumineux.

À la place, on peut remplacer toute l’implémentation par un simple appel à mx.fast.rms_norm. Le code est plus simple, le graphe de calcul n’a qu’un seul nœud et le calcul lui-même sera beaucoup plus rapide.

MLX propose une API pour ajouter des noyaux Metal personnalisés, utile si votre fonction nécessite une implémentation sur mesure qui n’existe pas encore dans mx.fast. Vous écrivez le noyau personnalisé et MLX gère tout le reste, y compris la compilation et l’exécution juste-à-temps. Ces noyaux sont écrits en Metal, un langage et une API développés par Apple pour exécuter des fonctions sur les GPU Apple. Vous créez le noyau en fournissant une chaîne de code Metal en entrée, ainsi que quelques informations sur les entrées et sorties.

Vous appelez le noyau en spécifiant la taille de la grille et les formes et types de sortie. MLX traite l’appel du noyau comme n’importe quelle autre opération. Il crée un nœud dans le graphe de calcul, qui sera évalué de manière paresseuse.

Un autre outil pour accélérer vos charges de travail en apprentissage automatique est la quantification. Les grands modèles nécessitent beaucoup de mémoire et de bande passante pour être rapides. Dans bien des cas, la précision requise pour l’entraînement est bien plus élevée que celle nécessaire en inférence pour obtenir une qualité équivalente ou presque. Réduire la précision permet de faire tenir des modèles plus grands en mémoire et de les exécuter plus rapidement.

Si votre modèle est en précision flottante 32 bits, vous pouvez passer au bfloat16 ou float16 pour réduire de moitié les besoins en mémoire. Quand 16 bits sont encore trop, MLX propose des routines intégrées pour quantifier les tableaux à une taille encore plus réduite et effectuer des opérations avec eux. Par exemple, vous pouvez quantifier en utilisant 4 bits par élément pour réduire encore plus les besoins en mémoire.

Pour quantifier une matrice, utilisez mx.quantize. Vous spécifiez le nombre de bits à utiliser par élément et la taille du groupe. La taille de groupe détermine combien d’éléments de la matrice quantifiée partagent une même valeur d’échelle et de biais. Moins vous utilisez de bits et plus la taille du groupe est grande, plus le résultat est compact et rapide. MLX propose plusieurs options de précision et de taille de groupe pour vous offrir un maximum de flexibilité.

Vous pouvez multiplier n’importe quel vecteur ou matrice non quantifié par votre matrice quantifiée avec mx.quantized_matmul.

Utilisez mx.dequantize pour récupérer une approximation de l’entrée d’origine.

MLX.mn propose aussi un utilitaire pratique pour quantifier un module en une seule commande. Supposons que vous ayez un modèle qui est une pile d’une couche d’intégration suivie de plusieurs couches linéaires. Vous pouvez quantifier le modèle avec nn.quantize. La commande de quantification accepte également un rappel facultatif pour vous permettre d’avoir un contrôle plus précis sur les couches à quantifier et la précision à utiliser pour une couche donnée.

Lors de la génération de texte avec de grands modèles de langage, la quantification réduit considérablement l’utilisation de la mémoire et augmente les jetons générés par seconde. Dans certains cas, une seule machine ne suffit pas. Vous pouvez vouloir générer du texte avec un grand modèle qui ne tient pas en mémoire sur une seule machine. Ou vous pouvez affiner ou évaluer un modèle sur un grand jeu de données, deux tâches faciles à paralléliser et nettement plus rapides sur plusieurs machines.

MLX est prêt à l’emploi avec la possibilité d’exécuter des calculs arbitraires sur plusieurs machines. Les machines peuvent être connectées via Ethernet ou Thunderbolt.

Vous utilisez le sous-package mx.distributed pour répartir les calculs sur plusieurs machines. MX distribué est principalement un ensemble d’opérations de communication.

Par exemple, all_Sum ajoute les tableaux d’entrée sur toutes les machines. La sortie de all_Sum est l’entrée additionnée, qui est la même pour chaque machine. Voyons de plus près comment additionner un tableau sur plusieurs machines. Initialisez le backend distribué à l’aide de mx.distributed.init. Cette étape est facultative, mais elle sert à accéder au groupe de communication si nécessaire. Le groupe de communication contient des informations utiles, comme le nombre total de processus et l’indice du processus en cours.

Puis, créez un tableau avec une seule valeur sur chaque processus et utilisez mx.distributed.all_sum pour additionner ces tableaux sur tous les processus.

MLX dispose d’un lanceur pratique pour exécuter un programme MLX sur plusieurs machines. Pour exécuter un programme sur 4 machines, utilisez mlx.launch avec les adresses IP des 4 hôtes. Tout ce que nous avons vu jusqu’ici utilise MLX en Python. Dans de nombreux cas, vous préférerez la simplicité et la flexibilité de Python. Dans d’autres cas, vous pourriez préférer Swift. C’est pourquoi MLX propose une API complète en Swift. Il est construit sur Metal et fonctionne parfaitement sur macOS, iOS, iPadOS, visionOS, et d’autres plates-formes. Commencer avec MLX en Swift est aussi simple que de l’ajouter en tant que package à votre projet Xcode. Cliquez sur le projet, puis sur le signe + dans l’onglet des dépendances des packages. Entrez ensuite l’URL du référentiel GitHub MLX Swift et cliquez sur le bouton Add package. C’est tout ce qu’il faut pour commencer à développer avec MLX en Swift. Pour faciliter au maximum la transition entre Python et Swift, les API des deux langages sont volontairement similaires. Voici une comparaison côte à côte d’un extrait de code Python vu précédemment avec son équivalent en MLX Swift. Créer des tableaux, effectuer des opérations dessus et consulter leurs métadonnées se fait presque de la même manière en Swift qu’en Python. Toutes les fonctionnalités principales vues avec MLX en Python, ainsi que les optimisations abordées dans la section sur l’accélération de MLX, s’appliquent également à MLX Swift. Nous avons passé en revue de nombreuses fonctionnalités clés de MLX. Pour en savoir plus sur le framework, consultez le site web de MLX, qui contient des liens vers la documentation, des exemples, etc. Les API Python et Swift ont toutes deux un référentiel d’exemples, contenant de nombreux cas d’usage courants en apprentissage automatique, comme l’entraînement et la génération de modèles de langage, la génération d’images, la reconnaissance vocale, et bien d’autres. Ces exemples sont un excellent moyen d’en savoir plus sur MLX et un bon point de départ pour intégrer MLX dans vos propres projets. Merci de votre attention, j’ai hâte de voir ce que vous allez créer avec MLX.