Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures

Le papier présente Mozart, un cadre de co-conception algorithme-hardware qui optimise l'entraînement des modèles de langage à base d'experts (MoE) sur des architectures de puces 3,5D en exploitant leur modularité pour améliorer l'allocation des experts, le recouvrement communication-calcul et l'utilisation des ressources.

L'essentiel

🎻 Mozart : Le Chef d'Orchestre des Super-Ordinateurs

Imaginez que vous essayez de faire jouer un orchestre symphonique géant (un modèle d'intelligence artificielle très puissant) dans une petite chambre d'hôtel. C'est le problème actuel des modèles d'IA modernes : ils sont immenses, mais les ordinateurs qui les font tourner sont souvent mal organisés, ce qui crée du bruit, des retards et de la gaspillage d'énergie.

Les chercheurs ont créé Mozart, un nouveau système qui combine un nouveau logiciel et un nouveau matériel (des puces électroniques spécialisées) pour résoudre ce chaos. Le but ? Faire tourner ces géants de l'IA beaucoup plus vite et plus efficacement.

Voici comment cela fonctionne, en utilisant des analogies simples :

1. Le Problème : L'Orchestre "MoE" (Mélange d'Experts)

Les modèles d'IA modernes, comme ceux qui écrivent des textes ou répondent à des questions, utilisent une architecture appelée MoE (Mélange d'Experts).

• L'analogie : Imaginez un hôpital avec des milliers de médecins spécialistes (les "experts"). Quand un patient arrive (une phrase de texte), le système ne fait pas venir tous les médecins. Il en choisit seulement quelques-uns (par exemple, un cardiologue et un neurologue) pour traiter le cas.
• Le problème : Dans les ordinateurs actuels, ces médecins sont dispersés dans des bâtiments différents (des puces séparées). Pour qu'ils puissent travailler ensemble, il faut envoyer des dossiers par la poste entre les bâtiments. C'est lent, coûteux en énergie et crée des embouteillages. De plus, certains médecins sont très demandés, d'autres pas du tout, ce qui crée des goulots d'étranglement.

2. La Solution Matérielle : Le "Wafer-Scale" (La Grande Table de Cuisine)

Au lieu d'avoir des petits ordinateurs séparés, Mozart utilise une architecture appelée 3.5D Wafer-Scale.

• L'analogie : Au lieu d'avoir plusieurs petites tables de cuisine séparées par la pièce, Mozart prend une énorme dalle de verre unique (un "wafer" de silicium) et y installe tous les éléments directement dessus.
• La structure :
  • Les "îlots" (Chiplets) : Sur cette grande dalle, il y a des zones spécialisées. Certaines zones sont pour la mémoire (les bibliothèques), d'autres pour le calcul (les cuisiniers).
  • L'autoroute (NoP-Tree) : Au lieu de routes sinueuses, ils ont construit une autoroute en forme d'arbre très rapide qui relie tout le monde. Le centre de l'arbre gère les décisions, et les feuilles de l'arbre sont les experts qui travaillent.

3. La Solution Logicielle : L'Intelligence de Mozart

Mozart n'est pas seulement du matériel, c'est aussi une façon intelligente d'organiser le travail.

• A. Regrouper les amis (Clustering) :

  • Le concept : Mozart observe les habitudes des médecins. Il remarque que le cardiologue et le neurologue sont souvent appelés ensemble pour certains patients.
  • L'action : Au lieu de les laisser dans des bâtiments différents, Mozart les installe côte à côte sur la même petite puce. Ainsi, ils peuvent se passer des dossiers en chuchotant (communication ultra-rapide) au lieu d'envoyer des emails. Cela réduit énormément le temps perdu.

• B. Le flux continu (Streaming) :

  • Le concept : Imaginez un chef qui doit cuisiner un repas. S'il attend d'avoir tous les ingrédients avant de commencer à couper, il perd du temps.
  • L'action : Mozart utilise une technique de "flux". Pendant que le premier groupe d'ingrédients est en train d'être cuisiné, le système commence déjà à charger le deuxième groupe d'ingrédients. Il fait le calcul et le transport en même temps, comme un chef qui coupe des légumes pendant que l'eau bout. Cela évite que les cuisiniers restent inactifs en attendant les ingrédients.

4. Les Résultats : Plus Vite, Plus Économe

En combinant cette "grande dalle" intelligente avec cette organisation logique, les chercheurs ont obtenu des résultats impressionnants :

• Vitesse : Le système est 2 fois plus rapide que les méthodes actuelles pour entraîner ces modèles.
• Efficacité : Il gaspille beaucoup moins d'énergie car il évite les allers-retours inutiles de données.

En Résumé

Mozart est comme un chef d'orchestre génial qui a réorganisé une symphonie chaotique.

1. Il a mis tous les musiciens sur la même scène géante (le wafer) pour qu'ils soient proches.
2. Il a regroupé les musiciens qui jouent souvent ensemble sur le même pupitre pour qu'ils n'aient pas à crier pour se parler.
3. Il a fait en sorte que la musique ne s'arrête jamais, en faisant passer les notes (les données) pendant que les musiciens jouent.

Le résultat ? Une intelligence artificielle qui apprend beaucoup plus vite, consomme moins d'électricité et peut être déployée sur des systèmes de taille industrielle sans s'effondrer.

Résumé technique

Titre : Mozart : Entraînement Modulaire et Efficace de MoE sur des Architectures de Chiplets à Échelle de Wafer 3.5D

1. Problématique

Les modèles de langage de grande taille (LLM) basés sur l'architecture Mélange d'Experts (MoE) offrent une efficacité computationnelle supérieure en activant dynamiquement des sous-réseaux spécialisés (experts) plutôt que l'ensemble du modèle. Cependant, leur déploiement sur du matériel conventionnel (GPU/CPUs) rencontre des obstacles majeurs :

• Problèmes de localité mémoire : La nature dynamique et dispersée des activations des experts rend difficile la gestion efficace de la mémoire.
• Surcharge de communication : Le mécanisme de routage "All-to-All" (chaque token doit être envoyé à plusieurs experts répartis sur différents nœuds) crée un goulot d'étranglement majeur en termes de bande passante et de latence.
• Utilisation inefficace des ressources : Les charges de travail hétérogènes et irrégulières des MoE entraînent une sous-utilisation des ressources de calcul et des déséquilibres de charge entre les unités de traitement.
• Limites de l'intégration monolithique : Les puces traditionnelles atteignent leurs limites physiques (taille du photoréticule, mise à l'échelle des transistors), rendant nécessaire l'adoption d'architectures modulaires basées sur des chiplets.

2. Méthodologie : Co-conception Algorithme-Hardware

Mozart propose un cadre de co-conception (algorithme + matériel) spécifiquement adapté aux architectures 3.5D à échelle de wafer (intégration de chiplets sur une grande surface de silicium).

A. Côté Algorithme :

1. Clustering et Allocation des Experts :
   • Au lieu d'une répartition aléatoire, Mozart analyse les motifs d'activation (fréquence et co-activation) des experts sur un jeu de données de pré-entraînement.
   • Il regroupe les experts fréquemment co-activés sur les mêmes chiplets ou chiplets adjacents.
   • Cela permet d'équilibrer la charge de travail entre les groupes de chiplets et de réduire le volume de données échangées lors de la communication "All-to-All" (car moins de tokens doivent traverser les interconnexions si les experts sont colocalisés).
2. Ordonnancement Fin (Fine-grained Scheduling) :
   • Pour masquer la latence de communication DRAM, Mozart introduit un mécanisme de streaming (flux continu) des tokens et des experts.
   • Il découpe les lots de tokens en micro-lots et chevauche le chargement des poids des experts depuis la DRAM avec le calcul sur puce.
   • Les experts les plus activés sont chargés en priorité pour maximiser le recouvrement calcul/communication.

B. Côté Matériel (Architecture 3.5D) :

1. Topologie NoP-Tree 2.5D :
   • L'architecture utilise une topologie arborescente où les chiplets d'Attention (goulot de mémoire) sont placés au centre (nœuds de dispatch) et les chiplets d'Experts (goulot de calcul) en périphérie (feuilles).
   • Des modules de commutation (switches) intégrés au réseau permettent une agrégation locale des sorties des experts, réduisant le trafic global.
2. Hiérarchie Mémoire 3D (Logic-on-Memory) :
   • Chaque chiplet de calcul intègre une pile verticale de 3D (die logique + die SRAM) via un collage hybride (hybrid bonding).
   • Cela permet de mettre en cache les activations (données temporaires fréquemment réutilisées) très près du calculateur, réduisant drastiquement l'accès à la DRAM externe.
   • Les poids du modèle, statiques pendant une itération, restent en DRAM distribuée.

3. Contributions Clés

• Stratégie d'Allocation d'Experts : Une méthode basée sur l'analyse des motifs de co-activation pour placer les experts sur des chiplets spécifiques, optimisant ainsi la localité des données et réduisant la communication inter-chiplets.
• Architecture Chiplet 3.5D Spécifique : Conception d'une architecture matérielle hétérogène combinant des empilements 3D (mémoire/logique) et une interconnexion 2.5D en arbre, conçue spécifiquement pour les charges de travail MoE.
• Ordonnancement par Streaming : Un mécanisme logiciel permettant de chevauler efficacement les transferts de données (DRAM) et les calculs, atténuant les goulots d'étranglement de la bande passante.
• Framework de Co-conception : Une approche unifiée qui aligne la logique de routage des MoE avec la topologie physique du matériel.

4. Résultats Expérimentaux

Les auteurs ont évalué Mozart sur trois modèles MoE populaires de différentes échelles : Qwen3-30B-A3B, OLMoE-1B-7B, et DeepSeek-MoE-16B.

• Accélération : Mozart démontre une accélération significative par rapport aux méthodes de base (sans optimisations) :
  • 1.92x pour Qwen3-30B-A3B.
  • 2.37x pour OLMoE-1B-7B.
  • 2.17x pour DeepSeek-MoE-16B.
• Réduction de Latence : L'optimisation de la communication "All-to-All" et le chevauchement calcul/communication réduisent considérablement la latence d'entraînement, même avec des séquences longues (jusqu'à 512 tokens).
• Impact de la Bande Passante : Les gains sont encore plus prononcés avec une mémoire haute bande passante (HBM2) par rapport à une mémoire SSD, confirmant que l'architecture permet de mieux exploiter la bande passante disponible grâce au streaming.
• Analyse des Goulots : L'étude montre que le système est principalement limité par le chargement séquentiel des poids (bound mémoire) plutôt que par le calcul, ce que Mozart atténue efficacement via son ordonnancement.

5. Signification et Impact

Mozart représente une avancée majeure pour le déploiement de modèles de langage modulaires à grande échelle.

• Évolutivité : En exploitant les architectures à échelle de wafer, Mozart contourne les limites physiques des puces monolithiques, permettant l'entraînement de modèles avec des centaines de milliards de paramètres.
• Efficacité Énergétique et Temporelle : La réduction des mouvements de données inutiles et l'optimisation de l'utilisation des ressources de calcul améliorent l'efficacité énergétique globale.
• Nouveau Paradigme : Ce travail valide l'hypothèse que l'inspiration tirée de la modularité du cerveau humain (via les MoE) doit être couplée à une architecture matérielle modulaire (chiplets) pour atteindre un potentiel maximal. Il ouvre la voie à des systèmes de formation de LLM plus rapides, moins coûteux et plus scalables.

En résumé, Mozart résout le problème de la "désynchronisation" entre la nature dynamique et dispersée des modèles MoE et la rigidité des architectures matérielles traditionnelles, en offrant une solution logicielle et matérielle intégrée pour l'ère post-entraînement des grands modèles.