Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

Dit paper introduceert Deep Optimizer States, een nieuwe techniek die de GPU-geheugenuitbuiting optimaliseert door de optimizer-parameters dynamisch tussen CPU en GPU te wisselen tijdens het trainen van grote taalmodellen, wat resulteert in 2,5 keer snellere iteraties dan bestaande methoden.

De kern

Stel je voor dat je een gigantische puzzel probeert op te lossen. Deze puzzel is zo groot dat hij de hele vloer van je huis beslaat. In de wereld van kunstmatige intelligentie (AI) is deze puzzel een groot taalmodel (zoals een slimme chatbot), en de puzzelstukjes zijn de parameters (de kennis) die het model moet leren.

Het probleem? Je hebt niet genoeg ruimte op je werktafel (de GPU, de krachtige computerchip die de zware rekenwerk doet) om alle puzzelstukjes tegelijkertijd neer te leggen. Je moet dus steeds stukjes van de tafel halen en in een kast (de CPU-geheugen of host-geheugen) zetten, om er later weer bij te kunnen.

Dit is wat er gebeurt bij het trainen van moderne AI-modellen: ze zijn zo groot dat ze niet meer in het snelle geheugen van de grafische kaart passen.

Het Oude Probleem: De Trage Koerier

In de traditionele manier (zoals gebruikt door DeepSpeed, een populaire software), gebeurt het volgende:

1. De GPU doet het rekenwerk (het "vooruit" en "achteruit" kijken in de puzzel).
2. Als het tijd is om de puzzelstukjes aan te passen (de optimizer update), moet de GPU wachten tot de CPU de stukjes uit de kast haalt, aanpast, en weer terugbrengt.
3. Het probleem: De CPU is als een trage koerier die op een fiets rijdt, terwijl de GPU een supersnelle raceauto is. De raceauto staat urenlang stil te wachten tot de fietsman de pakketjes heeft bezorgd. Bovendien is de deur tussen de werktafel en de kast (de PCIe-verbinding) vaak te smal, waardoor er veel wachttijd ontstaat.

De Oplossing: "Deep Optimizer States" (De Slimme Uitwisseling)

De auteurs van dit paper hebben een nieuwe methode bedacht, genaamd Deep Optimizer States. Ze gebruiken een slimme truc om die wachttijden te elimineren.

Stel je voor dat je niet wacht tot de fietsman alles heeft bezorgd, maar dat je een concurrentie organiseert:

1. Verdeel en heers: In plaats van de hele puzzel als één blok te behandelen, delen ze de puzzel op in kleine groepjes (subgroepen).
2. Tandem-werk:
   • Terwijl de trage fietsman (CPU) bezig is met het aanpassen van groepje A en B, haalt de raceauto (GPU) alvast groepje C uit de kast en past die direct aan.
   • Terwijl de raceauto bezig is met groepje C, haalt de fietsman alvast groepje D en E klaar.
3. De "Interleaved" (Verweven) Techniek: Ze laten de CPU en de GPU tegelijkertijd werken aan verschillende onderdelen van de puzzel. Ze wisselen elkaar af op een manier die precies past bij de snelheid van de deur (de bandbreedte).

De Creatieve Analogie: De Keuken en de Magazijn

Laten we dit vergelijken met een drukke keuken:

• De Chef-kok (GPU): Heeft een super-snelle snijtafel en kan duizenden groenten per minuut snijden. Maar zijn tafel is klein.
• Het Magazijn (CPU/Host): Is enorm groot, maar ligt ver weg en de deur is smal.
• De Oude Methode: De chef snijdt alles, stopt de groenten in een mand, loopt naar het magazijn, wacht tot de magazijnmedewerker (CPU) de groenten snijdt, en wacht dan tot die terugkomt. De chef staat de hele tijd stil.
• De Nieuwe Methode (Deep Optimizer States):
  • De chef snijdt een hoop groenten (groepje 1).
  • Terwijl de chef doorgaat met groenten (groepje 2), rent de magazijnmedewerker alvast naar de deur om groenten (groepje 3) op te halen en snijdt die alvast.
  • De chef pakt de klaarliggende groenten van groepje 3 en snijdt die, terwijl de magazijnmedewerker alvast groepje 4 haalt.
  • Resultaat: Niemand staat stil. De deur wordt constant gebruikt, en zowel de chef als de magazijnmedewerker werken hard.

Waarom werkt dit zo goed?

De onderzoekers hebben ontdekt dat tijdens het trainen van AI er momenten zijn waarop de GPU-geheugen ruimte vrijkomt (bijvoorbeeld als de "activaties" van een berekening klaar zijn). Ze gebruiken die lege ruimte om tijdelijk de "slimme" GPU te laten werken aan een deel van de aanpassingen, terwijl de CPU het andere deel doet.

Ze hebben ook een slim rekenmodel (een soort verkeersregelaar) bedacht dat precies berekent: "Hoeveel groepjes moet de GPU doen en hoeveel de CPU, zodat niemand wacht?"

Het Resultaat

Door deze slimme dans tussen de snelle GPU en de trage CPU:

• Wordt de training 2,5 keer sneller.
• Kunnen mensen met minder dure hardware (zoals een enkele server met 4 GPU's) toch enorme AI-modellen trainen.
• Wordt de "dode tijd" van de computer bijna volledig opgeheven.

Kortom: In plaats van wachten tot de trage koerier alles bezorgt, laten ze de snelle raceauto en de trage fietsman samenwerken aan verschillende pakketjes tegelijk. Zo wordt de hele bezorging veel sneller, zonder dat je een nieuwe, duurdere auto hoeft te kopen.

Technische samenvatting

Probleemstelling

De training van grote taalmodellen (LLMs) en Transformers wordt beperkt door een "geheugengrens" (memory wall). Hoewel technieken zoals 3D-parallelisme (data, pipeline, tensor) helpen, zijn de optimizer-staten (momentum, variantie en parameters in FP32) vaak te groot om volledig in het GPU-geheugen te passen.

• Huidige oplossing: State-of-the-art frameworks (zoals DeepSpeed ZeRO-Offload) lossen dit op door de optimizer-staten naar het host-geheugen (CPU) te verplaatsen (offloading).
• De beperkingen:
  1. I/O-bottleneck: De bandbreedte tussen CPU en GPU (via PCIe) is beperkt (vaak 25-50 GB/s), wat veel trager is dan GPU-communicatie.
  2. Rekenkracht: CPUs zijn orders of magnitude trager dan GPUs bij het updaten van parameters.
  3. Suboptimaal gebruik: Bestaande hybride methoden (zoals DeepSpeed TwinFlow) gebruiken een statische verdeling: een vast percentage van de optimizer staat op de GPU, de rest op de CPU. Dit leidt tot onderbenutting van de GPU-geheugencapaciteit tijdens de forward/backward passes en zorgt voor blokkerende I/O-wachttijden, omdat de CPU en GPU vaak niet gelijktijdig werken.

Methodologie: Deep Optimizer States

De auteurs introduceren Deep Optimizer States, een middleware-techniek die dynamische, verweven (interleaved) offloading toepast om de training te versnellen. De kern van de methode is het benutten van fluctuaties in het GPU-geheugengebruik en de PCIe-bandbreedte.

Belangrijkste ontwerpprincipes:

1. Verweven Updaten (Interleaved Updates): In plaats van statisch vast te leggen welke optimizer-subgroepen op de CPU of GPU zitten, worden updates dynamisch verdeeld. Tijdens de update-fase worden sommige subgroepen bijgewerkt op de GPU en andere op de CPU, en dit wordt afwisselend gedaan.
2. Overlappende Beweging en Executie: Het systeem overlapt drie processen om idle-tijd te minimaliseren:
   • CPU-berekeningen voor het updaten van subgroep $i$.
   • Asynchrone H2D-overdracht (Host-naar-Device) van de optimizer-staat voor subgroep $i+1$ naar de GPU.
   • Asynchrone D2H-overdracht (Device-naar-Host) van de reeds bijgewerkte parameters van subgroep $i-1$ van de GPU.
3. Efficiënt Gradient Management: Gradients die op de GPU worden gegenereerd (vaak in FP16) worden direct op de GPU omgezet naar FP32 (in-place conversie) voordat ze naar de host worden gestuurd. Dit voorkomt trage conversie op de CPU en benut de hoge bandbreedte van de GPU voor conversie.
4. Performance Model: De auteurs hebben een wiskundig model ontwikkeld om de optimale verhouding ($k$) te bepalen tussen CPU- en GPU-updates. Dit model balanceert de doorvoersnelheid van de CPU, de GPU en de PCIe-bandbreedte om de maximale overlapping te bereiken.
   • Formule: $k$ wordt berekend op basis van de tijd die nodig is voor CPU-updates, downscaling en data-overdracht, zodat de PCIe-link continu bezet is zonder dat de CPU of GPU wacht.

Implementatie:
De techniek is geïntegreerd in DeepSpeed (specifiek de ZeRO-3 engine) en Megatron-LM. Het gebruikt CUDA-streams voor asynchrone data-overdracht en vermijdt blokkering van de Python Global Interpreter Lock (GIL) door C++-extensies.

Kernbijdragen

1. Observatie: Het analyseren van trainingstijden toont aan dat GPU-geheugenuitbesteding tijdens de update-fase sterk fluctueert en PCIe-links vaak onderbenut zijn.
2. Nieuwe Architectuur: Een ontwerp dat statische offloading vervangt door dynamische, verweven offloading van optimizer-subgroepen.
3. Performance Model: Een algoritme dat de optimale "update stride" (hoe vaak een update op de GPU moet gebeuren) berekent op basis van systeemspecificaties.
4. Middleware: Een open-source implementatie die naadloos werkt met bestaande LLM-training frameworks.
5. Validatie: Uitgebreide experimenten die aantonen dat de methode significant sneller is dan bestaande state-of-the-art oplossingen.

Resultaten

De auteurs hebben hun aanpak getest op een systeem met 4x NVIDIA H100 GPUs (80 GB) en Intel Xeon CPU's, trainend op modellen tot 20 miljard parameters.

• Snelheidswinst: Deep Optimizer States is 2,5x sneller in iteratietijd vergeleken met DeepSpeed ZeRO-3 (volledige CPU-offloading) en 1,7x tot 2,3x sneller dan DeepSpeed TwinFlow (statische hybride offloading).
• Update-Doorvoer: De optimizer-update-doorvoer is gemiddeld 70% hoger dan bij ZeRO-3.
• Efficiëntie: De methode bereikt een GPU-gebruik van bijna 100% en PCIe-gebruik van ~40% van de piekbandbreedte, terwijl CPU-gebruik rond de 60-70% blijft (door concurrentie op het geheugen, maar toch veel efficiënter dan statische methoden).
• Schalbaarheid: De prestaties schalen goed met toenemende data-parallelisme en verschillende micro-batch maten. Zelfs bij een 20B-parameter model is de totale trainingstijd vergelijkbaar met die van een 7B-model met traditionele methoden.
• Correctheid: Het performance model bleek accuraat; de optimale verhouding bleek vaak $k=2$ (elke tweede subgroep op de GPU), wat leidde tot de hoogste doorvoer.

Betekenis en Toekomstperspectief

De paper is significant omdat het een fundamenteel probleem in het trainen van grote modellen oplosbaar maakt zonder de noodzaak van extreem dure hardware-upgrades (meer GPU-geheugen).

• Resource-efficiëntie: Het stelt onderzoekers in staat om LLM's te trainen op systemen met beperkt GPU-geheugen (bijv. single-node setups) door het CPU-geheugen effectiever te benutten zonder de snelheid te verliezen.
• Toekomstige hardware: De methode is nog relevanter voor toekomstige systemen zoals de NVIDIA Grace Hopper Superchip, die een zeer hoge bandbreedte (200 GB/s) tussen CPU en GPU hebben. Dit zou de prestaties van deze dynamische offloading nog verder kunnen verbeteren.
• Generalisatie: Hoewel getest op Transformers, is de aanpak generiek en toepasbaar op andere grote modellen (zoals vision-modellen of domeinspecifieke modellen).

Kortom, Deep Optimizer States transformeert de "memory wall" van een blokkade in een dynamisch, geoptimaliseerd proces door de idle-tijd van CPU's en GPU's slim te synchroniseren en de beperkte I/O-bandbreedte maximaal te benutten.