Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

1. Problemstellung

Das Training und die Inferenz großer Sprachmodelle (LLMs) stoßen auf zwei primäre Speicherengpässe:

1. Parameterspeicher: Der Speicherbedarf für Modellgewichte, Gradienten und Optimiererzustände.
2. Aktivierungsspeicher: Der Speicherbedarf für Zwischenaktivierungen, der sich in Aufmerksamkeitsmechanismen quadratisch mit der Sequenzlänge ($S^2$) skaliert.

Bestehende Parallelisierungsstrategien adressieren diese Probleme isoliert oder mit Kompromissen:

• Tensor-Parallelismus (TP): Teilt Modellgewichte über Geräte auf, reduziert den Parameterspeicher, lässt den Aktivierungsspeicher jedoch vollständig repliziert. Bei langen Sequenzlängen wird es aufgrund des hohen Aktivierungsspeichers und des Kommunikationsaufkommens ineffizient.
• Sequenz-Parallelismus (SP): Teilt die Eingabesequenz über Geräte auf, reduziert den Aktivierungsspeicher, lässt die Modellgewichte jedoch vollständig repliziert. Er ist wirkungslos, wenn der Parameterspeicher der Engpass ist.
• Konventioneller TP+SP: Kombiniert beides, indem sie orthogonalen Achsen in einem Device-Mesh zugewiesen werden (z. B. ein $T \times \Sigma$-Gitter). Obwohl effektiv, verbraucht dies eine große Anzahl von Geräten für den Modellparallelismus ($T \times \Sigma$), wodurch weniger Geräte für den Datenparallelismus übrig bleiben. Darüber hinaus zwingt dies Parallelgruppen oft dazu, langsamere interknotenverbindungen zu überbrücken, was hochbandbreitige intraknotenverbindungen (wie NVLink oder AMD Infinity Fabric) unterauslastet.

Die Kernherausforderung: Wie lässt sich sowohl der Parameter- als auch der Aktivierungsspeicher auf einer einzigen Device-Achse gleichzeitig reduzieren, um die Speichereffizienz zu maximieren und Geräte für den Datenparallelismus freizugeben, ohne prohibitiv hohe Kommunikationskosten zu verursachen?

2. Methodik: Tensor- und Sequenz-Parallelismus (TSP)

Die Autoren schlagen TSP vor, eine Strategie, die Tensor- und Sequenz-Parallelismus auf eine einzelne Device-Achse „faltet". Anstatt TP und SP orthogonalen Dimensionen zuzuweisen, weist TSP jedem Rang (GPU) sowohl einen Shard der Modellgewichte als auch einen Shard der Sequenz zu.

Schlüsselkomponenten des Designs:

• Vereinheitlichtes Sharding: Jedes Gerät hält $1/D$ der Modellparameter und $1/D$ der Sequenz-Token. Dies reduziert sowohl den gewichtsproportionalen Speicher als auch den Aktivierungsspeicher um einen Faktor $D$ auf einer einzigen Achse.
• Laufzeitpläne: Um die Tatsache zu handhaben, dass Gewichte und Aktivierungen unterschiedlich geschert sind, setzt TSP spezifische Kommunikationspläne für verschiedene Schichten ein:
  • Aufmerksamkeits-Schichten: Verwendet einen Broadcast-Plan.
    • Rangs iterieren über Gewichts-Shards (verpackte $W_Q, W_K, W_V, W_O$).
    • In jeder Iteration broadcastet ein Rang seinen Gewichts-Shard an alle anderen.
    • Alle Rangs berechnen lokale Projektionen auf ihren Sequenz-Shards.
    • Eine All-Gather-Operation rekonstruiert den vollständigen Key/Value-Kontext über die Sequenzdimension hinweg, um FlashAttention auszuführen.
    • Ein Zickzack-Partitionierungsschema wird verwendet, um die Arbeitslast der kausalen Aufmerksamkeit auszugleichen (sicherstellend, dass spätere Token nicht unverhältnismäßig mehr Arbeit verursachen).
  • Gated-MLP-Schichten: Verwendet einen Ring-Kommunikations-Plan.
    • Gewichts-Shards ($W_1, W_2, W_3$) zirkulieren Punkt-zu-Punkt um den Ring.
    • Jeder Rang akkumuliert lokale Teilergebnisse.
    • Dies eliminiert die Notwendigkeit einer All-Reduce-Operation (standardmäßig in TP), da die Sequenz lokal bleibt und nur Gewichte bewegt werden.
• Hardware-Bewusstsein: Das Design ist für Topologien wie AMDs MI300X (Infinity Fabric) optimiert, bei denen die kollektive Bandbreite mit der Anzahl der beteiligten GPUs skaliert. TSP stellt sicher, dass Kommunikationsprimitive (Broadcast und Ring) die gesamte Gruppe einbeziehen oder mit Berechnungen überlappt werden, um die Effizienz zu maximieren.

3. Hauptbeiträge

1. Neuartige Parallelisierungsstrategie: Einführung von TSP, das zwei orthogonale Parallelismusachsen in eine zusammenfasst und eine gleichzeitige Reduzierung von Parameter- und Aktivierungsspeicher ermöglicht.
2. Theoretische Analyse: Ein rigoroses mathematisches Modell, das TSP mit DP, TP, SP und TP+SP vergleicht hinsichtlich:
   • Speicher: TSP ist das einzige Schema, das sowohl gewichtsproportionalen als auch Aktivierungsspeicher um $1/D$ reduziert, ohne $D^2$ Geräte zu benötigen.
   • Kommunikation: Obwohl TSP gewichtsbewegende Kommunikation einführt (die TP+SP vermeidet), beweisen die Autoren, dass bei langen Sequenzen oder moderaten Batch-Größen die aktivierungsproportionale Kommunikation dominiert, was TSP konkurrenzfähig macht.
   • FLOPs: TSP behält die gleiche pro-Gerät-Berechnungseffizienz wie TP und SP bei.
3. Implementierung & Überlappung: Detaillierte Implementierung von TSP-Aufmerksamkeits- und MLP-Blöcken mit expliziten Strategien zur Überlappung von Kommunikation und Berechnung (Pipelining von Gewichtsübertragungen hinter GEMM-Berechnungen).
4. Topologie-Optimierung: Demonstration, dass TSP es ermöglicht, die gesamte Modell-Replik innerhalb eines einzelnen Hochbandbreiten-Knotens (intraknoten) unterzubringen, wodurch die Leistungsnachteile vermieden werden, die oft durch das Überbrücken langsamerer interknotenverbindungen bei 2D TP+SP-Meshes entstehen.

4. Ergebnisse

Die Autoren bewerteten TSP auf einem Cluster mit bis zu 1.024 AMD MI300X GPUs (128 Knoten) mit einem dichten Transformer-Modell mit 7 Milliarden Parametern.

• Speichereffizienz:
  • TSP erreichte konsistent den niedrigsten Spitzenspeicher über alle getesteten Sequenzlängen (16K bis 128K Token).
  • Bei kurzen Kontexten entsprach TSP TP (dominiert durch Parameterspeicher).
  • Bei langen Kontexten übertraf TSP TP und TP+SP erheblich (dominiert durch Aktivierungsspeicher) und verhinderte Out-Of-Memory (OOM)-Fehler, bei denen andere Strategien versagten.
• Durchsatz:
  • Trotz eines höheren theoretischen Kommunikationsaufkommens (aufgrund von Gewichtsverschiebungen) behielt TSP einen konkurrenzfähigen Durchsatz mit TP+SP bei.
  • Die zusätzliche Kommunikation wurde erfolgreich durch Pipelining hinter Berechnungen versteckt.
  • Durchsatzvorteile wuchsen mit steigenden Parallelismusgraden ($D$) und Sequenzlängen.
• Skalierbarkeit:
  • TSP ermöglichte größere Mikro-Batches aufgrund des reduzierten Speicherfußabdrucks, was den Durchsatz weiter verbesserte.
  • Es ermöglichte, Modell-Repliken vollständig innerhalb eines einzelnen Knotens (8 GPUs) unterzubringen, wodurch Ränge für Datenparallelismus (DP) in größeren Clustern freigegeben wurden.

5. Bedeutung

• Hardware-bewusste Skalierung: TSP bietet eine praktische Lösung für moderne HPC-Cluster, bei denen die intraknoten-Bandbreite hoch ist, die interknoten-Bandbreite jedoch ein Engpass darstellt. Es maximiert die Nutzung teurer intraknotenverbindungen.
• Tauglichkeit für lange Kontexte: Da sich Modelle zu längeren Kontextfenstern (100k+ Token) bewegen, wird der Aktivierungsspeicher zum primären Engpass. TSP ist einzigartig positioniert, dies ohne den massiven Geräteaufwand von 2D-Parallelismus zu bewältigen.
• Flexibilität: TSP kann orthogonal mit anderen Parallelisierungsschemata (Pipeline-Parallelismus, Expert-Parallelismus, Datenparallelismus) kombiniert werden und bietet Systemdesignern eine flexible „Achse", um spezifische Hardware-Topologien zu optimieren.
• Kosteneffizienz: Durch die Reduzierung der für den Modellparallelismus benötigten Geräte ermöglicht TSP die Nutzung mehrerer Geräte für den Datenparallelismus, was potenziell die für einen bestimmten Trainingslauf erforderliche gesamte Clustergröße verringert oder die effektive Batch-Größe erhöht.

Zusammenfassend stellt TSP eine bedeutende Evolution in verteilten Trainingsstrategien dar, die eine kontrollierbare Erhöhung des Kommunikationsaufkommens gegen erhebliche Gewinne in der Speichereffizienz und Hardwareauslastung tauscht, wodurch es zu einem entscheidenden Werkzeug für das Training von LLMs der nächsten Generation mit langen Kontexten wird.