EasyBalance: Cross-Layer Load Balancing in Distributed MoE Inference

Stel je voor dat je een enorme, razendsnelle bibliotheek runt waar duizenden studenten (tokens) antwoorden moeten vinden uit een gigantische encyclopedie. Deze encyclopedie is niet door één persoon geschreven; het is een "Mixture of Experts" (MoE), wat betekent dat er honderden verschillende gespecialiseerde schrijvers (experts) binnenin zitten. Wanneer een student een vraag stelt, beslist een slimme bibliothecaris (de router) snel welke paar schrijvers het meest geschikt zijn om het antwoord te geven. Om dit super snel te laten gaan, splitst de bibliotheek deze schrijvers over vele verschillende computers (devices) die parallel werken.

Maar er is een addertje onder het gras: niet alle vragen zijn hetzelfde. Soms stelt een enorme menigte studenten vragen waarover slechts één specifieke schrijver het antwoord weet. De computer van die schrijver raakt dan overbelast en vertraagt, terwijl de andere computers, waarvan de schrijvers minder druk zijn, maar liefst zitten te wachten tot de trage partij klaar is. Dit "wachtspel" verspilt een enorme hoeveelheid energie en tijd. Jarenlang was de oplossing om meer schrijvers in te huren of hen te verplaatsen om de werkdruk te balanceren, maar dat kost te veel geheugen en is moeilijk uit te voeren tijdens het proces.

Maak kennis met EasyBalance, een slimme nieuwe strategie die dit wachtspel oplost zonder nieuwe mensen aan te nemen of meubels te verplaatsen. In plaats van te proberen de schrijvers te fixen, verandert EasyBalance wanneer de studenten hun vragen stellen. Het realiseert zich dat hoewel de bibliotheek vragen meestal laag voor laag verwerkt, het eigenlijk mogelijk is om studenten uit verschillende lagen van het proces tegelijkertijd te laten werken. Door deze groepen te mengen en te matchen, krijgt de drukke schrijver een pauze, omdat de "zware" vragen van de ene groep worden uitgebalanceerd door de "lichte" vragen van de andere groep. Het resultaat? De computers blijven druk bezig, de wachttijd verdwijnt en de hele bibliotheek draait veel sneller.

Het Probleem: De "Wacht op de Langzaamste" Regel

In de wereld van AI, specifiek bij deze "Mixture of Experts"-modellen, is het systeem ontworpen om ongelooflijk efficiënt te zijn. Er wordt slechts een klein aantal "experts" geactiveerd voor elk stukje data. Maar in een gedistribueerde opstelling — waarbij deze experts verspreid zijn over meerdere grafische kaarten (GPU's) — wordt het rommelig.

Stel je een estafette voor waarbij de wisselstok pas aan de volgende loper wordt overhandigd wanneer iedereen zijn etappe heeft voltooid. Als één loper een zware rugzak draagt (een "hete" expert met te veel tokens), vertraagt hij het hele team. De andere lopers, die hun lichte lasten al hebben afgelegd, moeten stilstaan en wachten. In de termen van het paper wordt dit load imbalance (belastingonbalans) genoemd. Het systeem wordt gegijzeld door de enkelvoudige langzaamste device, waardoor de anderen ongebruikt blijven.

Eerdere pogingen om dit op te lossen omvatten expert replicatie (extra kopieën van de drukke schrijvers inhuren) of expert migratie (schrijvers naar andere computers verplaatsen). Hoewel deze werken, hebben ze grote nadelen: ze vreten veel geheugen op, kosten extra tijd voor communicatie en zijn rigide. Als je het type vragen dat de bibliotheek krijgt verandert (de taak), kan het oude plan volledig falen.

De Oplossing: EasyBalance

De auteurs van dit paper, Yize Wu en collega's, stellen een frisse aanpak voor genaamd EasyBalance. Hun grote idee is om te stoppen met het proberen te fixen van de mapping (toewijzing) van schrijvers aan computers, en in plaats daarvan de planning van het werk te fixen.

Ze vertrouwen op twee belangrijke inzichten:

1. Cross-Layer Redundantie: Zelfs als een specifieke laag van het model een specifieke set experts heeft, zitten de experts in andere lagen al in het geheugen van de computer, klaar om te gaan. Ze zijn "natuurlijk redundant" voor de huidige taak. Je hoeft geen nieuwe kopieën in te huren; je moet alleen de exemplaren te gebruiken die je al hebt.
2. Workload Combinatie: Het paper suggereert dat je micro-batches (kleine groepen vragen) van verschillende lagen van het model tegelijkertijd kunt draaien. Hoewel het model zaken meestal één laag per keer verwerkt, laat de wiskunde zien dat het combineren van deze groepen veilig is. Sterker nog, het is vaak beter. Als Groep A een zware last heeft op Computer 1, maar Groep B heeft een zware last op Computer 2, dan balanceert het gezamenlijk draaien van beide groepen het totale gewicht. Het "worst-case" scenario (waarbij beide groepen op dezelfde computer terechtkomen) is statistisch gezien zeldzaam, vooral naarmate je meer computers toevoegt.

Hoe het werkt: De "Slimme Planner"

EasyBalance fungeert als een slimme verkeersregelaar. In plaats van alle studenten van Laag 1 te sturen, en daarna alle studenten van Laag 2, kijkt het naar de huidige menigte. Het kiest een mix van studenten uit verschillende lagen om samen te draaien.

• Het selecteert een deelverzameling van deze micro-batches om direct uit te voeren.
• Het stelt anderen uit (wacht op hen) als zij een knelpunt zouden veroorzaken.
• Dit gebeurt zonder te veranderen aan welke computer welke expert woont.

Dit betekent dat het systeem zich onmiddellijk kan aanpassen aan elk nieuw type taak zonder de hardware of het geheugen te hoeven herconfigureren. Het is als een restaurantkeuken die besluit om tegelijkertijd een burger en een salade te bereiden omdat de grill druk is maar het saladestation vrij is, in plaats van te wachten tot de grill alles heeft afgerond voordat de salade wordt gestart.

De Resultaten: Sneller en Minder Verspilling

De onderzoekers hebben EasyBalance getest op verschillende grote modellen, waaronder Qwen3-30B en Moonlight-16B, draaiend op 8 GPU's. Ze gebruikten een benchmark genaamd LongBench, die veel verschillende soorten taken dekt, zoals leesbegrip en codegeneratie.

De bevindingen waren consistent en indrukwekkend:

• Verminderde Stilstand: EasyBalance verminderde de GPU "onderbenutting" (de tijd dat computers niets doen) met grotende groepen van meer dan 40%. In veel gevallen daalde de stilstandtijd van ongeveer 0,35 (35% verspilling) naar ongeveer 0,2 (20% verspilling).
• Snellere Inferentie: Omdat de computers niet op elkaar wachtten, nam de totale tijd om een antwoord te krijgen (end-to-end latency) aanzienlijk af.
• Flexibiliteit: In tegen tegenstelling tot eerdere methoden die de specifieke taak vooraf moesten kennen om de experts in te stellen, werkte EasyBalance net zo goed over alle 13 geteste taken heen, van trivia tot codegeneratie.

Het paper onderzocht ook verschillende "scheduling" strategieën (hoe het systeem beslist welke groepen te mengen). Ze ontdekten dat een strategie genaamd MaxUtil (die probeert het gebruik van elke GPU te maximaliseren) het beste werkte, maar zelfs simpelere, snellere strategieën zoals CumUtil (batches één voor één toevoegen als ze helpen) waren nog steeds veel beter dan niets doen.

Waarom het ertoe doet

Het meest opwindende deel van EasyBalance is dat het geen extra geheugen of complexe herconfiguratie vereist. Het werkt met de bestaande opstelling. Naarmate AI-modellen groter en complexer worden, zal het probleem dat sommige computers wachten terwijl anderen werken alleen maar groter worden. Dit paper suggereert dat door simpelweg slimmer te zijn over wanneer we het werk uitvoeren, in plaats van waar we de experts plaatsen, we deze enorme AI-systemen aanzienlijk efficiënter kunnen maken.

De auteurs merken op dat hoewel hun methode zeer effectief is, deze steunt op de statistische waarschijnlijkheid dat zware belastingen niet altijd tegelijkertijd op dezelfde computer terechtkomen. In hun tests over diverse modellen en taken heen, versnelde deze strategie de inferentie consequent, wat bewijst dat de beste manier om een knelpunt op te lossen soms is om de werkers elkaar over de finishlijn te laten helpen, in plaats van te proberen de werkers zelf te verplaatsen.