FAST: An Efficient Scheduler for All-to-All GPU Communication

Dit paper introduceert FAST, een efficiënte scheduler voor All-to-All(v)-communicatie die schaalbaarheid en snelheid verbetert door workload-scheefheid en congestie aan te pakken, waardoor synthetisatietijden drastisch worden verkort en prestaties op diverse GPU-clusters worden geoptimaliseerd.

De kern

Stel je voor dat je een gigantische, super-snelle keuken hebt met honderden chefs (de GPU's) die samen een enorme maaltijd bereiden. Dit is wat er gebeurt in moderne kunstmatige intelligentie (AI) systemen, zoals die gebruikt worden voor chatbots of medische diagnoses.

Om deze maaltijd te maken, moeten de chefs constant informatie uitwisselen. Soms moet iedere chef met elke andere chef praten. In de tech-wereld noemen ze dit "All-to-All" communicatie.

Het Probleem: De Chaos in de Keuken

In het verleden was dit relatief makkelijk. Maar nu zijn de keukens veranderd:

1. Twee soorten gangen: Binnen één keuken (een server) lopen de chefs via supersnelle, korte gangen (de "scale-up" links). Maar om met chefs in andere keukens te praten, moeten ze door een trage, lange tunnel (de "scale-out" links).
2. Oneerlijke werkverdeling: Niet elke chef heeft evenveel werk. Sommige chefs krijgen plotseling 100 orders, terwijl anderen maar 1 order hebben. Dit heet "skew" (scheefheid).
3. De "Incast" crash: Stel je voor dat 50 chefs tegelijk proberen hun bord door de smalle tunnel naar één specifieke chef te duwen. De tunnel raakt verstopt, de borden vallen, en iedereen staat vast. Dit noemen ze "incast".
4. Snel veranderend werk: Het werk verandert elke paar honderd milliseconden. Een chef die nu overbelast is, is straks misschien helemaal leeg.

Huidige oplossingen zijn als een manager die een week van tevoren een perfect rooster maakt. Maar als het werk elke seconde verandert, is dat rooster binnen een seconde al verouderd. Het kost de manager uren om een nieuw rooster te maken, terwijl de chefs al urenlang inactief staan.

De Oplossing: FAST (De Slimme Chef)

De onderzoekers hebben FAST bedacht. Dit is een slimme, snelle planner die niet probeert het perfecte rooster voor de hele eeuw te maken, maar gewoon nu de beste beslissingen neemt.

FAST werkt in drie simpele stappen, met een paar slimme trucs:

Stap 1: De "Interne Ruil" (Het binnen de keuken oplossen)

Stel je voor dat Chef A 100 borden moet afleveren bij Chef B, maar Chef C heeft maar 2 borden. Chef A zou de hele tunnel verstoppen.
FAST's truc: Chef A geeft een deel van die 100 borden aan Chef C. Omdat Chef A en C in dezelfde keuken zitten, kunnen ze die borden via de supersnelle gangen uitwisselen.

• Het resultaat: Nu moet Chef A nog maar 50 borden door de trage tunnel sturen, en Chef C ook 50. De trage tunnel wordt niet meer overbelast door één persoon. De "scheefheid" is binnen de keuken opgelost.

Stap 2: De "Eén-op-Één Dans" (Het door de tunnel sturen)

Nu de werkverdeling binnen de keukens gelijk is, moeten de chefs door de trage tunnel.
FAST's truc: In plaats van dat iedereen tegelijk probeert te praten (wat leidt tot de crash), organiseert FAST een perfecte dans.

• Chef A praat alleen met Chef X.
• Chef B praat alleen met Chef Y.
• Niemand praat met dezelfde persoon als iemand anders.
  Dit zorgt ervoor dat er nooit een file ontstaat in de tunnel. Iedereen heeft precies één gesprekspartner, en iedereen is even druk. Dit heet een "one-to-one matching".

Stap 3: De "Glijbaan" (Pipelining)

Meestal wacht je tot alles klaar is voordat je begint met de volgende stap. FAST doet dit niet.
Terwijl de chefs in de tunnel (scale-out) al aan het werk zijn, beginnen de chefs binnen de keuken (scale-up) al met de volgende ronde van uitwisseling. Het is alsof de chefs een glijbaan gebruiken: zodra je de top hebt verlaten, glijd je alweer naar beneden, terwijl er iemand nieuw de top op komt. Hierdoor staat de trage tunnel nooit stil.

Waarom is dit zo geweldig?

1. Snelheid: De huidige planners zijn als een wiskundige die uren zit te rekenen om een rooster te maken. FAST is als een ervaren chef die in microseconden (miljoenste van een seconde) een beslissing neemt. Het is snel genoeg om mee te gaan met de snel veranderende AI-werkstromen.
2. Efficiëntie: Door de interne gangen slim te gebruiken, wordt de trage tunnel nooit verstopt. De chefs blijven de hele tijd werken in plaats van te wachten.
3. Resultaat: In tests met de nieuwste hardware (NVIDIA en AMD) was FAST tot 4,5 keer sneller dan de beste bestaande systemen bij zware, oneerlijke werklasten.

Samenvattend

FAST is als een super-efficiënte conciërge in een drukke kantoorgebouw.

• Als iedereen tegelijk naar de lift wil (de trage tunnel), zorgt de conciërge ervoor dat mensen eerst even in de gangen (de snelle interne links) wisselen zodat de lift niet overvol raakt.
• Hij zorgt ervoor dat mensen in de lift altijd één-op-één zitten met iemand die ze moeten spreken, zodat er geen ruzie of files zijn.
• En hij doet dit zo snel, dat niemand in de lift hoeft te wachten.

Hierdoor kunnen de AI-modellen veel sneller leren en werken, zonder vast te lopen in de digitale files.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "FAST: An Efficient Scheduler for All-to-All GPU Communication" in het Nederlands.

1. Het Probleem

All-to-All (en specifiek alltoallv, waarbij verschillende hoeveelheden data worden uitgewisseld) is een kritieke communicatieprimitief in moderne machine learning-werklasten, met name bij Mixture-of-Experts (MoE) modellen. In deze modellen worden tokens dynamisch naar verschillende "experts" (neuronale netwerken) gerouteerd, wat frequente all-to-all communicatie tussen GPU's vereist.

De huidige uitdagingen voor het plannen van deze communicatie zijn:

• Werklast-scheefheid (Skewness): Niet alle GPU-paren wisselen evenveel data uit. Sommige experts worden veel vaker gekozen dan anderen, wat leidt tot "stragglers" (traagste GPU's die de hele groep vertragen) en overbelaste netwerkkaarten (NIC's).
• Dynamiek: Het verkeerspatroon verandert elke paar honderd milliseconden tijdens het trainen, omdat de gating-functie tokens dynamisch toewijst. Statische schema's werken hier niet.
• Heterogene Fabric: Moderne clusters hebben een twee-laags netwerk:
  • Scale-up: Snelle intra-server links (bijv. NVLink, Infinity Fabric).
  • Scale-out: Langzamere inter-server links (bijv. Ethernet, InfiniBand).
    De schaalbaarheid wordt beperkt door de langzamere scale-out laag.
• Incast-congestie: Bij dichte communicatiepatronen sturen veel zenders tegelijkertijd data naar dezelfde ontvanger, wat leidt tot wachtrijen in switches en verlies aan doorvoer.
• Schaalbaarheid van bestaande oplossingen: Bestaande planners (zoals TACCL, TE-CCL) gebruiken complexe optimalisatieproblemen (vaak NP-hard) die minuten tot uren nodig hebben om een schema te genereren. Dit is te traag voor dynamische MoE-werklasten.

2. Methodologie: De FAST-benadering

FAST (Fast All-to-All Scheduler) lost deze problemen op door het probleem te vereenvoudigen in plaats van de complexiteit van de planner te vergroten. De kerninzicht is dat de scale-up (intra-server) bandbreedte ongeveer een orde van grootte hoger is dan scale-out. Daarom kan de snelle scale-up laag worden gebruikt om scheefheid lokaal te absorberen voordat data de langzamere scale-out laag bereikt.

Het ontwerp bestaat uit twee fasen:

Fase 1: Intra-server Scheduling (Balancering en Herdistributie)

Het doel is om de werklast binnen elke server gelijk te verdelen voordat data de server verlaat.

• Sender-balancering: Zware GPU's sturen overtollige data naar lichtere GPU's binnen dezelfde server via de snelle scale-up links. Hierdoor heeft elke NIC in de server dezelfde uitgaande last naar een specifieke doelserver.
• Receiver-balancering: In plaats van dat een GPU direct data ontvangt van de bron-GPU, ontvangt een "proxy"-GPU in de doelserver de data van de bron-GPU. Deze proxy deelt de data vervolgens via scale-up links uit naar de juiste doel-GPU's.
• Resultaat: De scheefheid binnen de server wordt geëlimineerd. De scale-out laag ziet nu een vereenvoudigde, gebalanceerde werklast op server-niveau.

Fase 2: Inter-server Scheduling (Gebalanceerde 1-op-1 Transfers)

Zodra de intra-server scheefheid is opgelost, wordt de resterende server-tot-server werklast gescheduld.

• Birkhoff-decompositie: FAST past een wiskundige stelling toe (Birkhoff's decomposition) om het verkeersmatrix op te splitsen in een reeks van permutatiematrices.
• Eén-op-één matching: Elke permutatiematrix vertegenwoordigt een transfer-fase waarin elke zender precies één ontvanger heeft en vice versa. Dit garandeert dat er geen incast optreedt (geen overbelasting van een enkele ontvanger).
• Optimaliteit: Deze methode zorgt ervoor dat de "bottleneck"-servers (die de meeste data moeten sturen/ontvangen) continu actief blijven op hun maximale snelheid totdat hun werk voltooid is. Dit benadert de theoretische ondergrens voor voltooiingstijd.

Pipelining

Om de totale doorvoertijd te minimaliseren, worden de fasen gepipelined:

• Terwijl de scale-out transfers (fase 2) plaatsvinden, worden de intra-server balancering en herdistributie (fase 1) voor de volgende fase in de achtergrond uitgevoerd.
• Dit verbergt de overhead van de snelle intra-server operaties achter de langzamere scale-out operaties.

3. Belangrijkste Bijdragen

1. Polynoom-tijd Planner: FAST is de eerste planner die alltoallv-werklasten in polynoomtijd ($O(N^5)$) kan schedulen, waardoor het geschikt is voor dynamische werklasten die elke paar honderd milliseconden veranderen.
2. Architectuur-afhankelijk Ontwerp: Het benut specifiek de heterogeniteit van moderne GPU-clusters (snelle scale-up vs. trage scale-out) om het probleem te decomponeren.
3. Geen Incast: Door het afdwingen van 1-op-1 matchingen via Birkhoff-decompositie wordt netwerkcongestie proactief voorkomen.
4. Implementatie: De auteurs hebben FAST geïmplementeerd en geëvalueerd op zowel NVIDIA H200 als AMD MI300X clusters.

4. Resultaten

De evaluatie toont aanzienlijke verbeteringen ten opzichte van state-of-the-art oplossingen (NCCL, RCCL, DeepEP, TACCL, TE-CCL):

• Doorvoersnelheid (Throughput):
  • Op NVIDIA clusters: 1.01x tot 1.3x sneller dan de beste baselines (NCCL/DeepEP) bij scheve werklasten.
  • Op AMD clusters: 1.5x tot 2.8x sneller dan de beste baselines.
  • In geïntegreerde MoE-training (Megatron-LM op AMD): Een 4.48x verbetering in end-to-end training throughput vergeleken met RCCL, die zwaar lijdt onder incast.
• Scheduling Tijd:
  • FAST is extreem snel: het genereert een schema voor 64 GPU's in slechts 221 µs.
  • Ter vergelijking: Bestaande solver-gebaseerde planners (zoals SyCCL) hebben seconden tot minuten nodig voor veel kleinere clusters, wat ze onpraktisch maakt voor dynamische MoE-werklasten.
• Schaalbaarheid: De planner schaalbaar tot 320 GPU's met een overhead van slechts 77 ms.

5. Betekenis en Impact

FAST lost een fundamenteel knelpunt op in het trainen van grote MoE-modellen. Door de planningstijd met ordes van grootte te verminderen en de netwerkcongestie te elimineren, maakt het efficiënter gebruik van dure GPU-hardware mogelijk.

• Praktische toepasbaarheid: Het is snel genoeg om "on-the-fly" te worden gebruikt tijdens het trainen, waarbij het verkeerspatroon continu verandert.
• Toekomstbestendig: Het ontwerp is robuust voor de groeiende kloof tussen intra- en inter-server bandbreedte, wat waarschijnlijk zal aanhouden in toekomstige hardware-generaties.
• Open Source: De code is beschikbaar gesteld, wat de adoptie in bestaande frameworks (zoals Megatron-LM) faciliteert.

Samenvattend introduceert FAST een nieuw paradigma voor collectieve communicatie: in plaats van te proberen een perfect, universeel schema te berekenen (wat te duur is), wordt het probleem lokaal opgelost met snelle links en globaal geoptimaliseerd met een efficiënt wiskundig model voor de langzame links.