CRAFT: Cost-aware Expert Replica Allocation with Fine-Grained Layerwise Estimations

Dit paper introduceert CRAFT, een kostenefficiënt framework voor de replica-allokatie van experts in Mixture-of-Experts-modellen dat door middel van fijne-granulariteitsinschattingen per laag de belastingverdeling optimaliseert binnen een gegeven geheugenbudget, waardoor de doorvoersnelheid aanzienlijk wordt verhoogd zonder extra training of modelwijzigingen.

De kern

Stel je voor dat je een enorm groot restaurant runt, een "Super-Restaurant" genaamd MoE (Mixture of Experts). Dit restaurant is zo groot dat het duizenden koks (de "experts") heeft.

In een normaal restaurant werken alle koks aan hetzelfde menu. Maar in dit Super-Restaurant werkt het anders: elke klant (een stukje tekst of een vraag) krijgt een speciale menukaart. Een slimme ober (de "router") kijkt naar de bestelling en zegt: "Ah, deze klant wil een Italiaans gerecht, dus ik stuur die naar kok 45. Die klant wil sushi, dus die gaat naar kok 12."

Het probleem?
Soms zijn bepaalde gerechten (bijvoorbeeld pizza) extreem populair. Plotseling staan er honderden mensen in de rij bij Kok 45, terwijl Kok 12 (de sushi-kok) en Kok 99 (de dessert-kok) alleen maar naar hun handen zitten te staren.

• Kok 45 staat in de stress en maakt alles langzaam.
• De andere koks staan te wachten.
• Het hele restaurant vertraagt omdat iedereen moet wachten op de drukste kok.

In de wereld van kunstmatige intelligentie (AI) noemen we dit ongelijkmatige last. De computerchips (GPU's) die deze koks bedienen, raken overbelast, terwijl andere chips niets doen.

Het oude probleem: Te veel kopieën maken

Om dit op te lossen, dachten de managers van het restaurant: "Laten we gewoon meer koks voor de populaire gerechten in dienst nemen!"
Dus als Kok 45 het druk heeft, huren ze nog 10 extra koks in die exact hetzelfde doen als Kok 45. Ze noemen dit Expert Replication (het kopiëren van experts).

Maar hier zit een addertje onder het gras:
Elke nieuwe kok kost geld en ruimte in de keuken. De keuken (het computergeheugen) is al heel vol. Als je te veel extra koks huurt, is er geen ruimte meer voor de ingrediënten (de data) die nodig zijn om de bestellingen daadwerkelijk te maken. Het restaurant wordt zelfs langzamer, omdat de koks nu te veel tijd kwijt zijn aan het zoeken van een plekje in de volle keuken.

De oude methode (genaamd EPLB) was heel simpel: "Laten we bij elke laag van het restaurant precies evenveel extra koks inhuuren, ongeacht of ze het nodig hebben."
Dit resulteerde in een keuken vol met koks die niets te doen hadden, terwijl er toch nog steeds drukke plekken waren.

De nieuwe oplossing: CRAFT

De onderzoekers in dit paper hebben CRAFT bedacht. CRAFT staat voor iets als "Slimme Kostenefficiënte Koks-planning".

In plaats van blindelings overal extra koks in te huren, doet CRAFT het volgende:

1. Het Kijken naar de Historie:
   CRAFT kijkt eerst heel precies naar de bestellingen van gisteren. "Welke gerechten waren echt populair? Welke lagen (etages) in het restaurant hadden het meest te doen?"
   Ze ontdekten dat sommige etages al heel goed verdeeld waren (daar hoef je niets te doen), terwijl andere etages extreem scheef zaten (daar is echt hulp nodig).

2. Slimme Verdeling (De "Bakkerij"-analogie):
   Stel je voor dat je een bakkerij hebt met 10 ovens.

   • Oven 1 (de pizza-oven) staat altijd vol.
   • Oven 2 (de taart-oven) staat vaak leeg.
   • Oven 3 (de brood-oven) is soms vol, soms leeg.

   De oude methode zou zeggen: "Laten we bij elke oven 2 extra ovens bouwen." Dat kost enorm veel geld en ruimte.
   CRAFT zegt: "Wacht even. Oven 1 heeft 4 extra ovens nodig. Oven 3 heeft er 1 nodig. Oven 2 heeft helemaal geen extra ovens nodig, die staat al rustig."

   CRAFT verdeelt het budget (de ruimte in de keuken) precies daar waar het winst oplevert. Ze bouwen geen extra ovens waar ze niet nodig zijn.

3. Het Resultaat:

   • De keuken blijft minder vol (er is meer ruimte voor ingrediënten).
   • De drukke koks krijgen precies genoeg hulp.
   • Het restaurant draait soepeler en sneller.

Waarom is dit belangrijk?

In de echte wereld betekent dit dat bedrijven die enorme AI-modellen draaien (zoals die van Google of Meta) sneller kunnen werken en minder geld hoeven uit te geven aan dure computerchips.

• Vroeger: Je kocht 100 extra chips om de drukte op te vangen, maar je had er eigenlijk maar 60 nodig. De andere 40 stonden alleen maar te stoffen.
• Met CRAFT: Je koopt precies de 60 chips die je nodig hebt, en je plaatst ze slim. Je krijgt 15% meer snelheid voor hetzelfde geld, of zelfs meer snelheid voor minder geld.

Samenvattend in één zin:

CRAFT is als een slimme restaurantmanager die niet overal evenveel extra personeel huurt, maar precies weet waar de drukte zit en daar slim extra handen inzet, zodat het restaurant sneller draait zonder dat de keuken volloopt met onnodige koks.

Technische samenvatting

Probleemstelling

Mixture-of-Experts (MoE) modellen zijn de standaard geworden voor het schalen van grote taalmodellen (LLMs) omdat ze een hoge capaciteit bieden tegen een relatief lage rekenkost per token. Bij het implementeren van MoE-modellen wordt vaak Expert Parallelism (EP) gebruikt, waarbij experts (de feed-forward blokken) over verschillende GPU's worden verdeeld.

De kernproblemen die hierbij optreden zijn:

1. Token-Load Imbalans: De router in een MoE-model stuurt tokens dynamisch naar experts. In de praktijk volgt dit een Zipf-verdeling, waarbij een paar "hete" experts (hot experts) een overproportioneel aantal tokens ontvangen, terwijl andere "koude" experts (cold experts) weinig werk hebben.
2. Schaalproblemen bij EP: Deze onbalans leidt tot inactiviteit op GPU's met koude experts en netwerkcongestie, wat de inferentie-doorvoer (throughput) aanzienlijk verlaagt.
3. Beperkingen van bestaande oplossingen:
   • Expert Placement: Het slimmer toewijzen van experts aan apparaten helpt, maar faalt wanneer de load-verdeling extreem scheef is (één expert domineert).
   • Expert Replication (bijv. EPLB): Het kopiëren van hete experts naar meerdere apparaten lost de load-balans op, maar de huidige methoden (zoals EPLB) repliceren vaak uniform (één replica per laag per GPU). Dit leidt tot over-replicatie: veel replica's bieden weinig extra winst in load-balans, maar verbruiken wel kostbaar GPU-geheugen. Dit geheugenverbruik beperkt de grootte van de KV-cache, wat de doorvoer juist weer negatief beïnvloedt.

Methodologie: CRAFT

CRAFT (Cost-Aware Expert Replica Allocation with Fine-Grained Layerwise Estimations) is een framework dat de replicatie van experts optimaliseert binnen een vast geheugenbudget. Het doel is om de load-balans te maximaliseren zonder onnodig geheugen te verspillen.

Kernidee: In plaats van uniform te repliceren, analyseert CRAFT de load-verdeling per laag en allocateert replica's fijnkorrelig (per laag) op basis van het geschatte voordeel (benefit).

De workflow van CRAFT bestaat uit drie stappen:

1. Offline Profiling en Schatting:

   • CRAFT analyseert offline de expert-load distributie (op basis van een profiel van tokens).
   • Het schat voor elke laag het "replicatievoordeel" (balancedness gain) bij verschillende aantallen replica's. Hieruit blijkt dat sommige lagen (hoog-scheef) enorm profiteren van replicatie, terwijl andere lagen (laag-scheef) nauwelijks voordeel halen.

2. Optimalisatie (Dynamic Programming):

   • Het probleem wordt gemodelleerd als een Multiple-Choice Knapsack Problem (MCKP).
   • Doel: Kies het aantal replica's per laag zodat het totale replicatievoordeel wordt gemaximaliseerd, binnen het totale geheugenbudget (aantal replica's per GPU).
   • CRAFT gebruikt dynamische programmering om de optimale verdeling te vinden. Het kan ook automatisch het optimale replicatiefactor (R) bepalen dat de beste balans biedt tussen load-balans en geheugenverbruik.

3. Capaciteitsbewuste Toewijzing (Assignment & Placement):

   • Omdat het aantal replica's per laag nu verschilt, is het toewijzen aan GPU's complexer.
   • CRAFT gebruikt een greedy algoritme met twee doelen:
     • Primair: Zorg dat elke GPU ongeveer evenveel experts (capaciteit) heeft om geheugenfragmentatie te voorkomen.
     • Secundair: Verdeel experts over nodes in een "interleaved" patroon om node-level onbalans te minimaliseren.
   • Daarna wordt de standaard "greedy placement" toegepast om de zwaarst belaste experts op de minst belaste apparaten te plaatsen.

CRAFT kan naadloos worden geïntegreerd in bestaande frameworks (zoals SGLang) zonder extra training of modelwijzigingen.

Belangrijkste Bijdragen

1. Eerste gedetailleerde karakterisering: Het paper biedt het eerste inzicht in hoe doorvoer en load-balans schalen met het aantal replica's. Het toont aan dat er snel afnemende meeropbrengsten (diminishing returns) optreden en dat laag-specifiek gedrag cruciaal is.
2. Het CRAFT-framework: Een end-to-end oplossing die replica's toewijst op basis van geschatte prestatievoordelen. Het werkt voor elke MoE-architectuur en vereist geen modelwijzigingen.
3. Empirische validatie: Uitgebreide evaluatie op state-of-the-art LLM-serving frameworks met modellen van honderden miljarden tot een biljoen parameters.

Resultaten

De evaluatie is uitgevoerd op een cluster van 8 nodes (AWS p4de.24xlarge) met modellen zoals DeepSeek-R1 (671B) en Kimi-K2 (1000B).

• Doorvoerverbetering: CRAFT verhoogt de end-to-end doorvoer (goodput) gemiddeld met 1,14x (en tot 1,2x) ten opzichte van bestaande technieken zoals EPLB.
• Efficiëntie: CRAFT bereikt bijna dezelfde load-balans als EPLB, maar gebruikt aanzienlijk minder replica's (bijvoorbeeld 7,25x minder replica's voor DeepSeek-R1).
• Geheugenbesparing: Door minder replica's te gebruiken, blijft er meer geheugen over voor de KV-cache. Dit voorkomt dat de doorvoer daalt door een te kleine cache, wat bij EPLB vaak het geval is bij grote clusters.
• Robuustheid: CRAFT presteert consistent goed op zowel sterk scheve datasets (waar load-balans kritiek is) als minder scheve datasets. Bij EPLB leidt over-replicatie op minder scheve datasets juist tot een doorvoerverlies.
• Schaalbaarheid: CRAFT schaalt effectief naarmate het cluster groter wordt (van 6 naar 12 nodes), terwijl de prestaties van een "placement-only" baseline sterk verslechteren door load-imbalance.

Significantie

CRAFT lost een kritiek probleem op in de productie van grote MoE-modellen: de afweging tussen load-balans en geheugenverbruik. Bestaande methoden zijn vaak te "royaal" met geheugen, wat leidt tot inefficiëntie. CRAFT bewijst dat fijnkorrelige, kostenbewuste replicatie de sleutel is tot maximale doorvoer.

Het paper toont aan dat men niet hoeft te kiezen tussen load-balans en geheugenoptimalisatie; door slim te alloceren kan men beide maximaliseren. Dit maakt CRAFT een waardevolle toevoeging voor cloud-aanbieders en organisaties die grote taalmodellen draaien, omdat het direct leidt tot lagere operationele kosten en betere servicekwaliteit (lagere latency, hogere doorvoer) zonder de complexiteit van modelhertraining.