Pipeline Parallelism: Daten flussaufwärts pumpen

In unserer vorherigen Untersuchung der Tensor Parallelism (TP) haben wir das Modell horizontal aufgeteilt und einzelne Gewichtsmatrizen auf mehrere GPUs verteilt. Obwohl dies effektiv ist, benötigt TP häufige Kollektivoperationen, weshalb schnelle Verbindungen oft hilfreich sind. TP kann auch über Knoten laufen, wenn Verbindung und Workload das rechtfertigen.

Um über einen einzelnen Node hinaus zu skalieren – also auf Cluster mit Hunderten oder Tausenden von GPUs – benötigen wir eine Technik, die weniger "gesprächig" ist und die geringere Bandbreite über Ethernet oder InfiniBand besser toleriert.

Hier kommt die Pipeline Parallelism (PP) ins Spiel. Anstatt die Schichten selbst zu teilen, schneidet PP das Modell vertikal und partitioniert die Schichten des neuronalen Netzes auf mehrere Geräte. Dadurch wird Ihr GPU-Cluster zu einer Fertigungsstraße, in der die Aktivierungen entlang der Kette weitergegeben werden.

Dieser Artikel analysiert die Hardware-Mechanik von PP, die erforderlichen Scheduling-Algorithmen, um das Silizium auszulasten, und die entscheidenden Kompromisse, die dabei eine Rolle spielen.

1. Die Hardware-Mechanik: Vertikales Schneiden

TP verteilt Operationen innerhalb einer Schicht; PP weist verschiedenen Stufen unterschiedliche Schichten zu. Beide verteilen Rechenarbeit und Parameterspeicher.

Die Einrichtung

Stellen Sie sich einen hypothetischen Transformer mit 32 Schichten vor, der auf vier GPUs (Stufen) verteilt ist.

GPU 0: Schichten 0–7
GPU 1: Schichten 8–15
GPU 2: Schichten 16–23
GPU 3: Schichten 24–31

Die Kommunikationsprimitive: Point-to-Point (P2P)

PP nutzt häufig Send/Recv zwischen Stufen. TP verwendet häufig Kollektivoperationen wie All-Reduce; All-Reduce kombiniert Tensoren und ist kein Broadcast.

Forward Pass: GPU 0 berechnet die Aktivierungen für Schicht 7 und sendet diese an GPU 1. GPU 1 empfängt sie und beginnt mit der Berechnung von Schicht 8.
Backward Pass: GPU 1 berechnet die Gradienten für Schicht 8 und sendet den Gradienten bezüglich des Eingangs zurück an GPU 0.

Hardware-Auswirkung: An einer einfachen Stufengrenze werden pro Micro-Batch ein Aktivierungstensor vorwärts und ein Aktivierungsgradient rückwärts übertragen. Ihre Größen hängen von Micro-Batch, Sequenzlänge, Hidden-Dimension und Datentyp ab. Die Kosten im Vergleich zu TP hängen von Workload und Topologie ab; PP kann über Knoten sinnvoll sein.

2. Das "Bubble"-Problem: Silicon Idle Time

Der grundlegende Fehler einer naiven Pipeline ist die sequentielle Abhängigkeit. Bei einem Micro-Batch und einfachem sequentiellem Ablauf warten nachfolgende Stufen auf Aktivierungen; der Rückwärtspfad läuft in umgekehrter Richtung. Die Wartezeit wird als Pipeline Bubble bezeichnet.

Um dem entgegenzuwirken, fügen wir Micro-Batches ein. Wir teilen den globalen Batch in $M$ kleine Stücke auf und schicken sie nacheinander durch die Pipeline.

Zeitplan A: GPipe (Alle-Vorwärts-Alle-Rückwärts)

Dies ist die einfachste Implementierung.

Füllen: Alle $M$ Micro-Batches werden durch den Vorwärtspfad geschickt.
Leeren: Sobald der letzte Micro-Batch das Ende erreicht, werden die Rückwärtspässe in umgekehrter Reihenfolge gestartet.

Die Hardware-Kosten: Speicher. Jede Stufe behält den Vorwärtszustand ihrer Micro-Batches bis zu deren Rückwärtspass, sofern Aktivierungen nicht neu berechnet werden. Der Aktivierungsspeicher kann daher mit der Micro-Batch-Anzahl wachsen.

Zeitplan B: 1F1B (Eins-Vorwärts-Eins-Rückwärts)

Dieser Zeitplan ist etwa in PyTorch und Megatron-LM implementiert. Anstatt auf alle Vorwärtspässe zu warten, wechselt eine GPU im Wesentlichen zwischen Vorwärts- und Rückwärtsaufgaben, sobald die Pipeline gefüllt ist (der "steady state").

Aufwärmen: GPUs führen genügend Vorwärtspässe aus, um die Pipeline zu füllen.
Steady State: 1 Vorwärts (neue Aktivierung erzeugen) $\rightarrow$ 1 Rückwärts (alte Aktivierung verbrauchen, Speicher freigeben).
Abkühlphase: Verarbeiten Sie die verbleibenden Backward-Passes.

Der Hardware-Vorteil: Durch das Ineinandergreifen der Operationen reduzieren wir den maximalen Speicherbedarf drastisch. Der Vorwärtszustand für Micro-Batch $k$ kann nach dessen Rückwärtspass freigegeben werden, statt alle Micro-Batches bis zum Ende des Trainingsschritts zu behalten.

3. Bare-Metal-Implementierung: Die 1F1B-Logik

Die Implementierung einer Pipeline erfordert die Verwaltung einer Zustandsmaschine auf jeder GPU. Der folgende Pseudocode zeigt die Reihenfolge auf einer mittleren Stufe. Er lässt Puffergrößen, passende Zeitpläne der Nachbarn, Backward-Buchhaltung, Loss-Berechnung und Synchronisation aus. Für ausführbaren Code ist eine unterstützte Pipeline-Runtime geeignet.

text

# Schematic synchronous 1F1B schedule for one intermediate stage.
# All stages must use a matching send/receive order; this is not runnable code.
for each warmup microbatch:
    activation = recv(previous_stage)
    output = local_forward(activation)
    send(output, next_stage)
    retain_backward_state(microbatch)
for each steady-state slot:
    forward_one_new_microbatch()          # receive, compute, send, retain
    gradient = recv(next_stage)           # oldest outstanding microbatch
    input_gradient = local_backward(gradient)
    send(input_gradient, previous_stage)
for each remaining outstanding microbatch:
    backward_one_microbatch()             # receive, compute, send
optimizer_step()                          # after all microbatches, all stages

Der Ablauf ist schematisch. Blockierende Sende- und Empfangsoperationen müssen mit den Nachbarstufen koordiniert werden.

4. Vorteile und Nachteile: Wann sollte man PP verwenden?

Vorteile

Inter-Node-Skalierung: Da P2P-Kommunikation leichtgewichtig ist (nur Aktivierungen an den Grenzen), kann PP zwischen Knoten eingesetzt werden, wenn Aktivierungstransfers und Pipeline-Bubbles zum Netzwerk und Workload passen.
Speichereffizienz: Durch das Aufteilen der Modellschichten werden Parameter-, Gradienten- und Optimizer-Zustände nach Stufen verteilt. Der Speicher pro GPU hängt von Schichtgrößen, Embeddings, geteilten Gewichten, Aktivierungen und der Balance der Aufteilung ab; vier Stufen bedeuten nicht automatisch je ein Viertel des Gesamtspeichers.

Nachteile

Die Bubble (Leerlaufzeit): Der größte Kostenfaktor. Während der Aufwärm- und Abkühlphase sind GPUs untätig. Wenn die Anzahl der Micro-Batches gering ist, sinkt die Effizienz drastisch. Für eine einfache, ausgeglichene Pipeline mit gleichen Vorwärts- und Rückwärtszeiten pro Stufe, ohne Kommunikationskosten, $M$ Micro-Batches und $N$ Stufen ergibt sich die idealisierte Auslastung $M/(M+N-1)$ . Ungleichgewicht und Kommunikation senken die tatsächliche Auslastung.
Veraltete Gewichte / Komplexität: Manche asynchronen Zeitpläne verwenden mehrere Gewichtsversionen. Ein synchroner 1F1B-Zeitplan mit Flush aktualisiert die Gewichte nach allen Micro-Batches eines Trainingsschritts; für die Gleichwertigkeit zum Training ohne Pipeline sind die Details entscheidend.
Statische Formen: Variable Formen erfordern die Übertragung von Form-Metadaten oder feste Puffer mit Padding; die Unterstützung hängt von der Runtime ab.

Zusammenfassung: Der 3D-Parallelismus-Stack

Pipeline Parallelism befindet sich in der Mitte der 3D-Skalierungshierarchie:

Tensor Parallelism: Innerhalb des Knotens (intra-layer). Schnelle Verbindungen sind hilfreich.
Pipeline Parallelism: Zwischen den Knoten (inter-layer). Kann je nach Workload und Netzwerk zwischen Knoten laufen.
Data Parallelism: Repliziert den gesamten TP+PP-Cluster, um die Batchgröße zu skalieren.

Durch die Kombination dieser Ansätze können wir Modelle mit Billionen von Parametern auf Tausenden von GPUs trainieren.