🤖 machine learning

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

Das Papier stellt Asteria vor, ein Laufzeitsystem, das eine praktische, skalierbare Optimierung zweiter Ordnung für große Sprachmodelle ermöglicht, indem es Optimiererzustände dynamisch über heterogene Speicherhierarchien verteilt und teure Präkonditionierberechnungen vom kritischen GPU-Trainingspfad entkoppelt, um den Overhead zu reduzieren und die Konvergenz zu beschleunigen.

Ursprüngliche Autoren: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Veröffentlicht 2026-05-18

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen, einem riesigen, superschlauen Roboter (einem Large Language Model) beizubringen, Geschichten zu schreiben. Dafür benötigen Sie einen „Lehrer" (einen Optimierer), der die Fehler des Roboters betrachtet und ihm sagt, wie er sich verbessern kann.

Derzeit nutzen die meisten Lehrer eine einfache Methode namens AdamW. Sie ist wie ein Schüler, der seine Hausaufgaben überprüft, ein paar falsche Antworten sieht und kleine, schnelle Korrekturen vornimmt. Es ist schnell und funktioniert gut, aber es erfordert viele Übung (Millionen von Beispielen), um wirklich gut zu werden.

Es gibt eine schlauere, fortgeschrittenere Lehrermethode namens Second-Order Optimization (wie SOAP oder Shampoo). Dieser Lehrer schaut nicht nur darauf, was falsch war; er analysiert die Form der Fehler, um das Problem tiefgründig zu verstehen. Er lernt viel schneller und benötigt weniger Beispiele. Allerdings gibt es einen riesigen Haken: Dieser clevere Lehrer ist unglaublich schwer. Er erfordert enorme Mengen an Speicher und Rechenleistung für seine komplexen Mathematikaufgaben, was oft dazu führt, dass der Computer abstürzt oder so langsam läuft, dass der einfache Lehrer die Arbeit tatsächlich zuerst beendet.

Hier kommt Asteria ins Spiel: Der „intelligente Logistikmanager"

Die Forscher an der Universität Oxford haben ein neues System namens Asteria entwickelt. Denken Sie an Asteria nicht als neuen Lehrer, sondern als einen brillanten Logistikmanager, der das gesamte Klassenzimmer neu organisiert, damit der clevere Lehrer endlich seine Arbeit verrichten kann, ohne die Schule zu zerstören.

So löst Asteria die drei größten Probleme, wobei einfache Analogien verwendet werden:

1. Das Problem „Möbel in einem winzigen Raum" (Speicher)

Das Problem: Der clevere Lehrer muss riesige, komplexe Diagramme (Matrizen) in seinem Kopf (dem GPU-Speicher) behalten. Auf einem Standardcomputer ist nicht genug Platz vorhanden. Es ist, als würde man versuchen, ein Kingsize-Bett, einen Esstisch und einen Kleiderschrank in ein Studio-Apartment zu quetschen. Der Raum wird voll, und der Lehrer muss aufgeben.
Die Asteria-Lösung: Asteria agiert wie ein Experte für klappbare Möbel. Er erkennt, dass der Lehrer nicht alles gleichzeitig in den Händen halten muss.

Er behält die aktivsten Teile der Diagramme auf der GPU (dem Schreibtisch).
Er verlagert die schweren, weniger häufig verwendeten Teile auf die CPU (den Flur) oder sogar auf eine Festplatte (die Garage).
Entscheidend ist, dass er die schweren Teile nur genau dann wieder auf den Schreibtisch bringt, wenn sie benötigt werden. Dies ermöglicht es dem cleveren Lehrer, auf einem winzigen Laptop genauso gut zu arbeiten wie auf einem Supercomputer.

2. Das Problem „Stau" (Geschwindigkeit)

Das Problem: Alle paar Schritte muss der clevere Lehrer eine massive, komplexe Berechnung durchführen (wie das Lösen eines riesigen Puzzles), um seine Diagramme zu aktualisieren. Dies dauert lange und hält die ganze Klasse davon ab, zu arbeiten. Es ist wie ein Lieferwagen, der die gesamte Autobahn blockiert, während er ein einziges Paket entlädt. Die GPU (das Gehirn des Computers) sitzt untätig da und wartet.
Die Asteria-Lösung: Asteria führt eine parallele Fließbandproduktion ein.

Anstatt den Hauptunterricht zu unterbrechen, um die schwere Mathematik zu erledigen, schickt Asteria die „schwere Arbeit" an ein Team von Arbeitern im Hinterzimmer (die CPU).
Während die Hauptklasse (die GPU) weiter dem Roboter Unterricht erteilt, lösen die Arbeiter im Hinterzimmer leise die komplexen Puzzles im Hintergrund.
Bis die Klasse die neuen Diagramme benötigt, hat das Hinterzimmer sie bereits fertiggestellt und herübergeschoben. Die Hauptklasse muss niemals anhalten. Der „Stau" verschwindet.

3. Das Problem „Gruppenchat" (Verteiltes Training)

Das Problem: Wenn gleichzeitig mit vielen Computern trainiert wird, müssen alle normalerweise anhalten und sich auf exakt dieselben Informationen einigen, bevor sie zum nächsten Schritt übergehen. Dies ist wie ein Gruppenchat, in dem alle warten müssen, bis die langsamste Person antwortet, bevor jemand wieder tippen kann.
Die Asteria-Lösung: Asteria verwendet eine „Gut genug"-Politik.

Anstatt darauf zu warten, dass alle perfekt synchronisiert sind, erlaubt es Computern, für kurze Zeit mit leicht „veralteten" (etwas veralteten) Informationen zu arbeiten.
Es sendet nur Updates, wenn sie wirklich notwendig sind.
Dies hält die Gruppe schnell in Bewegung, wie ein Team, das einander vertraut, weiterzuarbeiten und Updates später nachzuholen, anstatt das gesamte Projekt alle fünf Minuten zu stoppen.

Die Ergebnisse: Was haben sie herausgefunden?

Die Forscher testeten Asteria auf echter Hardware, einschließlich eines leistungsstarken einzelnen Computers (Nvidia DGX Spark) und eines großen Clusters von Computern (Nvidia GH200).

Es funktioniert auf kleinen Maschinen: Sie konnten ein großes Sprachmodell (1 Milliarde Parameter) auf einer einzelnen Maschine trainieren, die zuvor die Speicheranforderungen dieses cleveren Lehrers nicht bewältigen konnte.
Es ist in Echtzeit schneller: Da es die „Staus" versteckt, endet das Training in weniger tatsächlicher Uhrzeit, obwohl die Mathematik schwieriger ist.
Es spart Energie: Indem es das Gehirn des Computers (GPU) beschäftigt hält und nicht zulässt, dass es untätig auf Berechnungen wartet, verbraucht Asteria tatsächlich weniger Gesamtenergie, um dasselbe Ergebnis zu erzielen, verglichen mit den alten, klobigen Methoden.
Es verliert keine Qualität: Das Modell lernt genauso gut wie mit dem „nativen" (nicht optimierten) cleveren Lehrer, aber es kommt viel schneller und günstiger dorthin.

Zusammenfassung:
Asteria erfindet keine neue mathematische Formel. Stattdessen überdenkt es, wie der Computer diese Mathematik ausführt. Es trennt die schwere Arbeit von der Hauptarbeit, nutzt allen verfügbaren Speicherplatz klug aus und lässt Computer asynchron arbeiten. Dies verwandelt eine „theoretisch großartige, aber praktisch unmögliche" Methode in ein praktisches Werkzeug für das Training der nächsten Generation von KI.

Technischer Zusammenfassung: Asteria – Laufzeit-orchestrierte Optimierung zweiter Ordnung für skalierbares LLM-Training

1. Problemstellung

Optimierungsmethoden zweiter Ordnung (z. B. Shampoo, SOAP, KL-Shampoo) bieten im Vergleich zu Optimierern erster Ordnung wie AdamW eine überlegene statistische Effizienz und schnellere Konvergenz beim Training von Large Language Models (LLMs). Ihre praktische Bereitstellung im großen Maßstab wird jedoch durch schwerwiegende systemische Barrieren behindert:

Speicherbeschränkungen: Methoden zweiter Ordnung erfordern das Speichern großer Kovarianzmatrizen ( $L$ und $R$ ) und die Durchführung dichter Eigenwertzerlegungen, was zu einem Speicherbedarf von $O(N^2)$ und Rechenkosten von $O(N^3)$ führt. In speicherbeschränkten Umgebungen (z. B. Single-Node-Workstations mit Unified Memory Architectures) erschöpfen diese Zustände den physischen Speicher, was zu Out-of-Memory (OOM)-Fehlern führt oder eine Verringerung der Modellgröße und der Batch-Größe erzwingt.
Latenzspitzen: Die periodische Berechnung von inversen Wurzeln (z. B. $L^{-1/4}$ ) beinhaltet Matrixoperationen mit kubischen Kosten. In bestehenden Implementierungen blockieren diese synchronen Updates den GPU-Trainingspipeline, was zu schweren periodischen Latenzspitzen führt, die die für effizientes verteiltes Training (z. B. FSDP) essentielle feingranulare Überlappung von Berechnung und Kommunikation stören.
Synchronisationsaufwand: Aktuelle verteilte Methoden zweiter Ordnung verlassen sich auf starre, synchrone globale Zustandsupdates (AllGather) über alle Ranks hinweg. Dies koppelt schnelle kommunikation innerhalb des Knotens mit langsamer Synchronisation zwischen Knoten und verhindert, dass das System asynchrone, topologiebewusste Staleness-Grenzen ausnutzt.

Bestehende Lösungen, wie z. B. DTensor-basiertes Sharding, versagen bei der Behebung dieser Probleme in ressourcenbeschränkten Single-Node-Umgebungen oder wenn der Kommunikationsaufwand der globalen Synchronisation die Vorteile überwiegt.

2. Methodik: Die Asteria-Laufzeitumgebung

Asteria ist ein hardware-software-co-designtes Laufzeitsystem, das auf PyTorch FSDP aufbaut. Es entkoppelt die Logik der Optimierung zweiter Ordnung vom kritischen GPU-Trainingspfad durch drei Kernmechanismen:

A. Architekturadaptive asymmetrische Speichertiering

Im Gegensatz zu Optimierern erster Ordnung, die Zustände als homogene Vektoren behandeln, erkennt Asteria den asymmetrischen Lebenszyklus von Zuständen zweiter Ordnung:

GPU: Speichert factor_matrices (Kronecker-Faktoren), die während des Trainings häufig aktualisiert werden.
Host (CPU/UVM): Speichert inv_factor_matrices (inverse Faktoren). Diese werden im CPU-zugänglichen Speicher gehalten, um Speicherdruck auf der GPU zu vermeiden.
NVMe (Optional): Bietet eine Staging-Ebene für kapazitätsbegrenzte Ausführungen, wodurch Zustände über serialisierte Berechnung oder JIT-Paging ausgelagert werden können.
Mechanismus: Die Laufzeitumgebung erfasst Host-seitige Snapshots von Faktormatrizen, delegiert asynchrone CPU-Worker zur Berechnung inverser Wurzeln und schreibt die Ergebnisse zurück in den Host-Speicher. Bei Bedarf greift die GPU über Unified Virtual Memory (UVM)-Ansichten oder explizites Prefetching darauf zu, wodurch redundante Host-zu-Gerät-Kopien vermieden werden.

B. Hook-orchestrierte Shadow-State-Pipeline

Um Latenzspitzen zu eliminieren, entfernt Asteria die $O(N^3)$ -Berechnung inverser Wurzeln aus dem kritischen GPU-Pfad:

Shadow Stream: Ein CUDA-Stream mit niedriger Priorität übernimmt das GPU-seitige State-Staging (Prefetching und Paging), ohne den primären Trainingsstream zu blockieren.
Asynchrone Ausführung: Kostspielige Updates inverser Wurzeln werden asynchron auf CPU-Worker-Threads ausgeführt.
Hook-Integration: Leichte PyTorch-Hooks (register_forward_hook, register_full_backward_pre_hook) fungieren als Scheduling-Signale. Sie lösen das Ablaufen bereiter Zustandsübertragungen und das Prefetching bald zu verbrauchender inverser Faktoren aus, wodurch sichergestellt wird, dass die GPU frische Konditionierer hat, ohne die Hauptschleife des Trainings zu blockieren.

C. Topologiebewusste Staleness-begrenzte Kohärenz

Für verteiltes Training ersetzt Asteria starre globale Synchronisation durch ein selektives Kohärenzprotokoll:

Staleness-Verfolgung: Ein CoherenceRegistry verfolgt die Frische jedes Konditionierer-Blocks. Nur Blöcke, die ein konfiguriertes Staleness-Budget überschreiten, werden synchronisiert.
Hierarchische Synchronisation: Die Synchronisation erfolgt in zwei Stufen: zuerst innerhalb von Knoten-lokalen Prozessgruppen, dann über Knoten-Repräsentanten hinweg. Dies minimiert den Verkehr zwischen Knoten und vermeidet unnötige CPU-zu-GPU-Rundreisen, indem die Kommunikation im host-zugänglichen Speicher gehalten wird.
Begrenzter Barrier: Die Laufzeitumgebung erzwingt eine lokale Barriere, die sicherstellt, dass, falls eine Version eines Konditionierers zu alt wird, während ein CPU-Job läuft, der Optimierer auf den Abschluss wartet, wodurch Frische und Durchsatz ausbalanciert werden.

3. Hauptbeiträge

Laufzeit-Entkopplung: Asteria zeigt, dass Optimierung zweiter Ordnung praktisch gemacht werden kann, nicht durch Vereinfachung des Algorithmus, sondern durch ein Umdenken im State-Management, der Hintergrundberechnung und der Synchronisation auf Laufzeitumgebungsebene.
Speicherskalierbarkeit: Es ermöglicht das Training zweiter Ordnung auf Single-Node-Hardware mit Speicherbeschränkungen (z. B. eine einzelne GB10-GPU mit 128 GB Unified Memory), wo traditionelles Sharding unmöglich ist.
Latenzversteckung: Durch die Verlagerung der Berechnung inverser Wurzeln auf den Host und die Nutzung einer Shadow-Pipeline glättet Asteria periodische Latenzspitzen und erreicht ein glattes Hardware-Ausführungsprofil, das mit Methoden erster Ordnung vergleichbar ist.
Energieeffizienz: Das System verbessert den Energie-Verlust-Tradeoff und erzielt eine größere Verlustreduktion pro Energieeinheit durch die Reduzierung von GPU-Idle-Zeiten und die Optimierung der Überlappung von Host-seitigen Berechnungen.

4. Experimentelle Ergebnisse

Die Autoren bewerteten Asteria an OLMo-Familienmodellen (660M, 1B, 7B Parameter) unter Verwendung des OLMo-Trainingsframeworks auf zwei Plattformen:

Single-Node (DGX Spark, GB10 GPU):
- Durchsatz: Asteria-SOAP und Asteria-KL-Shampoo eliminierten die periodischen Schrittzeit-Spitzen, die in nativen Implementierungen zu sehen waren (Reduktion der Spitzen-Schrittzeit von ~96s auf ~66,5s).
- Energie: Während native Methoden zweiter Ordnung insgesamt ~17–20 % mehr SoC-Energie als AdamW verbrauchten, reduzierte Asteria diese Lücke erheblich (z. B. verbrauchte Asteria-KL-Shampoo nur ~7,9 % mehr Energie als AdamW).
- Effizienz: Asteria-KL-Shampoo erreichte die höchste normalisierte Verlustreduktionseffizienz (3,3010) und übertraf sowohl native Methoden zweiter Ordnung als auch AdamW.
Multi-Node (GH200 Cluster):
- Konvergenz: Asteria-Varianten behielten Konvergenzverläufe bei, die ihren nativen Gegenstücken zweiter Ordnung nahe kamen, und übertrafen AdamW in der Wandzeit signifikant, um Zielverlustniveaus zu erreichen.
- Staleness-Toleranz: Experimente mit variierendem Staleness-Budget ( $S$ ) zeigten, dass eine Erhöhung von $S$ von 1 auf 5 die Trainingszeit erheblich reduzierte, ohne die endgültige Modellqualität nennenswert zu beeinträchtigen.
- Strong Scaling: Bei 7B-Modellen, die von 2 auf 16 Knoten skalierten, erzielte Asteria im Vergleich zu nativen Baselines einen höheren Speedup und eine geringere Zeit pro Schritt, was sich auch als wirksam erwies, wenn die relativen Kosten der Zustandspflege zunahmen.

5. Bedeutung und Behauptungen

Die Arbeit argumentiert, dass das Haupthindernis für das praktische Training von LLMs mit Methoden zweiter Ordnung nicht die mathematische Kosten der krümmungsbewussten Updates sind, sondern die System-Inkompatibilität zwischen der Zustandspflege zweiter Ordnung und modernen Trainingsstapeln.

Neudefinition des Problems: Die Autoren behaupten, dass Optimierung zweiter Ordnung als Co-Design-Problem betrachtet werden sollte, das Algorithmen, Laufzeit-Scheduling und heterogene Speichersysteme umfasst.
Durchführbarkeit: Sie zeigen, dass großer aggregierter GPU-Speicher keine grundlegende Voraussetzung für das Training zweiter Ordnung ist; durch lebenszyklusbewusste Platzierung und gestuften Zugriff können diese Methoden auch bei engen Speicherbudgets ausgeführt werden.
Praktikabilität: Indem Asteria die Frische von Konditionierern und die Kommunikation als laufzeitverwaltete Ressourcen behandelt und nicht als starre globale Zustände, verringert es die Kluft zwischen der statistischen Versprechen der Optimierung zweiter Ordnung und ihrer realweltlichen Einsetzbarkeit.

Die Arbeit schließt, dass Asteria die Optimierung zweiter Ordnung zu einer gangbaren Systemwahl für das LLM-Training macht und bessere Konvergenz sowie Energieeffizienz bietet, ohne algorithmische Vereinfachungen zu erfordern.