Datenstrom nach oben pumpen: Eine hardwarezentrierte Tiefenanalyse der Pipeline-Parallelität

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, stößt TP an eine harte physikalische Grenze: die Lichtgeschwindigkeit. Es erfordert so häufige Synchronisation (zweimal pro Transformer-Block), dass es praktisch auf den ultraschnellen NVLink-Bereich eines einzelnen Compute-Nodes beschränkt ist.

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

Während sich Tensor Parallelism auf compute-bound Skalierung (Aufteilung der Berechnungen) bezieht, geht es bei Pipeline Parallelism um memory-bound Skalierung (Aufteilung des Zustands).

Die Einrichtung

Stellen Sie sich ein Llama-3-Modell mit 32 Schichten vor, das auf 4 GPUs (Stages) 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)

Im Gegensatz zu TP, das auf aufwändige All-Reduce-Operationen (Broadcast von Daten an alle) angewiesen ist, basiert PP auf Send/Recv (Point-to-Point)-Operationen.

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: Das übertragene Datenvolumen entspricht lediglich der Größe der Aktivierungen (Batch $\times$ Sequence $\times$ Hidden), nicht den Modellgewichten oder vollständigen Gradienten. Dies ist deutlich geringer als bei TP, wodurch PP über Knoten hinweg, die durch Standard-Netzwerkverbindungen verbunden sind, praktikabel wird.

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

Der grundlegende Fehler einer naiven Pipeline ist die sequentielle Abhängigkeit. Wenn GPU 0 rechnet, wartet GPU 1 auf Daten. Bei einem Durchlauf mit nur einem Batch ist jeweils nur eine GPU aktiv. 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. GPU 0 muss die Aktivierungen (für den Rückwärtspfad) für jeden einzelnen Micro-Batch zwischenspeichern, während sie darauf wartet, dass das Signal das Ende der Pipeline erreicht und zurückkommt. Dies führt zu enormen VRAM-Spitzen.

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

Dies ist der Industriestandard (verwendet in DeepSpeed und Megatron-LM). 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. Wir geben den Aktivierungsspeicher für Micro-Batch $k$ sofort nach Abschluss des Backward-Passes frei, anstatt ihn für die gesamte Dauer der Epoche zu halten.

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

Die Implementierung einer Pipeline erfordert die Verwaltung einer Zustandsmaschine auf jeder GPU. Hier ist eine konzeptionelle PyTorch-Implementierung eines einzelnen 1F1B-Schritts für einen intermediären GPU-Knoten (nicht der erste oder letzte).

import torch
import torch.distributed as dist
 
def train_step_1f1b(model, micro_batches, my_rank, prev_rank, next_rank):
    # Setup communication buffers
    # We need to know the shape of activations ahead of time (static shapes preferred)
    activations_recv = torch.zeros(SHAPE, device='cuda')
    grads_recv = torch.zeros(SHAPE, device='cuda')
    
    # 1. Warmup Phase: Fill the pipe
    # Calculate how many FWD passes we do before the first BWD arrives
    num_warmup = get_warmup_steps(my_rank, len(micro_batches))
    
    fwd_cache = [] # We must stash activations for our own BWD pass later
 
    for i in range(num_warmup):
        # Recv activation from previous stage
        dist.recv(activations_recv, src=prev_rank)
        
        # Compute Forward
        # We detach to allow autograd graph to be separated per microbatch
        inputs = activations_recv.clone().requires_grad_(True)
        output = model(inputs)
        
        # Send activation to next stage
        dist.send(output.data, dst=next_rank)
        
        # Cache for backward
        fwd_cache.append((inputs, output))
 
    # 2. Steady State: 1 Forward, 1 Backward
    remaining_fwd = len(micro_batches) - num_warmup
    
    for i in range(remaining_fwd):
        # --- Forward Step ---
        dist.recv(activations_recv, src=prev_rank)
        inputs = activations_recv.clone().requires_grad_(True)
        output = model(inputs)
        dist.send(output.data, dst=next_rank)
        fwd_cache.append((inputs, output))
        
        # --- Backward Step ---
        # Recv gradients from next stage
        dist.recv(grads_recv, src=next_rank)
        
        # Pop the oldest microbatch from cache
        my_inputs, my_outputs = fwd_cache.pop(0)
        
        # Run local backward
        # This computes gradients for model weights AND input activations
        torch.autograd.backward(
            tensors=[my_outputs], 
            grad_tensors=[grads_recv]
        )
        
        # Send input gradients back to previous stage
        dist.send(my_inputs.grad, dst=prev_rank)
 
    # 3. Cooldown: Finish remaining backwards
    while fwd_cache:
        dist.recv(grads_recv, src=next_rank)
        my_inputs, my_outputs = fwd_cache.pop(0)
        torch.autograd.backward([my_outputs], [grads_recv])
        dist.send(my_inputs.grad, dst=prev_rank)

Hinweis: Dieser Code vereinfacht die Fehlerbehandlung und Optimierung (wie die Verwendung von isend/irecv für asynchrone Überlappung), veranschaulicht jedoch den grundlegenden Datenfluss.

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

Vorteile

Inter-Node-Skalierung: Da P2P-Kommunikation leichtgewichtig ist (nur Aktivierungen an den Grenzen), funktioniert PP auch über langsamere Netzwerke (Ethernet) gut. Es ist die Standardmethode, um ein Modell über mehrere Serverknoten zu skalieren.
Speichereffizienz: Durch das Aufteilen der Modellschichten teilen Sie die Anforderungen an Parameter-, Gradienten- und Optimizer-State-Speicher perfekt gleichmäßig auf. Ein 100GB Modell auf 4 GPUs benötigt nur 25GB pro GPU (zuzüglich Aktivierungs-Overhead).

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. Formel für die Effizienz: $\frac{M}{M + N - 1}$ wobei $M$ die Micro-Batches und $N$ die Pipeline-Stufen sind. Sie benötigen viele Micro-Batches (hohe globale Batchgröße), um die Bubble zu verstecken.
Veraltete Gewichte / Komplexität: Einige fortgeschrittene Zeitpläne (wie PipeDream) verwenden veraltete Gewichte, um Bubbles zu vermeiden, was die Konvergenz beeinflussen kann. 1F1B erhält die mathematische Korrektheit, erfordert jedoch komplexes Zustandsmanagement.
Statische Formen: PP-Implementierungen sind in der Regel problematisch bei dynamischem Kontrollfluss oder variablen Sequenzlängen, da P2P-Kommunikation feste Puffergrößen erwartet.

Zusammenfassung: Der 3D-Parallelismus-Stack

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

Tensor Parallelism: Innerhalb des Knotens (intra-layer). Schnelles NVLink erforderlich.
Pipeline Parallelism: Zwischen den Knoten (inter-layer). Langsameres Netzwerk wird toleriert.
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.