Technisches Resümee: DOPPLER – Dual-Policy Learning für die Gerätezuweisung in asynchronen Datenflussgraphen

1. Problemstellung

Das Paper adressiert das Problem der Gerätezuweisung (Device Assignment) in Multi-GPU-Systemen, die komplexe Machine-Learning-Workloads (z. B. Transformer, Mixture-of-Experts) ausführen, welche als statische Datenflussgraphen dargestellt werden. Das Ziel besteht darin, Operationen (Knoten) auf Geräte abzubilden, um die Gesamtausführungszeit in einem work-conserving (WC) System zu minimieren.

Im Gegensatz zu herkömmlichen Bulk-Synchronous-Frameworks (z. B. Standard-PyTorch oder TensorFlow), bei denen alle Geräte auf die langsamste Operation warten, bevor sie fortfahren können, planen WC-Systeme Operationen asynchron, sobald deren Datenabhängigkeiten erfüllt sind. Während WC-Systeme ein erhebliches Potenzung für Beschleunigungen durch die Überlappung von Kommunikation und Berechnung bieten, führen sie einzigartige Herausforderungen mit sich:

• Stochastische Ausführungsreihenfolge: Oh[ne globale Synchronisation ist die Reihenfolge der Ausführung dynamisch und reagiert empfindlich auf Hardware-Heterogenität und Ressourcenkonflikte.
• Komplexe Optimierung: Die Ausführungszeit lässt sich aufgrund dieser Stochastik nicht als geschlossene Funktion der Zuweisung ausdrücken.
• Widersprüchliche Ziele: Eine effektive Zuweisung muss gleichzeitig die GPU-Last ausbalancieren und die Inter-GPU-Kommunikation minimieren – eine Aufgabe, die dadurch erschwert wird, dass die Lastverteilung in asynchronen Umgebungen inhärent temporal ist.

Bisherige lernbasierte Ansätze sind durch ihre Abhängigkeit von Bulk-Synchronous-Annahmen, die Verwendung einer einzelnen Platzierungsstrategie ohne Modellierung der Systemdynamik und den Mangel an Optimierung während des tatsächlichen Einsatzes begrenzt.

2. Methodik: Das DOPPLER-Framework

Die Autoren schlagen DOPPLER vor, ein Reinforcement Learning (RL)-Framework, das ein Dual-Policy-sequentielles Entscheidungsverfahren und eine dreistufige Trainingspipeline verwendet.

2.1 Dual-Policy-Formulierung

Anstatt eine einzelne monolithische Policy zu lernen, zerlegt DOPPLER den Zuweisungsprozess in zwei distinkte Policies, die als episodischer Markov-Entscheidungsprozess (MDP) modelliert werden:

1. SEL (Selection Policy): Wählt die nächste Operation (Vertex) zur Zuweisung aus der Menge der Kandidatenknoten aus. Diese Policy durchläuft den teilweise zugewiesenen Graphen, um den nicht-deterministischen Fluss der „Zeit“ zu approximieren, und lernt effektiv eine Ausführungsreihenfolge.
2. PLC (Placement Policy): Bestimmt das spezifische GPU-Gerät für die ausgewählte Operation mit dem Ziel, die Last zu balancieren und die Kommunikation zu minimieren.

Architektur:

• Beide Policies nutzen Graph Neural Networks (GNNs) mit Message Passing, um die Struktur des Datenflussgraphen, die Abhängigkeiten und die kritischen Pfade zu kodieren.
• SEL aggregiert Informationen aus kritischen Pfaden (Bottom-Level- und Top-Level-Pfade) sowie Knotenmerkmale, um den nächsten Knoten auszuwählen.
• PLC kombiniert Knotenrepräsentationen, Gerätemerkmale (z. B. aktuelle Last, Endzeit) und den Satz bereits platzierter Knoten, um das optimale Gerät auszuwählen.
• Effizienzoptimierung: Um die prohibitiven Trainingskosten zu vermeiden, die mit dem erneuten Durchlaufen von Message Passing bei jedem Schritt verbunden sind, führt DOPPLER das Message Passing einmal pro Episode durch und aktualisiert die Gerätemerkmale dynamisch, ohne die Graph-Embeddings bei jedem Schritt neu zu berechnen.

2.2 Dreistufige Trainingspipeline

DOPPLER nutzt ein „Learning-by-Doing“-Paradigma über drei Stufen, um Konvergenz und Adaptivität zu gewährleisten:

1. Stufe I (Imitation Learning / Offline): Die Policies werden mittels Supervised Learning vortrainiert, um eine heuristische Lehrkraft (speziell die Critical Path Heuristik) nachzuahmen. Dies beschleunigt die Konvergenz durch Bereitstellung einer starken initialen Policy, bevor die Exploration beginnt.
2. Stufe II (Simulationsbasiertes RL / Offline): Die Policies werden mittels Reinforcement Learning innerhalb eines Software-Simulators (der die WC-Ausführungslogik implementiert) verfeinert. Die Belohnung (Reward) leitet sich aus der simulierten Ausführungszeit ab. Diese Stufe ermöglicht eine schnelle Iteration ohne die Latenz echter Hardware.
3. Stufe III (Real-System RL / Online): Die trainierten Policies werden in einem echten WC-System eingesetzt. Sie werden kontinuierlich über Online-RL unter Verwendung tatsächlicher Ausführungszeiten als Belohnung verfeinert. Diese Stufe passt die Policy an spezifische Hardwarecharakteristika (z. B. NVLink-Jitter, Kontention) an, die ein Simulator nicht perfekt erfassen kann, ohne zusätzlichen Overhead für den Nutzer zu erzeugen.

3. Zentrale Beiträge

1. Problemformulierung: Das Paper untersucht die Gerätezuweisung speziell im Kontext von asynchronen, work-conserving Multi-GPU-Systemen und adressiert die stochastische Natur der Ausführungsreihenfolge.
2. Dual-Policy-Ansatz: Es führt eine Trennung der Zuständigkeiten ein, bei der eine Policy die ungefähre Traversierungsreihenfolge der Knoten lernt und die andere die Geräteplatzierung übernimmt, wodurch sowohl die Ausführungsdynamik als als auch die Hardwarebeschränkungen erfasst werden.
3. Dreistufiges Framework: DOPPLER schlägt eine hybride Trainingsstrategie vor, die Imitation Learning, simulationsbasiertes RL und Online-RL im realen System kombiniert. Dieser Ansatz beschleunigt die Konvergenz und ermöglicht eine kontinuierliche Optimierung während des Betriebs.
4. Skalierbare GNN-Implementierung: Die Autoren schlagen eine effiziente Message-Passing-Approximation vor, die die Trainingszeit signifikant reduziert, ohne die Konvergenz zu beeinträchtigen.

4. Experimentelle Ergebnisse

Die Experimente wurden auf vier NVIDIA Tesla P100 GPUs mit verschiedenen Workloads durchgeführt: Chain Matrix Multiplications (CHAINMM), Feedforward Neural Networks (FFNN) und Llama Transformer-Blöcke/Layer.

• Performance vs. Baselines: DOPPLER (insbesondere das voll trainierte DOPPLER-SYS) übertraf konsistent die Baselines, einschließlich der Critical Path Heuristik, PLACETO (Single-Policy RL), GDP (Graph Embedding RL) und eines benutzerdefinierten enumerativen Optimierers.
  • DOPPLER erreichte bis zu 52,7 % geringere Ausführungszeiten im Vergleich zur besten Baseline (GDP).
  • Es reduzierte die Laufzeit um bis zu 78,2 % im Vergleich zur Critical Path Heuristik.
  • Es übertraf den benutzerdefinierten enumerativen Optimierer um bis zu 13,8 %.
• Ablationsstudien:
  • Dual Policy: Die Kombination aus SEL- und PLC-Policies lieferte die besten Ergebnisse, insbesondere für komplexe Modelle, was die Notwendigkeit unterstreicht, sowohl die Reihenfolge als auch die Platzierung zu lernen.
  • Trainingsstufen: Ein Training allein auf dem realen System führte zu langsamer Konvergenz und Instabilität. Die Kombination aus Imitation Learning (Stufe I) und Simulation (Stufe II) war entscheidend für eine stabile und schnelle Konvergenz vor der Online-Verfeinerung (Stufe III).
• Transfer Learning: DOPPLER zeigte eine starke Generalisierung. Auf einfachen Graphen trainierte Policies (FFNN) konnten mit wenigen Shot-Episoden (2k–4k) so feinjustiert werden, dass sie bei ungesehenen Architekturen (Llama) und unterschiedlichen Hardwarekonfigurationen (z. B. Transfer von 4 P100s auf 8 V100s) gut abschnitten.
• Skalierbarkeit: Die Trainings- und Inferenzzeiten von DOPPLER skalieren linear mit der Größe des Datenflussgraphen und sind somit effizienter als andere RL-basierte Baselines.

5. Bedeutung und Ansprüche

Das Paper behauptet, dass DOPPLER einen bedeutenden Schritt zur Optimierung moderner ML-Workloads darstellt, indem es über statische, Bulk-Synchronous-Annahmen hinausgeht. Die Bedeutung liegt in:

• Praktischer Effizienz: Durch die Ermöglichung der work-conserving Ausführung kann DOPPLER erhebliche GPU-Stunden in groß angelegten Implementierungen einsparen (z. B. die Einsparung von Millionen GPU-Stunden jährlich für ChatGPT-Skala-Workloads).
• Adaptivität: Im Gegensatz zu statischen Heuristiken oder rein simulationsbasierten Modellen ermöglicht die Online-Lernfähigkeit von DOPPLER die Anpassung an die spezifischen, verrauschten Realitäten von Produktions-Hardwareumgebungen.
• Methodische Innovation: Die Dual-Policy-Dekomposition bietet einen nuancierteren Ansatz für das kombinatorische Problem der Gerätezuweisung, indem sie das „Was als Nächstes auszuführen ist“ vom „Wo es auszuführen ist“ trennt, was besser mit der temporalen Natur der asynchronen Ausführung korreliert.

Die Autoren betonen, dass ihr Ansatz keine Änderungen an den zugrunde liegenden ML-Frameworks erfordert, sondern vielmehr die Scheduling-Schicht optimiert, was ihn zu einer praktischen Lösung für bestehende Multi-GPU-Infrastrukturen macht.