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

Das Paper stellt „Deep Optimizer States" vor, eine neue Technik, die durch dynamisches Verschieben von Optimierer-Zuständen zwischen CPU und GPU auf Basis von Nutzungsschwankungen die Speichereffizienz beim Training großer Transformer-Modelle verbessert und so im Vergleich zu bestehenden Methoden eine 2,5-fache Beschleunigung der Iterationen ermöglicht.

Das Wesentliche

🚀 Das Problem: Der riesige Rucksack und die schmale Brücke

Stell dir vor, du möchtest einen riesigen, intelligenten Roboter (ein KI-Modell wie ein Chatbot) trainieren. Dieser Roboter hat Milliarden von „Gedanken" (Parameter), die er lernen muss.

• Der Rucksack (GPU-Speicher): Um den Roboter zu trainieren, brauchst du einen sehr schnellen Arbeitsbereich, den wir „GPU-Speicher" nennen. Das ist wie ein kleiner, aber superschneller Schreibtisch.
• Der Keller (CPU-Speicher): Wenn der Schreibtisch zu voll ist, musst du Dinge in den Keller (den normalen Arbeitsspeicher des Computers) stellen. Das ist viel größer, aber viel langsamer zu erreichen.

Das Dilemma:
Bei modernen KI-Modellen ist der Rucksack (der GPU-Speicher) einfach zu klein. Der Roboter braucht nicht nur Platz für seine aktuellen Gedanken, sondern auch für einen riesigen „Lern-Protokoll" (den Optimierer-Zustand), der ihm sagt, wie er seine Fehler korrigieren soll. Dieser Protokoll ist oft dreimal so groß wie der Roboter selbst!

Da der Rucksack nicht reicht, müssen wir den Protokoll in den Keller legen. Aber hier liegt das Problem:

1. Der Flur ist eng: Der Weg zwischen Schreibtisch und Keller ist ein schmaler Flur (die PCIe-Leitung). Er ist überlastet.
2. Der Helfer ist langsam: Der Mann im Keller (die CPU) ist zwar fleißig, aber viel langsamer als der Assistent auf dem Schreibtisch (die GPU).

Das Ergebnis: Der Assistent auf dem Schreibtisch steht oft nur herum und wartet, bis der Mann im Keller die neuen Anweisungen durch den engen Flur geschleppt hat. Das Training wird extrem langsam.

---

💡 Die Lösung: „Deep Optimizer States" – Ein cleverer Takt

Die Autoren dieses Papiers haben eine neue Methode namens Deep Optimizer States entwickelt. Sie nennen es „interleaved offloading" (verflochtene Auslagerung), aber wir nennen es einfach: Der Taktwechsel.

Stell dir vor, du hast eine Gruppe von Arbeitern, die einen riesigen Haufen Steine (die Daten) bewegen müssen.

Wie es bisher funktioniert (Die alte Methode):

Die alten Systeme (wie DeepSpeed TwinFlow) haben gesagt: „Okay, wir legen die ersten 20 % der Steine auf den schnellen Schreibtisch und den Rest in den Keller."

• Das Problem: Während die Arbeiter auf dem Schreibtisch die 20 % bearbeiten, warten sie oft auf den Keller. Und während die im Keller arbeiten, warten sie auf den Schreibtisch. Es gibt viel Leerlauf. Der Flur wird nicht voll ausgenutzt.

Wie es jetzt funktioniert (Die neue Methode):

Die Autoren sagen: „Nein! Wir machen es wie einen gut getakteten Tanz."

1. Zerlegung in kleine Gruppen: Sie teilen den riesigen Lern-Protokoll in viele kleine, unabhängige Gruppen (Subgruppen) auf.
2. Der Tanz (Interleaving):
   • Während der Keller-Mann gerade an Gruppe A und B arbeitet, holt sich der Schreibtisch-Assistent bereits Gruppe C aus dem Keller (während der Keller-Mann noch arbeitet!).
   • Sobald der Keller-Mann fertig ist, schickt er die Ergebnisse durch den Flur.
   • Gleichzeitig bearbeitet der Schreibtisch-Assistent Gruppe C, während der Keller-Mann schon mit Gruppe D beginnt.
   • Der Clou: Der Flur (PCIe) wird in beide Richtungen gleichzeitig genutzt. Der Schreibtisch arbeitet, während der Keller arbeitet. Niemand wartet auf den anderen.

Die Analogie:
Stell dir eine Waschmaschine und einen Trockner vor.

• Alt: Du wäschst eine Ladung, wartest, bis sie fertig ist, lädst sie in den Trockner, wartest, bis der Trockner fertig ist, und fängst erst dann mit der nächsten Ladung an.
• Neu (Deep Optimizer States): Du legst die erste Ladung in die Waschmaschine. Während sie wäscht, legst du die zweite Ladung in den Trockner (oder holst die erste aus der Maschine, während die zweite gewaschen wird). Alles läuft parallel. Die Maschinen stehen nie still.

---

🛠️ Was macht das System noch smarter?

Das Papier beschreibt noch zwei weitere Tricks, die wie ein Schweizer Taschenmesser funktionieren:

1. Der perfekte Rhythmus (Performance-Modell):
   Das System berechnet genau, wie viele Gruppen auf dem Schreibtisch und wie viele im Keller bearbeitet werden sollen. Es fragt sich: „Wie schnell ist mein Keller-Mann? Wie schnell ist mein Flur?" Basierend darauf stellt es den Takt so ein, dass keine Maschine jemals leerläuft. Es ist wie ein Dirigent, der das Orchester so leitet, dass keine Note verpasst wird.

2. Der schnelle Umzug (Präzision):
   Normalerweise müssen Daten beim Umzug vom Keller zum Schreibtisch ihre „Kleidung" wechseln (von 16-Bit auf 32-Bit Precision), was Zeit kostet. Das neue System macht diesen Wechsel direkt auf dem Schreibtisch, während die Daten noch unterwegs sind. Das spart Zeit und verhindert, dass der Flur durch das Umkleiden blockiert wird.

---

🏆 Das Ergebnis: Warum ist das wichtig?

Durch diesen „Taktwechsel" erreichen die Autoren beeindruckende Ergebnisse:

• 2,5-mal schneller: Das Training läuft bis zu 2,5-mal schneller als bei den besten bisherigen Methoden.
• Effizientere Nutzung: Der teure GPU-Speicher und die schnellen Prozessoren werden zu 100 % ausgelastet, statt zu warten.
• Zukunftssicher: Auch wenn die Modelle noch größer werden (Milliarden von Parametern), bleibt das Training machbar, ohne dass man riesige, teure Supercomputer braucht.

Zusammenfassung in einem Satz:

Deep Optimizer States ist wie ein genialer Choreograf, der dafür sorgt, dass der schnelle Computer (GPU) und der große Speicher (CPU) nicht mehr aufeinander warten, sondern in einem perfekten, durchgehenden Tanz zusammenarbeiten, um KI-Modelle viel schneller zu lernen.

Technische Zusammenfassung

1. Problemstellung

Das Training von großen Sprachmodellen (LLMs) und Transformer-Architekturen stößt zunehmend auf eine „Speicherwand" (Memory Wall). Die Modelle wachsen auf Hunderte von Milliarden Parametern, was den Speicherbedarf für Modellparameter, Gradienten, Aktivierungen und vor allem den Optimierer-Zustand (Optimizer State) übersteigt, selbst wenn mehrere GPUs kombiniert werden.

Um dies zu lösen, offladen State-of-the-Art-Ansätze (wie DeepSpeed ZeRO-Offload) den Optimierer-Zustand teilweise oder vollständig auf den Host-Speicher (CPU-RAM). Dies führt jedoch zu zwei Hauptengpässen:

• I/O-Engpass: Der Datentransfer zwischen GPU und CPU über PCIe-Links ist langsam (typisch 25–50 GB/s) und wird oft nicht vollständig ausgelastet.
• Rechenengpass: CPUs sind in der Lage, Optimierer-Updates (insbesondere bei adaptiven Optimierern wie Adam) um Größenordnungen langsamer durchzuführen als GPUs.
• Ineffiziente Überlappung: Bestehende hybride Ansätze (z. B. DeepSpeed TwinFlow) nutzen statische Partitionierung. Ein fester Teil des Optimierers bleibt auf der GPU, der Rest auf der CPU. Dies führt zu Leerlaufzeiten, da während der GPU-Berechnungen die CPU oft wartet und umgekehrt, und die PCIe-Bandbreite während der Backward- und Update-Phasen unterausgelastet bleibt.

2. Methodik und Systemdesign

Die Autoren schlagen Deep Optimizer States vor, eine Middleware-Lösung, die auf dem Prinzip des interleaved offloading (verschachteltes Auslagern) basiert. Statt statischer Zuweisung wird der Optimierer-Zustand dynamisch zwischen CPU und GPU verschoben, um die Nutzung von Speicher und Bandbreite zu maximieren.

Kernprinzipien:

• Granulare Subgruppen: Der Optimierer-Zustand wird in kleine Subgruppen (Subgroups) unterteilt (basierend auf ZeRO-3 Sharding). Diese können unabhängig voneinander aktualisiert werden, da die Updates für verschiedene Subgruppen keine Abhängigkeiten zueinander haben (embarrassingly parallel).
• Dynamische Scheduling-Strategie: Anstatt einen festen Anteil auf der GPU zu halten, wird für jede Iteration entschieden, welche Subgruppen auf der GPU und welche auf der CPU aktualisiert werden.
• Überlappung von Berechnung und Transfer: Das System nutzt asynchrone Datenbewegungen (CUDA Streams), um folgende Operationen parallel auszuführen:
  • CPU berechnet Updates für Subgruppe $i$.
  • GPU berechnet Updates für Subgruppe $j$.
  • Asynchroner Transfer (H2D/D2H) von Gradienten und Parametern für Subgruppe $k$.
• Präzisions-Optimierung: Um teure Speicherzuweisungen und On-the-Fly-Konvertierungen zu vermeiden, werden Gradienten in FP32 auf der GPU konvertiert (wo die Bandbreite hoch ist) und dann als FP32 über PCIe transferiert, statt FP16 zu transferieren und auf der CPU zu konvertieren.
• Leistungsmodell: Ein mathematisches Modell (basierend auf Durchsatzraten von CPU, GPU und PCIe) berechnet den optimalen „Update-Stride" ($k$). Dies bestimmt, wie viele Subgruppen auf der CPU aktualisiert werden müssen, bevor eine auf der GPU aktualisiert wird, um die Pipeline optimal zu füllen.

Algorithmus:
Der Algorithmus (Algorithm 1 im Paper) steuert den Update-Prozess:

1. Prüfung, ob eine Subgruppe statisch auf der GPU liegt oder zum nächsten GPU-Update-Schritt gehört.
2. Falls GPU-Update: Asynchrones Prefetching der nächsten Subgruppe und Flushen der vorherigen.
3. Falls CPU-Update: Berechnung und asynchrones Downscaling der Parameter.
4. Nutzung dedizierter CUDA-Streams für gleichzeitige H2D- und D2H-Transfers (Full-Duplex).

3. Wichtige Beiträge

1. Analyse der Systemcharakteristika: Detaillierte Untersuchung zeigt, dass die GPU-Speichernutzung während der Update-Phase drastisch sinkt (da Aktivierungen freigegeben werden) und die PCIe-Bandbreite in allen Phasen unterausgelastet ist.
2. Neues Design-Paradigma: Einführung des „Interleaved Offloading", das statische Partitionierung durch dynamische, iterativ gesteuerte Verschiebung ersetzt.
3. Leistungsmodell: Entwicklung eines Modells zur Berechnung des optimalen Verhältnisses von CPU- zu GPU-Updates ($k$), um Überlappung zu maximieren.
4. Implementierung: Integration in DeepSpeed und Megatron-LM als Open-Source-Middleware, die mit ZeRO-3 kompatibel ist.
5. Ergebnisse: Nachweis signifikanter Beschleunigungen durch extensive Experimente.

4. Ergebnisse

Die Evaluation wurde auf einem Testsystem mit 4× NVIDIA H100 GPUs (80 GB) und 2× Intel Xeon Prozessoren durchgeführt.

• Geschwindigkeit: Deep Optimizer States erreicht bis zu 2,5-fach schnellere Iterationen im Vergleich zu DeepSpeed ZeRO-3 (vollständiges CPU-Offloading) und DeepSpeed TwinFlow (statisches Hybrid-Offloading).
• Update-Durchsatz: Die Rate der aktualisierten Parameter pro Sekunde steigt im Durchschnitt um 70 %.
• Speichereffizienz: Die Methode ermöglicht das Training von Modellen mit bis zu 20 Milliarden Parametern auf einem einzelnen Node mit begrenztem GPU-Speicher, ohne dass die Iterationszeit linear mit der Modellgröße wächst.
• Ressourcennutzung:
  • GPU-Auslastung steigt auf fast 100 % (im Vergleich zu ~30 % bei reinem CPU-Offloading).
  • PCIe-Auslastung wird effizienter genutzt (~40 % der Peak-Bandbreite).
  • Die CPU-Auslastung sinkt leicht (auf ~60 %), da die PCIe-Transfers und die GPU-Berechnungen die Engpässe verschieben, was jedoch zu einem höheren Gesamtdurchsatz (TFLOPs) führt.
• Skalierbarkeit: Die Methode skaliert gut mit verschiedenen Datenparallelismus-Graden (bis zu 4-fache Beschleunigung bei niedrigem Parallelismus, immer noch >2,5-fach bei höherem Parallelismus).

5. Bedeutung und Ausblick

Die Arbeit adressiert eine kritische Hürde beim Training großer Modelle auf ressourcenbeschränkter Hardware. Durch die intelligente Ausnutzung der Fluktuationen im Speicherverbrauch und der Bandbreite ermöglicht Deep Optimizer States:

• Effizienteres Fine-Tuning von LLMs auf einzelnen Nodes (wichtig für spezialisierte Anwendungen).
• Bessere Auslastung heterogener Systeme (CPU + GPU).
• Eine Grundlage für zukünftige Architekturen mit noch höherer CPU-GPU-Bandbreite (z. B. NVIDIA Grace Hopper Superchips), wo dynamisches Offloading noch leistungsfähiger sein wird.

Zusammenfassend transformiert Deep Optimizer States das Problem des Speicherengpasses von einem statischen Limit in ein dynamisches Ressourcenmanagement-Problem, das durch geschickte Überlappung von I/O und Berechnung gelöst wird.