Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Training

Kareus ist ein Trainingssystem, das den dynamischen und statischen Energieverbrauch beim Training großer Modelle gleichzeitig optimiert, indem es das gemeinsame Optimierungsproblem in partitionsbasierte Teilprobleme zerlegt und einen Multi-Pass-Multi-Objective-Algorithmus einsetzt, wodurch im Vergleich zu aktuellen State-of-the-Art-Methoden signifikante Reduktionen entweder der Trainingszeit oder des Energieverbrauchs erreicht werden.

Das Wesentliche

Das große Problem: Die Energie-Rechnung der KI

Stellen Sie sich vor, das Training eines massiven KI-Modells ist wie der Betrieb einer riesigen, hochmodernen Fabrik. Diese Fabrik benötigt zwei Dinge, um am Laufen zu bleiben: Elektrizität (Energie) und Zeit.

Momentan bauen die Menschen diese Fabriken schneller auf, als wir die Elektrizität dafür liefern können. Die Arbeitspapier argumentiert, dass Energie zu einer knappen, teuren Ressource geworden ist. Wenn wir sie nicht besser verwalten, können wir die KI einfach nicht auf die nächste Stufe heben.

Die zwei Arten von „verschwendeter“ Energie

Um die Lösung zu verstehen, müssen wir zuerst verstehen, wie diese Fabriken (GPUs) Energie verschwenden. Die Autoren unterteilen dies in zwei Arten:

1. Dynamische Energie (Die „Arbeitskosten“): Dies ist die Elektrizität, die verbraucht wird, wenn die Maschine tatsächlich Mathematik betreibt. Denken Sie an einen Automotor, der aufheult, während Sie fahren. Je schneller Sie fahren (höhere Frequenz), desto mehr Benzin verbrennen Sie.
2. Statische Energie (Die „Leerlaufkosten“): Dies ist die Elektrizität, die die Maschine verbraucht, nur um eingeschaltet zu bleiben, selbst wenn sie nichts Nützliches tut. Denken Sie an einen Automotor, der im Leerlauf an einer roten Ampel steht. Wenn Sie 5 Minuten an einer roten Ampel warten, verbrennen Sie Benzin, obwohl Sie sich nicht bewegen. Beim KI-Training passiert dies, wenn der Computer darauf wartet, dass Daten eintreffen, oder wenn Teile des Chips einfach nur untätig sind.

Die alten Lösungen: Das eine reparieren, das andere ignorieren

Vor dieser Arbeit versuchten Forscher, diese Probleme zu lösen, aber sie betrachteten nur eine Seite der Medaille:

• Lösung A (Perseus): Sie drosselten den Motor (senkten die Frequenz), wenn die Maschine nicht auf dem „kritischen Pfad“ (dem wichtigsten Teil der Arbeit) war. Dies sparte dynamische Energie (Kraftstoff), löste aber das Problem der Zeit im Leerlauf nicht.
• Lösung B (Nanobatching): Sie ordneten die Arbeit so um, dass ein anderer Teil der Maschine bereits arbeiten konnte, während ein anderer Teil noch auf Daten wartete. Dies reduzierte die Gesamtzeit, was statische Energie sparte (weniger Leerlauf), aber sie passten die Geschwindigkeit des Motors nicht an.

Das Papier argumentiert, dass es ist, als würde man versuchen, ein leckendes Boot zu reparieren, indem man entweder nur die Löcher auf der linken Seite flickt oder nur das Wasser auf der rechten Seite herausschöpft. Man muss beides gleichzeitig tun.

Die Entdeckung: Alles ist miteinander verbunden

Die Autoren entdeckten, dass diese beiden Probleme tief miteinander verknüpft sind. Man kann nicht einfach entscheiden, den „Motor zu drosseln“ oder die „Arbeit umzuorganisieren“, ohne das andere zu berücksichtigen.

Die Analogie: Stellen Sie sich eine Küche mit einem Koch (Berechnung) und einem Kellner (Kommunikation) vor.

• Wenn der Koch zu schnell arbeitet, kommt der Kellner nicht hinterher, und der Koch muss anhalten und warten (Zeitverschwendung/Leerlauf).
• Wenn der Koch langsamer arbeitet, kommt der Kellner vielleicht zu früh fertig und steht nur da und wartet (ebenfalls Zeitverschwendung).
• Entscheidend ist: Die beste Art, die Küche zu organisieren, hängt davon ab, wie schnell der Koch arbeitet. Wenn der Koch langsam ist, möchte man vielleicht, dass der Kellner früher beginnt. Wenn der Koch schnell ist, möchte man vielleicht, dass der Kellner wartet.

Das Papier fand heraus, dass die Änderung der Frequenz (Motorgeschwindigkeit) die optimale Planung (Küchenlayout) verändert. Bestehende Werkzeuge behandelten diese als getrennte Probleme, was zu suboptimalen Ergebnissen führte.

Die Lösung: Kareus

Kareus ist ein neues System, das wie ein superintelligenter Fabrikmanager fungiert. Es wählt nicht einfach nur eine Einstellung, sondern betrachtet das Gesamtbild und findet das perfekte Gleichgewicht zwischen Geschwindigkeit, Arbeit und Leerlauf.

So funktioniert es, unter Verwendung einer Lego-Analogie:

1. Das Aufteilen (Partitionierung): Anstatt zu versuchen, das gesamte massive KI-Modell auf einmal zu optimieren (was zu komplex wäre), zerlegt Kareus den Trainingsprozess in kleine, sich wiederholende Lego-Blöcke, die „Partitionen“ genannt werden. Es löst das Rätsel zuerst für einen kleinen Block.
2. Die Multi-Ziel-Optimierung: Für jeden Lego-Block führt Kareus eine anspruchsvolle Suche durch. Es fragt: „Wenn ich die Motorgeschwindigkeit ändere, wie viele Arbeiter sollte ich dem Kellner zuweisen? Wann sollte der Kellner beginnen?“ Es testet tausende Kombinationen, um die „Pareto-Front“ zu finden – den absolut besten Kompromiss zwischen schnellem Abschluss und geringem Energieverbrauch.
3. Das Zusammenfügen: Sobald es weiß, wie jeder kleine Lego-Block am besten betrieben wird, fügt es sie wieder zusammen, um einen Plan für die gesamte Fabrik zu erstellen.

Die Ergebnisse: Ein Win-Win

Das Papier testete Kareus an echten KI-Modellen (wie Llama und Qwen) und verglich es mit den besten bestehenden Systemen. Die Ergebnisse waren beeindruckend:

• Szenario 1 (Gleiche Zeit): Wenn Sie möchten, dass die KI in exakt der gleichen Zeit wie zuvor fertig wird, verbraucht Kareus bis zu 28,3 % weniger Energie.
• Szenario 2 (Gleiche Energie): Wenn Sie ein striktes Energiebudget haben (wie eine feste Menge an Elektrizität), kann Kareus das Training bis zu 27,5 % schneller abschließen.

Warum das wichtig ist

Das Papier kommt zu dem Schluss, dass wir nicht einfach mehr Hardware in das Problem werfen können, um immer größere und intelligentere KIs zu bauen. Wir brauchen intelligentere Software, die versteht, dass Geschwindigkeit, Planung und Energie ein einziges, großes, zusammenhängendes Puzzle sind. Kareus ist das erste System, das dieses Puzzle gleichzeitig löst, wodurch enorme Mengen an Elektrizität und Zeit gespart werden.

Technische Zusammenfassung

Technisches Resümee: Kareus – Gemeinsame Reduktion von dynamischer und statischer Energie beim Training großer Modelle

Problemstellung

Das rasante Wachstum der KI-Rechenleistung hat das Energieangebot überholt, wodurch Energie zu einer teuren und umkämpften Ressource für das Training großer Modelle geworden ist. Bestehende Optimierungsbemühungen adressieren typischerweise entweder die dynamische Energie (die durch aktive Berechnung verbraucht wird, reduzierbar durch Frequenzskalierung) oder die statische Energie (die unabhängig von der Aktivität durch die Hardware verbraucht wird, reduzierbar durch Verkürzung der Ausführungszeit).

Bestehende Lösungen leiden unter einer kritischen Einschränkung: Sie optimieren diese Faktoren entweder isoliert oder durch naive Kombinationen.

• Ansätze zur Frequenzskalierung (z. B. Perseus) reduzieren die dynamische Energie durch Senkung der GPU-Frequenzen für nicht-kritische Pfade, versäumen es jedoch oft, Möglichkeiten zur Überlappung von Kommunikation und Berechnung zu nutzen, da sie sequenziellen Kernel-Ausführungen folgen.
• Feinkörnige Kernel-Scheduling-Ansätze (z. B. Nanobatching) reduzieren die statische Energie durch die Überlappung von Kommunikation und Berechnung, ignorieren jedoch die dynamische Energie und die Frequenzskalierung.

Das Paper identifiziert, dass SM-Allokation, Kernel-Launch-Timing und GPU-Frequenz interdependent sind. Die Änderung eines Faktors verschiebt die optimale Konfiguration der anderen. Beispielsweise macht eine niedrigere Frequenz die Kernel relativ rechengebunden (compute-bound), was die optimalen Ressourcenkonflikte und Überlappungsstrategien verändert. Folglich scheitert die Optimierung von Teilmengen dieser Faktoren oder deren naive Kombination daran, die Zeit–Energie-Tradeoff-Grenze (Pareto-Front) auf ihr theoretisches Limit zu drücken.

Methodik

Die Autoren schlagen Kareus vor, ein Trainingssystem, das darauf ausgelegt ist, SM-Allokation, Kernel-Launch-Timing und GPU-Frequenz gemeinsam zu optimieren, um die Pareto-optimale Zeit–Energie-Front zu finden.

1. Partitioniertes Overlap-Ausführungsmodell

Um die globale Optimierung handhabbar zu machen, zerlegt Kareus das Problem in lokale Teilprobleme unter Verwendung eines partitionierten Overlap-Ausführungsmodells:

• Partitionierung: Der Ausführungsgraph der Forward- und Backward-Passes wird in feste Partitionen unterteilt. Eine Partition besteht aus einem Kommunikations-Kernel (z. B. AllReduce) und der längsten zusammenhängenden Sequenz von Kompositions-Kernels eines anderen Microbatches, die keine Datenabhängigkeiten zum Kommunikations-Kernel aufweist.
• Generalisierung: Dieses Modell generalisiert Nanobatching-Techniken, indem es eine feinkörnige Kontrolle über SM-Allokation und Launch-Timing für jede Partition ermöglicht.
• Hierarchische Komposition: Das System optimiert jede Partition unabhängig und setzt diese lokalen Fronten dann hierarchisch zu einer globalen Microbatch- und Iterations-Front zusammen.

2. Multi-Pass Multi-Objective Bayesian Optimization (MBO)

Da die erschöpfende Suche im Konfigurationsraum (der SM-Anzahlen, Launch-Timings und Frequenzen umfasst) prohibitiv teuer ist (sie würde tausende GPU-Stunden erfordern und thermische Interferenzen verursachen), setzt Kareus einen spezialisierten MBO-Algorithmus ein:

• Surrogatmodelle: Es verwendet zwei separate Gradient-Boosted Decision Tree (XGBoost)-Modelle, um die Ausführungszeit bzw. die dynamische Energie jeweils separat vorherzusagen.
• Multi-Pass Kandidatenselektion: Um die Pareto-Front effektiv zu explorieren, nutzt der Algorithmus mehrere Akquisitionsfunktionen parallel:
  • Hypervolume Improvement (HVI): Drei verschiedene HVI-Funktionen leiten die Suche hin zu niedrigerer Gesamtenergie, niedrigerer dynamischer Energie (bevorzugt niedrigere Frequenzen) und niedrigerer statischer Energie (bevorzugt schnellere Ausführung).
  • Unsicherheits-Exploration: Ein Bootstrap-Ensemble von Surrogatmodellen schätzt die Vorhersageunsicherheit und priorisiert Kandidaten in unterexplorierten Regionen des Suchraums.
• Thermisch stabile Profilierung: Um genaue Energiemessungen zu gewährleisten, führt das System die Partitionen über ein Fenster von 5 Sekunden aus und legt zwischen den Konfigurationen eine 5-sekündige Abkühlphase ein, um thermische Bias zu verhindern.

3. Systemimplementierung

Kareus wird durch die Integration mit Megatron-LM und Perseus implementiert.

• Es nutzt MSCCL++ für die feinkörnige Kontrolle der SM-Allokation.
• Es verwendet eine benutzerdefinierte torch.autograd.Function, um Partitionen zu kapseln, was unterschiedliche Schedules für Forward- und Backward-Passes ermöglicht.
• Zur Laufzeit wählt es einen Ausführungsplan aus der vorberechneten Front basierend auf Benutzerbeschränkungen (z. B. Zeitdeadline oder Energiebudget) aus und passt die GPU-Frequenzen dynamisch über einen von Perseus adaptierten Controller an.

Kernergebnisse

Kareus wurde auf 14 Workloads evaluiert, einschließlich realen Trainings von Llama 3.2 3B und Qwen 3 1.7B sowie einer groß angelegten Emulation von Llama 3.3 70B auf 16 NVIDIA A100 GPUs.

• Zeit–Energie-Tradeoff: Im Vergleich zum State-of-the-Art-Baseline-System (Megatron-LM + Perseus) erreicht Kareus:
  • Bis zu 28,3 % Reduktion der Trainingsenergie bei gleicher Trainingszeit.
  • Bis zu 27,5 % Reduktion der Trainingszeit bei gleichem Energieverbrauch.
• Maximaler Durchsatz: Im Modus für maximalen Durchsatz reduziert Kareus die Iterationszeit um bis zu 14,9 % und die Energie um bis zu 22,1 % im Vergleich zur Baseline Megatron-LM.
• Ablationsstudien: Das Entfernen entweder der Frequenzskalierung oder des Kernel-Schedulings führt zu einer signifikanten Leistungsverschlechterung (z. B. führt das Entfernen des Kernel-Schedulings zu einer um 10,8 % höheren Energie), was die Notwendigkeit der gemeinsamen Optimierung bestätigt.
• Skalierbarkeit: Die Ergebnisse der groß angelegten Emulation für Llama 3.3 70B zeigen konsistente Trends mit den Ergebnissen des physischen Testbeds, was die Anwendbarkeit des Systems auf massive Modelle demonstriert.

Bedeutung und Ansprüche

Das Paper behauptet, dass Kareus die Zeit–Energie-Tradeoff-Front für das Training großer Modelle verschiebt, indem es zeigt, dass dynamische und statische Energie nicht unabhängig voneinander optimiert werden können. Der Kernbeitrag liegt in der Erkenntnis, dass Ausführungspläne (SM-Allokation und Timing) und GPU-Frequenz gemeinsam und interdependent bestimmt werden.

Die Autoren betonen:

1. Gemeinsame Kontrolle ist überlegen: Die gemeinsame Optimierung von Frequenz und Scheduling ist nachweislich energieeffizienter als die Überlassung der Frequenzsteuerung an Hardware-Leistungsregler oder die Optimierung des Schedulings ohne Berücksichtigung der Frequenz.
2. Interdependenz ist entscheidend: Die optimale SM-Allokation und das optimale Launch-Timing sind frequenzspezifisch; ein Schedule, der bei einer Frequenz optimal ist, kann bei einer anderen suboptimal sein.
3. Praktische Effizienz: Durch die Zerlegung des Problems in Partitionen und den Einsatz von MBO erreicht Kareus diese Gewinne, ohne die prohibitiven Kosten einer erschöpfenden Suche, was es zu einer praktischen Lösung für aktuelle und zukünftige großskalige KI-Trainings-Workloads macht.