veScale-FSDP: Flexible and High-Performance FSDP at Scale

Het paper introduceert veScale-FSDP, een herontworpen FSDP-systeem dat met zijn flexibele 'RaggedShard'-indeling en structure-bewuste planning native ondersteuning biedt voor geavanceerde trainingsmethoden, waardoor het zowel de doorvoer met 5-66% verhoogt als het geheugengebruik met 16-30% verlaagt ten opzichte van bestaande systemen.

De kern

Stel je voor dat je een gigantisch puzzelstuk hebt, zo groot als een stad. Je wilt dit stuk samen met duizenden vrienden oplossen. In de wereld van kunstmatige intelligentie (AI) is dit puzzelstuk een groot taalmodel (zoals de hersenen van een supercomputer), en de vrienden zijn de GPU's (de krachtige rekenchips).

De uitdaging? Hoe verdeel je die enorme puzzel eerlijk over duizenden mensen, zodat iedereen tegelijkertijd werkt zonder in de weg te lopen, en zonder dat iemand een berg papier (geheugen) nodig heeft die hij niet kan dragen?

Dit is precies wat veScale-FSDP oplost. Hier is hoe het werkt, vertaald naar alledaags taal:

1. Het oude probleem: De stijve doos

Stel je voor dat je de puzzelstukjes moet verdelen. De oude systemen (zoals FSDP1 en FSDP2) deden dit heel strikt: ze sneden de puzzel in exact gelijke vierkante blokjes.

• Het probleem: Sommige puzzelstukken zijn niet vierkant, maar langwerpig of hebben een speciale vorm (zoals een blokje dat 8x8 is in plaats van 1x1).
• De consequentie: Als je een langwerpig stukje in een vierkante doos probeert te stoppen, moet je er leeg papier bij plakken (zogenoemd padding) om het te laten passen. Of je moet het stukje kapot snijden, waardoor je het niet meer goed kunt gebruiken.
• Resultaat: De mensen (GPU's) staan te wachten op hun stukjes, of ze moeten veel tijd besteden aan het verplaatsen van stukjes heen en weer. Het is traag en verspillend.

2. De nieuwe oplossing: De flexibele doos (RaggedShard)

veScale-FSDP introduceert een nieuwe manier van verdelen, genaamd RaggedShard.

• De analogie: In plaats van te snijden in strakke vierkanten, laat je de mensen hun eigen vormige stukken pakken. Als iemand een lang stuk nodig heeft, krijgt hij een lang stuk. Als iemand een blokje nodig heeft, krijgt hij een blokje.
• Waarom is dit cool? Moderne AI-modellen gebruiken soms slimme trucs (zoals block-wise quantization of speciale rekenmethodes) die net die specifieke vormen nodig hebben. Met de oude systemen moest je je hele puzzel aanpassen aan de doos. Met veScale-FSDP past de doos zich aan aan de puzzel. Je hoeft je model niet meer "kapot te maken" om het te laten werken.

3. De logistiek: De slimme planner

Nu je flexibele stukjes hebt, is de volgende vraag: Hoe pak je ze in de vrachtwagens (de communicatie) zodat niemand wacht?

• Het probleem: Als je willekeurige stukjes in een vrachtwagen stopt, ontstaan er gaten. De vrachtwagen is niet vol, en je moet vaak stoppen om de lading te verschuiven.
• De oplossing: veScale-FSDP heeft een slimme planner (een algoritme). Deze planner kijkt naar alle puzzelstukjes en bedenkt de perfecte volgorde om ze in de vrachtwagens te laden.
• Het resultaat: Geen gaten, geen wachten, en de vrachtwagens zijn altijd vol. Dit zorgt ervoor dat de data razendsnel van de ene computer naar de andere gaat.

4. De magische koffer (Distributed Buffer)

Stel je voor dat elke persoon een eigen koffer heeft om zijn puzzelstukjes in te bewaren.

• Oude manier: Iedereen heeft een eigen kleine koffer. Als je een stukje wilt delen, moet je het uit je koffer halen, in een nieuwe doos doen, en het geven. Dat kost tijd en ruimte.
• veScale-FSDP: Ze gebruiken één grote, gedeelde koffer (de Distributed Buffer). Iedereen heeft direct toegang tot zijn stukje in die grote koffer zonder het fysiek te hoeven verplaatsen.
• Het voordeel: Het bespaart enorm veel ruimte (geheugen) en gaat veel sneller, omdat je niet hoeft te verpakken en verpakken.

Wat levert dit op?

Dankzij deze nieuwe aanpak heeft het team van ByteDance (de makers) fantastische resultaten geboekt:

• Snelheid: Het trainen van AI-modellen is 5% tot 66% sneller. Dat is alsof je een reis van 1 uur opeens in 40 minuten doet.
• Geheugen: Het systeem gebruikt 16% tot 30% minder geheugen. Dit betekent dat je met dezelfde hoeveelheid computers veel grotere modellen kunt bouwen, of dat je minder dure hardware hoeft te huren.
• Schaalbaarheid: Het werkt soepel, zelfs als je tienduizenden GPU's tegelijk gebruikt.

Conclusie

Kortom: veScale-FSDP is als het overstappen van een stijve, slecht ontworpen trein naar een hypermoderne, flexibele magneettrein. De trein past zich aan aan de passagiers (het model) in plaats van dat de passagiers zich moeten vervormen om in de trein te passen. Hierdoor reist iedereen sneller, comfortabeler en met minder brandstof (geheugen) nodig.

Dit maakt het mogelijk om in de toekomst nog grotere en slimmere AI's te bouwen, zonder dat we duizenden extra computers nodig hebben.

Technische samenvatting

1. Het Probleem

Fully Sharded Data Parallel (FSDP), ook bekend als ZeRO, is de standaard voor het trainen van grote taalmodellen (LLMs) omdat het flexibel is en weinig ingrijpende wijzigingen in modelcode vereist. Echter, bestaande FSDP-systemen (zoals DeepSpeed ZeRO, PyTorch FSDP1 en FSDP2) stuiten op twee fundamentele beperkingen die de schaalbaarheid en efficiëntie in de weg staan:

• Gebrek aan flexibiliteit voor structure-bewuste training: Moderne state-of-the-art modellen (zoals Gemini, Kimi K2, DeepSeek-V3) gebruiken geavanceerde technieken zoals matrix-gebaseerde optimalisatoren (bijv. Shampoo, Muon) en block-wise kwantisatie (bijv. 8-bit Adam). Deze technieken vereisen dat tensorblokken intact blijven als atomaire eenheden. Bestaande FSDP-implementaties shards (verdelen) parameters echter op element- of rij-niveau. Dit creëert grenzen die niet overeenkomen met de vereiste blokgrootte, wat leidt tot complexe randvoorwaarden, padding, en inefficiënte communicatie. Ontwikkelaars moeten hun modellen of optimalisatoren invasief aanpassen om dit te omzeilen.
• Inefficiëntie in communicatie en geheugengebruik: Bestaande systemen kampen met fragmentatie in geheugenbeheer en inefficiënte communicatiepatronen (zoals ongelijkmatige buffers en veelvuldige kopieeroperaties). Dit beperkt de schaalbaarheid tot enkele tienduizenden GPU's en zorgt voor een hoge geheugendruk, wat leidt tot "Out of Memory" (OOM) fouten of dure device-side vrijgaves.

2. Methodologie

De auteurs introduceren veScale-FSDP, een herontworpen FSDP-systeem dat flexibiliteit en prestaties combineert door drie kerncomponenten:

A. RaggedShard: Een flexibel sharding-formaat

In plaats van vaste element- of rij-sharding, introduceert veScale-FSDP RaggedShard.

• Dit formaat ondersteunt willekeurige sharding-granulariteit (bijv. blokken van 32x32 elementen) en willekeurige verdeling over apparaten.
• Het is gebaseerd op het concept van JaggedTensors (variabele lengte) en integreert naadloos met PyTorch's DTensor.
• Hierdoor kunnen blokken voor kwantisatie of matrix-optimalisatoren perfect uitgelijnd worden met de sharding-grenzen, zonder padding of complexe communicatie.

B. Structure-bewust planningsalgoritme

Om RaggedShard-tensors efficiënt te communiceren, moet men ze groeperen in communicatiebuffers. Een naïeve aanpak leidt tot inefficiëntie (niet-aaneengesloten geheugen, onevenwichtige belasting).

• De auteurs formuleren het groeperen van deze tensors als een NP-harde optimalisatieprobleem (gereduceerd tot het Partition-probleem).
• Ze ontwikkelen een heuristisch algoritme met dynamische programmering (DP) dat in polynoomtijd werkt. Dit algoritme herschikt de tensors en voegt padding toe tussen tensors in plaats van binnen tensors, zodat blokken aaneengesloten blijven en de belasting over GPU's in balans is.
• Dit minimaliseert padding-overhead en maximaliseert de netwerkbenutting.

C. Distributed Buffer (DBuffer)

Om de communicatie en geheugenefficiëntie verder te optimaliseren, introduceert het systeem DBuffer.

• Dit is een primitief dat een globale buffer biedt over een N-dimensionale apparaat-topologie.
• Zero-copy: DBuffer biedt directe toegang tot data zonder kopieeroperaties (copy-in/copy-out) die bij traditionele FSDP nodig zijn.
• Kernel-fusie: Het fuseert identieke operaties (zoals add, scale, zero) over meerdere tensors voordat communicatie plaatsvindt, wat de blokkeringstijd verlaagt.
• Geheugenbeheer: Het gebruikt batch-allocaties en deterministische stream-afhankelijkheden om geheugenfragmentatie te verminderen.

3. Belangrijkste Bijdragen

1. RaggedShard: Een nieuw DTensor-sharding-formaat dat structure-bewuste training (block-wise kwantisatie, niet-element-georiënteerde optimalisatoren) mogelijk maakt zonder invasieve code-wijzigingen.
2. Planningsalgoritme: Een praktische, snelle heuristiek voor het optimaliseren van communicatiebuffers voor onregelmatige tensorgroottes, wat leidt tot minimale padding.
3. DBuffer: Een hoogpresterende primitief voor zero-copy communicatie en geoptimaliseerd geheugenbeheer.
4. Schaalbaarheid: Het systeem is getest op productie-omgevingen met 10.000+ GPU's en ondersteunt modellen tot 2,4 biljoen parameters.

4. Resultaten

Uitgebreide evaluaties op modellen zoals LLaMA-3-70B, GPT-OSS-120B en interne MoE-modellen tonen de volgende verbeteringen ten opzichte van bestaande systemen (DeepSpeed ZeRO, PyTorch FSDP1/2, Megatron-FSDP):

• Doorvoer (Throughput): Een toename van 5% tot 66% (afhankelijk van het model en de schaal). Vooral bij MoE-modellen en het gebruik van geavanceerde optimalisatoren is de winst groot.
• Geheugengebruik: Een reductie in piek-geheugengebruik van 16% tot 30%. Dit maakt het mogelijk om grotere modellen te trainen op dezelfde hardware of meer modellen tegelijk te draaien.
• Schaalbaarheid: VeScale-FSDP toont lineaire schaalbaarheid tot 10.000 GPU's (zwakke schaalbaarheid) en tot 2,4T parameters.
• Ondersteuning voor geavanceerde technieken:
  • 8-bit Adam: Werkt naadloos met block-wise kwantisatie zonder extra communicatie-overhead.
  • Muon-optimalisator: Ondersteunt matrix-gebaseerde optimalisatoren (Newton-Schulz) die de volledige 2D-matrix vereisen, met een MFU (Model FLOPS Utilization) van 47,3% op 256 GPU's.
• Overhead: Het planningsalgoritme kost minder dan 0,3 seconden tijdens initialisatie, wat verwaarloosbaar is.

5. Betekenis en Impact

VeScale-FSDP is een doorbraak in het trainen van de volgende generatie AI-modellen:

• Ontkoppeling van Model en Systeem: Het systeem scheidt de modeldefinitie van de systeemoptimalisatie. Onderzoekers kunnen complexe architecturen (zoals MoE) en nieuwe optimalisatoren ontwerpen zonder zich zorgen te hoeven maken over de onderliggende parallelisatie-logica.
• Productie-klaar: Het systeem is al "battle-tested" in productie bij ByteDance en is portabel zonder afhankelijkheid van interne infrastructuur.
• Open Source: De code voor RaggedShard is open-source gemaakt (via veScale), wat bijdraagt aan het PyTorch-ecosysteem. De functionaliteit wordt zelfs overwogen als een officiële feature in de PyTorch-roadmap.
• Toekomstbestendig: Het lost de fundamentele beperkingen op die de schaalbaarheid van huidige FSDP-implementaties beperken, waardoor het mogelijk wordt om modellen met biljoenen parameters efficiënt te trainen op clusters van tienduizenden GPU's.

Kortom, veScale-FSDP combineert de flexibiliteit die nodig is voor innovatieve modelarchitecturen met de strikte prestatie-eisen die nodig zijn voor training op industriële schaal.