Technische Zusammenfassung: DisagMoE

1. Problemstellung

Groß angelegte Mixture-of-Experts (MoE)-Architekturen ermöglichen das Training von Large Language Models (LLMs) mit Billionen von Parametern, indem sie eine Teilmenge von Experten spärlich aktivieren. Das Training dieser Modelle stößt jedoch aufgrund der für die Expertenparallelität (EP) erforderlichen All-to-All-Kommunikation auf schwerwiegende Engpässe. Wenn Modelle über mehrere GPU-Knoten hinweg skaliert werden, verschärft die begrenzte Bandbreite des Knoten-zwischen-Netzwerks (z. B. InfiniBand) die Kommunikationslatenz, die bis zu 50 % der Zeit eines Trainingsschritts in Anspruch nehmen kann.

Bestehende Optimierungsstrategien konzentrieren sich hauptsächlich darauf, Kommunikation auf Operator-Ebene mit Berechnungen zu überlappen (z. B. durch Chunking von FFN-Berechnungen) oder Micro-Batch-Pipelining zu verwenden. Diese Ansätze leiden jedoch unter zwei grundlegenden Einschränkungen:

1. Restliche Stillstände: Inhärente Ungleichgewichte in den Verhältnissen von Berechnung zu Kommunikation zwischen Attention- und FFN-Schichten verhindern eine perfekte Überlappung. Die Attention-Berechnung skaliert quadratisch mit der Sequenzlänge, während FFN und Kommunikation linear skalieren. Dies erzeugt „Kommunikationsschwänze", die nicht verborgen werden können, insbesondere wenn die Sequenzlängen zunehmen.
2. Ressourcenineffizienz: Bisherige disaggregierte Ansätze (die Attention und FFN auf unterschiedliche Hardware trennen) erfordern oft, dass jede Gruppe die vollen Modellparameter hält, was bei Skalierung zu Out-of-Memory (OOM)-Fehlern führt, oder sie versagen darin, die Ressourcenallokation an die spezifischen Rechen-Kommunikations-Eigenschaften jeder Komponente anzupassen.

2. Methodik: DisagMoE

DisagMoE ist ein Trainingssystem, das entwickelt wurde, um All-to-All-Engpässe zu mindern, indem es durch eine disaggregierte Architektur die Platzierung und Planung des Modells gemeinsam optimiert.

2.1 Disaggregierte Komponentenplatzierung

Im Gegensatz zur traditionellen Expertenparallelität, bei der Attention- und FFN-Schichten auf derselben GPU-Gruppe koexistieren, partitioniert DisagMoE das Modell nach Komponententyp:

• Attention-Gruppen (A-Worker): Diese Gruppen enthalten Self-Attention-Schichten aus mehreren Transformer-Schichten. Sie replizieren dichte Komponenten und verteilen Trainingsdaten über Datenparallelität (DP).
• FFN-Gruppen (F-Worker): Diese Gruppen enthalten FFN-Schichten (Experten) aus mehreren Transformer-Schichten. Experten werden innerhalb der Gruppe über Expertenparallelität (EP) auf GPUs aufgeteilt (Sharding).
• Schichtzuweisung: Schichten werden in einer verschachtelten Weise den Gruppen zugewiesen (z. B. behandelt Gruppe 0 die Schichten 0, $p$, $2p$; Gruppe 1 behandelt die Schichten 1, $p+1$ usw.), was die Geräteauslastung unter Speicherbeschränkungen maximiert.

2.2 AF-Pipe: Mehrstufiges Pipeline mit Überlappung

DisagMoE führt AF-Pipe ein, einen Pipeline-Plan, der Kommunikation als eigenständige Stufe und nicht als Nebeneffekt behandelt.

• Viele-zu-Viele (M2N)-Kommunikation: Anstelle von Standard-All-to-All-Austauschen innerhalb einer einzelnen Knotengruppe etabliert AF-Pipe unidirektionale, viele-zu-viele-Token-Austausche zwischen den Attention- und FFN-Gruppen.
• Stufenabstimmung: Die Pipeline richtet Grenzen über die Gruppen hinweg aus und fusioniert die traditionellen „Dispatch"- und „Combine"-Operationen zu einer einzigen einheitlichen Kommunikationsstufe. Dies reduziert den Kommunikationsaufwand um etwa $1/k$ (wobei $k$ die Anzahl der Top-$k$-Experten ist) und eliminiert redundante Datentransfers.
• Asynchrone Ausführung: Das System nutzt drei koordinierte Streams (Vorwärts, Rückwärts und Kommunikation) innerhalb jedes Workers. Dies ermöglicht es der FFN-Gruppe, Gradienten (N2M) zu senden, während sie gleichzeitig den Vorwärtsdurchlauf auf verfügbaren Daten berechnet, was eine blockierungsfreie Ausführung und eine nahezu kontinuierliche GPU-Auslastung sicherstellt.

2.3 Adaptive Worker-Allokation

Um das Ungleichgewicht zwischen Attention (rechengebunden) und FFN (kommunikationsgebunden) zu adressieren, setzt DisagMoE einen adaptiven Allokator ein, der von einem Rechen-Kommunikations-Roofline-Modell geleitet wird.

• Modellierung: Das System modelliert den Wendepunkt ($\hat{I}$), an dem eine Komponente von netzwerkgebunden zu rechengebunden wechselt. Attention-Schichten erreichen mit zunehmender Sequenzlänge schneller die Rechengrenze als FFN-Schichten.
• Optimierung: Das System formuliert die Partitionierung von GPUs und Network Interface Cards (NICs) als gemischt-ganzzahliges lineares Programmierungsproblem (MILP). Es weist FFN-Gruppen dynamisch mehr Netzwerkbandbreite zu (um ihre Kommunikationsgrenze anzuheben) und Attention-Gruppen mehr Rechenressourcen, um die Pipeline-Stufen auszugleichen und Blasen zu minimieren.
• Suchstrategie: Ein Algorithmus verwendet das Roofline-Modell, um eine Startkonfiguration zu generieren, die dann durch lokale Suche unter Verwendung gemessener Zeiten pro Schritt verfeinert wird, um nicht modellierte Overheads zu berücksichtigen.

3. Hauptbeiträge

1. Disaggregierte Komponentenplatzierung: Eine neuartige Architektur, die Attention- und FFN-Schichten in disjunkte Worker-Gruppen trennt und damit unabhängiges Skalieren und Ressourcenallokation ermöglicht.
2. AF-Pipe-Plan: Eine mehrstufige Pipeline, die konventionelles All-to-All durch M2N-Austausche zwischen Gruppen ersetzt und Kommunikation systematisch sowohl mit Attention- als auch mit FFN-Berechnungen überlappt.
3. Adaptive Ressourcenallokation: Ein roofline-gesteuerter Mechanismus, der die Größe von GPU- und NIC-Ressourcen für Attention- und FFN-Gruppen gemeinsam bestimmt und das Verhältnis von Berechnung zu Kommunikation über verschiedene Sequenzlängen und Modellkonfigurationen hinweg effektiv ausgleicht.

4. Experimentelle Ergebnisse

Die Autoren bewerteten DisagMoE auf Clustern mit 16 Knoten (128× H800 GPUs) unter Verwendung von drei MoE-Modellen: DeepSeek-MoE, GPT-OSS-120B und Qwen3, mit Sequenzlängen im Bereich von 4K bis 32K.

• Durchsatzverbesserung: DisagMoE übertraf konsistent starke Baselines, einschließlich Megatron-LM (1F1B verschachtelt), Tutel, Comet und DualPipe.
  • Erzielte bis zu 1,81× Beschleunigung gegenüber Megatron-1F1B.
  • Erzielte bis zu 1,34× Beschleunigung gegenüber state-of-the-art MoE-Überlappungssystemen.
• Reduktion der Kommunikation: Im Vergleich zu Baselines reduzierte DisagMoE die nicht überlappte Kommunikationszeit um bis zu 88 % (vs. Tutel), 75 % (vs. Comet) und 45 % (vs. DualPipe).
• Skalierbarkeit: Das System zeigte sich bei verschiedenen Top-$k$-Werten und Größen der Expertenparallelität (EP) als effektiv.
• Ressourcensensitivität: Experimente bestätigten, dass sich die optimalen GPU-Verhältnisse zwischen Attention- und FFN-Gruppen mit der Sequenzlänge ändern (z. B. begünstigen längere Sequenzen mehr Attention-Ressourcen), was die Notwendigkeit einer adaptiven Allokation untermauert.

5. Bedeutung und Behauptungen

Die Arbeit behauptet, dass DisagMoE eine grundlegende Einschränkung im aktuellen MoE-Training adressiert: die Unfähigkeit, Kommunikation aufgrund des strukturellen Ungleichgewichts zwischen Attention- und FFN-Schichten perfekt zu überlappen. Durch Disaggregation dieser Komponenten und adaptive Allokation von Ressourcen basierend auf einem Roofline-Modell erweitert das System das effektive Überlappungsfenster über das hinaus, was eine symmetrische Bandbreitenbereitstellung erreichen kann.

Die Autoren betonen, dass ihr Ansatz insbesondere für groß angelegte Pre-Training-Workloads mit festen Sequenzlängen von erheblicher Bedeutung ist. Sie weisen darauf hin, dass ihr aktuelles Design zwar stabile Formen (feste Sequenzlänge und Micro-Batch) voraussetzt, die Kern-Einsicht jedoch – dass zukünftige Hardware und Software günstige Rechen-Kommunikations-Verhältnisse über Komponenten hinweg anvisieren sollten, anstatt einfach nur die Rohbandbreite zu skalieren – eine neue Richtung für die effiziente Skalierung des MoE-Trainings bietet. Das System ist auf Megatron-LM implementiert und erweist sich als praktikable Lösung für das Training von Modellen mit hunderten Milliarden Parametern ohne die Speicherbeschränkungen, die frühere disaggregierte Trainingsversuche plagten.