Extension efficace des modèles de langage à grande échelle avec un mélange d’experts et le calcul analogique en mémoire 3D

Contexte académique

Ces dernières années, les modèles de langage à grande échelle (Large Language Models, LLMs) ont démontré des capacités impressionnantes dans le traitement du langage naturel et la génération de texte. Cependant, à mesure que la taille des modèles augmente, les coûts d’entraînement et d’inférence augmentent également de manière significative, notamment en termes de mémoire utilisée, de latence de calcul et de consommation d’énergie. Cela constitue l’un des principaux freins à l’adoption généralisée des LLMs. L’architecture traditionnelle de von Neumann, utilisée pour gérer des paramètres à grande échelle, implique des déplacements fréquents de données entre la mémoire et les unités de calcul, ce qui entraîne un “goulot d’étranglement de von Neumann”, exacerbant ces défis.

Pour résoudre ce problème, les chercheurs ont exploré plusieurs pistes technologiques, dont l’architecture dite “Mélange d’Experts” (Mixture of Experts, MoE). Le MoE utilise un mécanisme de calcul conditionnel (conditional computing) pour activer dynamiquement une partie des paramètres du modèle, réduisant ainsi les besoins en calcul. Cependant, le déploiement des modèles MoE repose toujours sur des architectures matérielles traditionnelles, ne résolvant pas entièrement le problème d’accès aux paramètres. Parallèlement, le calcul analogique en mémoire (Analog In-Memory Computing, AIMC) est apparu comme une technologie émergente qui permet de calculer directement dans la mémoire, évitant ainsi les déplacements de données et offrant une meilleure efficacité énergétique. La combinaison du MoE avec l’AIMC, en particulier avec l’utilisation de la mémoire non volatile tridimensionnelle (3D Non-Volatile Memory, 3D NVM), pourrait ouvrir une nouvelle voie pour l’extension des LLMs.

Cet article s’inscrit dans ce contexte et explore comment déployer l’architecture MoE sur du matériel de calcul analogique en mémoire 3D, tout en évaluant son potentiel pour réduire les coûts d’inférence des modèles de langage à grande échelle.

Source de l’article

Cet article a été rédigé par Julian Büchel et Athanasios Vasilopoulos d’IBM Research Europe, en collaboration avec d’autres experts d’IBM Almaden Research Center, Micron Technology et d’autres institutions. Il a été publié en janvier 2025 dans la revue Nature Computational Science sous le titre Efficient Scaling of Large Language Models with Mixture of Experts and 3D Analog In-Memory Computing.

Processus de recherche et résultats

1. Objectifs et cadre de recherche

L’objectif principal de cet article est d’explorer le déploiement de l’architecture MoE sur du matériel de calcul analogique en mémoire 3D et d’évaluer son potentiel pour réduire les coûts d’inférence des modèles de langage à grande échelle. Les chercheurs ont d’abord analysé les limites des LLMs traditionnels sur l’architecture de von Neumann, en mettant l’accent sur les goulots d’étranglement liés à l’accès aux paramètres et aux déplacements de données. Ensuite, ils ont proposé de combiner le MoE avec l’AIMC 3D, considérant que cette combinaison pourrait résoudre efficacement le problème d’accès aux paramètres et réduire la consommation d’énergie et la latence lors de l’inférence.

2. Simulation du système AIMC 3D

Pour évaluer les performances du MoE sur du matériel AIMC 3D, les chercheurs ont conçu un cadre de simulation abstrait pour un système AIMC 3D. Ce système est composé de plusieurs unités de mémoire 3D (tiles), chacune contenant des réseaux de mémoire non volatile multicouches (tiers). Dans la simulation, les chercheurs ont mappé les paramètres du modèle MoE sur ces unités de mémoire et ont évalué les performances d’inférence et la consommation d’énergie du modèle par simulation.

• Conception du cadre de simulation : Le simulateur a été implémenté en Python en utilisant les bibliothèques PyTorch et torch.fx pour définir l’architecture du modèle et le flux de données. Les chercheurs ont développé des modules de simulation personnalisés pour prendre en charge le mappage et l’exécution des modèles MoE. Le simulateur peut enregistrer le temps d’inférence, la consommation d’énergie et les besoins de mémoire de pointe.
• Cartographie et ordonnancement des modèles : Les chercheurs ont mappé différentes couches du modèle MoE sur le matériel AIMC 3D et ont utilisé un algorithme glouton pour optimiser la stratégie de mappage. Grâce à la simulation, ils ont découvert que les modèles MoE, grâce à leur mécanisme de calcul conditionnel, peuvent mieux exploiter la grande capacité de mémoire de l’AIMC 3D, réduisant ainsi les conflits entre les unités de calcul.

3. Comparaison entre MoE et les modèles denses

Pour évaluer les avantages du MoE, les chercheurs l’ont comparé aux modèles denses traditionnels. Les résultats montrent qu’avec l’augmentation du nombre de paramètres, le temps d’inférence des modèles MoE reste presque constant, alors que celui des modèles denses augmente de manière significative. Cela indique que l’architecture MoE peut étendre la taille du modèle en augmentant le nombre d’experts sans augmenter de manière significative la latence de calcul.

• Performances d’inférence : Dans la simulation, le temps d’inférence des modèles MoE était bien inférieur à celui des modèles denses, en particulier lorsque le nombre de paramètres atteint des centaines de milliards. Les chercheurs ont également constaté qu’avec l’augmentation du nombre d’experts, le temps d’inférence des modèles MoE augmente lentement, démontrant ainsi leur supériorité sur le matériel AIMC 3D.
• Consommation d’énergie et besoins en mémoire : Comme le calcul est effectué directement dans la mémoire sur le matériel AIMC 3D, la consommation d’énergie et les besoins en mémoire des modèles MoE sont nettement inférieurs à ceux des modèles denses. Les chercheurs ont noté que les besoins en mémoire de pointe des modèles MoE ne dépassent pas 1 Mo, ce qui est bien inférieur aux dizaines de Go nécessaires pour les modèles denses.

4. Comparaison des performances avec les GPU

Pour valider davantage les avantages du matériel AIMC 3D, les chercheurs l’ont comparé au GPU NVIDIA A100. Les résultats montrent que pour les grands modèles MoE, le débit (throughput) du matériel AIMC 3D est six fois supérieur à celui du GPU. De plus, l’efficacité énergétique du matériel AIMC 3D dépasse celle du GPU d’un facteur de 1000, démontrant ainsi sa supériorité dans le traitement des modèles de langage à grande échelle.

5. Robustesse du MoE au bruit matériel

Pour évaluer la robustesse des modèles MoE au bruit sur le matériel de calcul analogique en mémoire, les chercheurs ont effectué un entraînement “hardware-aware” (conscient du matériel). Les résultats montrent que les modèles MoE maintiennent une précision équivalente à celle des calculs en virgule flottante (iso-performance) même avec un niveau de bruit atteignant 6,3 %, ce qui indique une bonne robustesse sur le matériel AIMC 3D.

Conclusion et signification

Cette étude montre que la combinaison de l’architecture MoE avec le matériel de calcul analogique en mémoire 3D peut réduire de manière significative les coûts d’inférence des modèles de langage à grande échelle, en particulier en termes de consommation d’énergie et de latence. Grâce à son mécanisme de calcul conditionnel, le MoE peut mieux exploiter la grande capacité de mémoire de l’AIMC 3D, réduisant ainsi les conflits entre les unités de calcul. Par rapport aux modèles denses traditionnels et aux GPU, la combinaison du MoE avec l’AIMC 3D présente des avantages significatifs en termes de débit, d’efficacité énergétique et d’efficacité surfacique.

Cette recherche ouvre de nouvelles perspectives pour l’extension des modèles de langage à grande échelle, en particulier dans le domaine des coûts matériels et du calcul à haute efficacité énergétique. En combinant l’architecture MoE avec la technologie AIMC 3D, les chercheurs pourraient développer des modèles de langage à grande échelle plus efficaces et économiques, favorisant ainsi leur déploiement dans des applications pratiques.

Points forts de la recherche

1. Combinaison innovante d’architectures : Cet article est le premier à combiner l’architecture MoE avec le matériel de calcul analogique en mémoire 3D, proposant ainsi une nouvelle méthode pour résoudre les goulots d’étranglement liés aux coûts d’inférence des modèles de langage à grande échelle.
2. Réduction significative des coûts : Les résultats montrent que la combinaison du MoE avec l’AIMC 3D peut réduire de manière significative la consommation d’énergie et la latence d’inférence, en particulier lorsque le nombre de paramètres atteint des centaines de milliards.
3. Robustesse au bruit matériel : Grâce à l’entraînement “hardware-aware”, les modèles MoE maintiennent une précision élevée même avec des niveaux de bruit élevés, démontrant leur robustesse sur le matériel de calcul analogique en mémoire.

Autres informations utiles

Cet article a également rendu public le simulateur et le code d’implémentation des modèles MoE, permettant à d’autres chercheurs et développeurs de les utiliser. Cette initiative contribuera à faire progresser la recherche et les applications dans ce domaine.