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

Dit paper introduceert Mozart, een co-ontwerp van algoritme en hardware dat de training van MoE-gedreven taalmodellen op 3.5D wafer-scale chiplet-architecturen optimaliseert door middel van een slimme experttoewijzing, fijnmazige planning en een hiërarchisch geheugenontwerp om communicatie- en geheugenproblemen op te lossen.

De kern

Stel je voor dat je een enorm groot brein wilt bouwen om slimme computers (zoals de AI die we vandaag gebruiken) te laten leren. Dit brein bestaat uit miljoenen kleine onderdelen die samenwerken. In de wereld van kunstmatige intelligentie noemen we dit een MoE-model (Mixture of Experts). Het idee is simpel: in plaats dat één gigantische "super-intelligentie" alles doet, heb je duizenden kleine "experts" die elk gespecialiseerd zijn in één ding. Als de computer een vraag krijgt, kijkt hij eerst welke experts nodig zijn en schakelt die dan in.

Het probleem is echter: deze experts zitten vaak verspreid over verschillende plekken in de computer. Het is alsof je een team van duizenden specialisten hebt, maar ze zitten allemaal in verschillende gebouwen in een enorme stad. Als ze moeten samenwerken, moeten ze constant bellen, post sturen en reizen. Dit kost veel tijd en energie, en het remt het leerproces van de AI af.

De auteurs van dit paper, Mozart, hebben een oplossing bedacht. Ze noemen hun idee "Mozart" (een knipoog naar de componist die bekendstaat om zijn complexe maar harmonieuze muziek). Ze hebben een nieuwe manier bedacht om deze AI-modellen te trainen, door de software (de regels) en de hardware (de fysieke computer) perfect op elkaar af te stemmen.

Hier is hoe ze dat doen, vertaald in alledaagse termen:

1. De Stad op een Chip: De "3.5D Wafer-Scale" Architectuur

Stel je een traditionele computerchip voor als een klein dorpje. Mozart bouwt in plaats daarvan een gigantische stad op één enkele plaat (een "wafer").

• Het idee: In plaats van dat de experts ver weg wonen, bouwen ze hun huizen (de chiplets) heel dicht bij elkaar.
• De structuur: Ze gebruiken een speciale 3D-structuur. Denk aan een wolkenkrabber waar de verdiepingen (de logica en het geheugen) direct op elkaar zijn gelijmd, in plaats van dat ze met lange kabels verbonden moeten worden. Dit maakt het reizen voor de data extreem snel.
• De straten: Ze hebben een slim stratenplan (een "NoP-Tree") ontworpen. De belangrijkste verkeersknooppunten (de "switches") zitten in het midden, zodat data niet hoeft te slenteren om van A naar B te komen.

2. De Slimme Buurman: Expert Clustering

In de oude wereld werden de experts willekeurig over de stad verdeeld. Soms zaten twee experts die vaak samenwerken, aan de andere kant van de stad.

• De Mozart-oplossing: Ze kijken eerst naar hoe de AI werkt. Ze zien dat bepaalde experts vaak samen worden ingeschakeld (bijvoorbeeld een expert voor "wiskunde" en een voor "grammatica" werken vaak samen).
• De strategie: Ze verplaatsen deze "buren" naar hetzelfde huisje of naar huizen die direct naast elkaar liggen. Hierdoor hoeven ze niet meer de hele stad te reizen om te praten; ze kunnen gewoon door de muur bellen. Dit bespaart enorm veel tijd.

3. De Vloeiende Stroom: Streaming en Overlappen

Stel je voor dat een vrachtwagen (de data) moet wachten tot de fabriek (de computer) klaar is met een product, en dan pas de volgende vrachtwagen mag komen. Dat is inefficiënt.

• De Mozart-oplossing: Ze maken het proces als een vloeiende stroom. Terwijl de ene groep experts aan het rekenen is, wordt de volgende groep experts alvast naar de fabriek gebracht.
• Het resultaat: De computer is nooit stil. Terwijl hij rekent, worden er alvast nieuwe gegevens opgehaald. Het is alsof een kok die terwijl hij een bord afwast, alvast de volgende ingrediënten snijdt. Alles gebeurt tegelijkertijd.

Waarom is dit belangrijk?

Tot nu toe was het trainen van deze super-slimme AI-modellen erg traag en duur, omdat de "reistijd" van de data te lang was.

• Het effect: Met Mozart kunnen deze modellen twee keer zo snel leren als met de oude methoden.
• De analogie: Het is alsof je van een oude, smalle landweg (waar je vastloopt in file) verhuist naar een supersnelle, geautomatiseerde magneettrein (de 3.5D-chip).

Kort samengevat:
Mozart is een slimme manier om de "experts" van een AI dichter bij elkaar te zetten en hun werk te coördineren, zodat ze niet hoeven te wachten op elkaar. Door de software en de fysieke computerchips samen te ontwerpen, maken ze het trainen van de slimste AI's van de wereld veel sneller, goedkoper en efficiënter. Het is alsof je een orkest organiseert waarbij alle muzikanten perfect op elkaar afgestemd zijn, zodat de muziek (de AI) vloeiend en snel klinkt.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures", vertaald en samengevat in het Nederlands.

Probleemstelling

Mixture-of-Experts (MoE) architecturen hebben de efficiëntie van Large Language Models (LLMs) aanzienlijk verbeterd door modulaire berekeningen te gebruiken. Echter, de inherente sparsiteit en dynamische aard van MoE-LLMs vormen grote uitdagingen voor traditionele hardware (zoals GPUs en CPUs):

1. Geheugenlokaliteit: De dynamische toewijzing van tokens aan experts leidt tot slechte cache-efficiëntie en hoge toegangstijden tot het geheugen.
2. Communicatie- overhead: De "All-to-All" communicatie tussen parallelle eenheden (waarbij tokens worden gerouteerd naar de juiste experts) is een grote bottleneck, vaak beperkt door bandbreedte.
3. Inefficiënt hulpbronnengebruik: Traditionele hardware is vaak ontworpen voor dichte, uniforme berekeningen, wat leidt tot onderbenutting bij de ongelijke en dynamische werklasten van MoE-modellen.
4. Schalingsgrenzen: Monolithische chips bereiken hun limieten qua oppervlakte (photoreticle) en transistor-schaling.

Methodologie: Mozart Framework

Mozart is een algoritme-hardware co-design framework dat specifiek is ontworpen voor het trainen van MoE-LLMs op 3.5D wafer-scale chiplet-architecturen. De aanpak combineert software-optimalisaties met een nieuwe hardware-topologie.

1. Algoritme-niveau Optimalisaties

• Expert Clustering en Allocatie:
  • In plaats van experts willekeurig te verdelen, analyseert Mozart eerst de "co-activatie patronen" van experts op een instructie-tuning dataset.
  • Experts die vaak samen worden geactiveerd (co-activated), worden gegroepeerd in clusters.
  • Deze clusters worden strategisch toegewezen aan specifieke chiplets om de All-to-All communicatie te minimaliseren. Als twee experts die samen nodig zijn op dezelfde chiplet zitten, hoeft er geen data over de chiplet-grenzen heen te worden verplaatst.
• Fin-grained Scheduling (Streaming):
  • Om de communicatie-overhead met het DRAM-geheugen te maskeren, introduceert Mozart een streaming-mechanisme.
  • Streaming Tokens: Tokens worden opgesplitst in micro-batches.
  • Streaming Experts: Experts worden geladen op basis van hun werklast (werklast-georiënteerd). Experts met een hoge activatiekans worden eerst geladen.
  • Dit zorgt voor een maximale overlap tussen communicatie (data laden uit DRAM) en berekening (on-chip verwerking), waardoor de wachttijd voor het geheugen wordt geminimaliseerd.

2. Hardware-niveau Architectuur (3.5D Wafer-Scale)

Mozart introduceert een nieuwe hardware-architectuur die bestaat uit:

• 3.5D Integratie: Verticale stapeling van logica en geheugen (Logic-on-Memory) via hybride bonding. Dit verhoogt de bandbreedte en verlaagt de latentie voor frequent gebruikte data (activaties).
• 2.5D NoP-Tree Topologie (Network-on-Package):
  • De chiplets zijn verbonden via een boom-achtige structuur.
  • Attention Chiplets (geheugen-gebonden) zitten centraal als dispatchers.
  • Expert Chiplets (berekenings-gebonden) zitten aan de randen als bladeren.
  • Switches met in-network computing capaciteit aggregatie van MoE-uitvoer, wat de communicatie tussen chiplets verder reduceert.
• Hiërarchisch Geheugen:
  • DRAM: Bewaart de statische modelgewichten (off-chip).
  • SRAM: Fungeert als lokaal cache voor tijdelijke activaties op de chiplets, wat snelle toegang mogelijk maakt.

Belangrijkste Bijdragen

1. Co-design Framework: Een uniek systeem dat algoritmen (expert-plaatsing, scheduling) en hardware (3.5D chiplets, NoP-Tree) naadloos integreert voor MoE-training.
2. Expert Collaboration Strategy: Een methode om experts te clusteren op basis van co-activatie patronen, wat de noodzakelijke data-overdracht tussen chiplets drastisch vermindert.
3. Nieuwe Hardware Topologie: Een wafer-scale 3.5D architectuur met een 2.5D NoP-Tree interconnect die specifiek is ontworpen voor de heterogene en dynamische aard van MoE-werklasten.
4. Fin-grained Scheduling: Een pipeline-strategie die communicatie en berekening overlapt, waardoor de latentie van het geheugen wordt gemaskeerd.

Resultaten

Het team heeft Mozart geëvalueerd op drie populaire open-source MoE-LLMs: Qwen3-30B-A3B, OLMoE-1B-7B-0924, en DeepSeek-MoE-16B-Base.

• Snelheidswinst: Mozart bereikte een 1.9x tot 2.37x versnelling in post-training latentie in vergelijking met baseline-methoden (zonder de specifieke Mozart-optimalisaties).
  • Bijvoorbeeld: 2.37x versnelling voor OLMoE-1B-7B-0924.
• Communicatie-efficiëntie: De "All-to-All" communicatiecomplexiteit ($C_T$) werd aanzienlijk verlaagd door de slimme expert-plaatsing, wat direct leidde tot lagere eind-tot-eind latentie.
• Robuustheid: De prestaties bleven stabiel bij toenemende sequentielengtes (van 128 tot 512 tokens), terwijl baseline-methoden hier sterk onder leden door communicatie-bottlenecks.
• Geheugen vs. Berekening: De analyse toonde aan dat het systeem primair geheugen-gebonden (memory-bound) is door het laden van gewichten, maar dat Mozart dit effectief beheert door parallelisme en overlapping.

Betekenis en Impact

Mozart adresseert een kritieke lacune in de huidige AI-hardware-ontwikkeling: het schaalbaar maken van MoE-LLMs.

• Schalbaarheid: Het biedt een pad naar het trainen van modellen met biljoenen parameters die niet meer passen op traditionele monolithische chips.
• Efficiëntie: Door de modulaire aard van MoE te matchen met de modulaire aard van chiplet-architecturen, wordt energie-efficiëntie en hulpbronbenutting gemaximaliseerd.
• Toekomstgerichtheid: Het paper demonstreert dat de toekomst van groot-schalige AI-training ligt in co-design, waarbij hardware en algoritmen gezamenlijk worden ontwikkeld om de beperkingen van de wet van Moore en de complexiteit van MoE-modellen te overwinnen.

Kortom, Mozart bewijst dat wafer-scale chiplet-architecturen, wanneer gecombineerd met slimme algoritmen voor expert-plaatsing en scheduling, een krachtige oplossing zijn voor de volgende generatie Large Language Models.