Fast MoE Inference via Predictive Prefetching and Expert Replication

Technische Samenvatting: Snelle MoE-inferentie via voorspellende prefetching en expert-replicatie

Probleemstelling
De Mixture of Experts (MoE)-architectuur is een hoeksteen geworden voor het schalen van Large Language Models (LLMs) door enorme parameteraantallen mogelijk te maken zonder evenredige toename van de rekenkosten. Echter, MoE-inferentie kampt met kritieke inefficiënties die een praktische implementatie belemmeren:

1. Suboptimale GPU-gebruik: Door de inherente spaarzaamheid van expert-activering is slechts een klein deel van de experts op elk moment actief, waardoor het merendeel van de GPU-bronnen inactief blijft.
2. Ongelijkmatige belasting en latentie: Dynamische routing leidt vaak tot een ongelijke verdeling van de belasting, waarbij meerdere tokens in de wachtrij staan voor dezelfde fysieke expert, wat resulteert in token-stalling en verhoogde latentie.
3. Routing-overhead: Traditionele op softmax gebaseerde routingmechanismen introduceren rekencomplexiteit en communicatie-overhead (all-to-all routing) die niet gunstig schalen met toenemende expertaantallen.

Bestaande optimalisatietechnieken, zoals distillatie of kernel-niveau verbeteringen, bieden vaak beperkte verbeteringen of gaan ten koste van de modelnauwkeurigheid.

Methodologie: MoE-MPMC
De auteurs stellen MoE-MPMC (Mixture of Experts met Multiple Producers Multiple Consumers) voor, een raamwerk dat deze inefficiënties aanpakt via voorspellende prefetching en dynamische expert-replicatie. Het systeem werkt met twee gelijktijdige threads:

1. Inferentiethread:

   • Haalt een vooraf berekende hash-tabel op die expert-toewijzingen en replicatieaantallen bevat voor de huidige batch.
   • Laadt dynamisch alleen de voorspelde actieve experts in het GPU-geheugen en verplaatst ongebruikte experts naar CPU-geheugen.
   • Voert de inferentie uit met een aangepaste MoE-laag waarbij actieve experts fysiek worden gerepliceerd op de GPU. Dit maakt het mogelijk dat meerdere tokens die dezelfde expert nodig hebben gelijktijdig worden verwerkt door distincte fysieke replica's, waardoor wachttijden worden geëlimineerd.

2. Hash-bouwingst thread:

   • Draait parallel met de inferentiethread en verwerkt de volgende batch tokens.
   • Gebruikt een Simple Recurrent Unit (SRU) met spaarzame attention om te voorspellen welke experts voor de aankomende batch zullen worden geactiveerd.
   • Vervangt de op LSTM gebaseerde predictor uit eerder werk (SiDA-MoE) door SRU om snellere trainings- en inferentiesnelheden te bereiken door verminderde onderlinge afhankelijkheden en eenvoudigere gating-mechanismen.
   • Berekent replicatieaantallen voor elke voorspelde expert op basis van token-vraag en werkt de hash-tabel bij, die vervolgens in de wachtrij wordt geplaatst voor de inferentiethread.

Belangrijkste Bijdragen
Het artikel schetst drie primaire bijdragen:

1. Op SRU gebaseerde Voorspelling: De introductie van de Simple Recurrent Unit (SRU) om expert-selecties voor elke token snel te voorspellen. Dit vervangt langzamere LSTM-predictors, wat de voorspellingsfase aanzienlijk versnelt terwijl de nauwkeurigheid behouden blijft.
2. Diepe Fysieke Replicatie: Een strategie om voorspelde experts fysiek op de GPU te repliceren op basis van token-vraag. Dit zorgt ervoor dat bijna geen enkele token wachttijden ondervindt tijdens batch-verwerking, waardoor parallelisme en GPU-gebruik worden gemaximaliseerd.
3. Capaciteitsbegrenzing: Een praktische mechanisme om expert-replicatieaantallen te beperken volgens vooraf gedefinieerde GPU-geheugenbeperkingen. Dit zorgt ervoor dat replicatie binnen hardwarelimieten wordt geoptimaliseerd om bijna optimale resourcegebruik te bereiken zonder overloop te veroorzaken.

Experimentele Resultaten
De auteurs hebben MoE-MPMC geëvalueerd op grote schaal SwitchTransformers-modellen (Switch-base-8, 64, 128 en 256) met datasets zoals SST2 en MRPC. Belangrijkste bevindingen zijn:

• GPU-gebruik: De methode bereikt een bijna volledig GPU-gebruik (tot wel ~100% in sommige configuraties), een aanzienlijke verbetering ten opzichte van de baseline SwitchTransformers (die een gebruik liet zien zo laag als 1,3% voor 128 experts) en de SiDA-MoE-aanpak.
• Inferentiesnelheid: De aanpak levert tot een 3x verbetering in inferentiesnelheid op ten opzichte van baseline-architecturen.
• Nauwkeurigheid: De methode behoudt ongeveer 90-95% van de prestaties van baseline-architecturen, met een betere behoud van nauwkeurigheid dan SiDA-MoE vanwege de superieure voorspellingscapaciteiten van SRU ten opzichte van LSTM.
• Doorvoer: De doorvoer neemt aanzienlijk toe naarmate het aantal experts toeneemt, gedreven door het vermogen om meerdere tokens parallel te verwerken via expert-replicatie.

Betekenis en Claims
Het artikel claimt dat MoE-MPMC een belangrijke stap voorwaarts is naar de praktische en economisch haalbare implementatie van op MoE gebaseerde grote taalmodellen. Door te verschuiven van reactieve load balancing naar een proactieve strategie van voorspellende prefetching en replicatie, doet de methode het volgende:

• Vermindert memory access-bottlenecks en token-stalling.
• Optimaliseert GPU-resourcegebruik zonder de onderliggende MoE-berekening te wijzigen.
• Toont aan dat het voorspellen van expert-selecties haalbaar en effectief is wanneer fine-tuning en pretraining gelijktijdig plaatsvinden met token-voorspellings taken.

De auteurs positioneren dit werk als een vooruitgang ten opzichte van reactieve methoden (zoals SiDA-MoE) door toekomstige vraag te anticiperen en resource-toewijzing dynamisch aan te passen, waardoor de typische valkuilen van inactief GPU-geheugen en ongelijke belasting in spaarzame architecturen worden aangepakt.