Pipeline-Parallelismus ist eine verteilte Trainingsstrategie für tiefe neuronale Netze, die das Modell in sequenzielle Stufen (Stages) aufteilt, die auf verschiedenen Geräten (z. B. GPUs) ausgeführt werden. Sie kombiniert Aspekte der Daten- und Modellparallelität, um den Durchsatz zu erhöhen und den Speicherbedarf pro Gerät zu reduzieren.
Hintergrund und Motivation
Das Training großer neuronaler Netze, insbesondere großer Sprachmodelle und anderer Deep-Learning-Modelle, erfordert enorme Rechenressourcen. Mit zunehmender Modellgröße in die Milliarden oder Billionen von Parametern wird das Training auf einem einzelnen Gerät aufgrund von Speicher- und Rechenlimits unpraktikabel. Datenparallelität repliziert das gesamte Modell auf jedem Gerät und partitioniert die Daten, scheitert jedoch, wenn das Modell nicht in den Speicher eines einzelnen Geräts passt. Modellparallelität partitioniert das Modell selbst über Geräte, führt jedoch oft zu Lastungleichgewichten und hohem Kommunikationsaufwand. Pipeline-Parallelismus adressiert diese Probleme, indem er das Modell in aufeinanderfolgende Stufen unterteilt, die jeweils auf einem anderen Gerät ausgeführt werden und Daten in einer Pipeline-Art weiterreichen.
Kernkonzepte
Die grundlegende Arbeitseinheit im Pipeline-Parallelismus ist der Mikro-Batch. Ein großer Trainings-Batch wird in kleinere Mikro-Batches aufgeteilt, die sequenziell durch die Pipeline-Stufen fließen. Jede Stufe führt Vorwärts- und Rückwärtsdurchläufe auf ihren zugewiesenen Schichten aus. Die Pipeline wird "gefüllt", indem Mikro-Batches nacheinander eingespeist werden. Nach einer anfänglichen Aufwärmphase können alle Stufen gleichzeitig arbeiten, was den Durchsatz erhöht.
Die Effizienz einer Pipeline wird oft durch die Blasenzeit (Bubble) gemessen, also den Anteil der Zeit, in der Stufen untätig sind. Für eine Pipeline mit p Stufen und m Mikro-Batches wird der ideale Durchsatz erreicht, wenn m viel größer als p ist, wodurch der relative Blasen-Overhead minimiert wird. Eine größere Anzahl von Mikro-Batches erhöht jedoch auch den Speicherbedarf für Aktivierungen, was einen Kompromiss zwischen Durchsatz und Speichernutzung darstellt.
Wichtige Schemata
Es wurden verschiedene Pipeline-Parallelismus-Schemata entwickelt, die sich in der Planung von Vorwärts- und Rückwärtsdurchläufen sowie im Umgang mit Gradientenaktualisierungen unterscheiden.
GPipe
Google führte GPipe im Jahr 2019 ein. Es partitioniert das Modell in Stufen und verwendet einen einfachen Plan: Jeder Mikro-Batch durchläuft alle Stufen in Vorwärtsrichtung und dann in Rückwärtsrichtung. GPipe verwendet synchrone Gradientenaktualisierungen, d. h., alle Stufen warten, bis alle Mikro-Batches verarbeitet sind, bevor die Gewichte aktualisiert werden. Dies gewährleistet konsistente Gradienten, erzeugt jedoch eine Blase, die mit der Anzahl der Stufen wächst. GPipe speichert entweder Aktivierungen für den Rückwärtsdurchlauf oder berechnet sie neu, was einen Kompromiss zwischen Speicher und Rechenzeit darstellt.
PipeDream
PipeDream, vorgestellt von Microsoft Research im Jahr 2019, verwendet einen asynchronen Plan. Jede Stufe verarbeitet Mikro-Batches in einer Round-Robin-Reihenfolge und führt Vorwärts- und Rückwärtsdurchläufe so schnell wie möglich aus, sobald die Eingaben verfügbar sind. Dies reduziert die Blase erheblich, führt jedoch zu "gewichteter Veraltung" (Weight Staleness): Verschiedene Stufen können unterschiedliche Versionen der Gewichte verwenden, was die Konvergenz beeinträchtigen kann. PipeDream führt eine Technik namens "Weight Stashing" ein, bei der mehrere Versionen der Gewichte für verschiedene Mikro-Batches gespeichert werden, um die Konsistenz zu verbessern, was jedoch den Speicherbedarf erhöht.
PipeDream-2BW und Varianten
PipeDream-2BW (2020) verbessert PipeDream, indem es zwei Gewichtspuffer verwendet: einen für Vorwärts- und einen für Rückwärtsdurchläufe. Dies reduziert den Speicherbedarf und verbessert die Konvergenz im Vergleich zum ursprünglichen PipeDream. Varianten wie V-Pipe und PipeMare untersuchen weitere Synchronisationsstrategien und Planungsoptimierungen, um die Blase weiter zu reduzieren und die Skalierbarkeit zu erhöhen.
Megatron-LM und verschachtelte Planung
NVIDIA's Megatron-LM (2019) kombinierte Tensor-Parallelismus (Aufteilung einzelner Schichten) mit Pipeline-Parallelismus. Megatron-2 (2020) führte einen verschachtelten Plan ein, bei dem das Modell in mehr Stufen als Geräte unterteilt wird und jedes Gerät mehrere Stufen bearbeitet. Dies reduziert die Blasengröße, da mehr Mikro-Batches gleichzeitig in der Pipeline sein können, erhöht jedoch den Kommunikationsaufwand.
Chimera und andere neuere Schemata
Chimera (2021) verwendet eine bidirektionale Pipeline, bei der einige Stufen in Vorwärts- und andere in Rückwärtsrichtung arbeiten, um die Blase weiter zu reduzieren. Andere neuere Schemata wie PTD-P (2021) und ZeroBubble (2023) zielen darauf ab, die Blase vollständig zu eliminieren oder zu minimieren, indem sie ausgefeiltere Planungs- und Synchronisationstechniken einsetzen.
Mathematische und praktische Überlegungen
Pipeline-Parallelismus wird häufig mit anderen Parallelismus-Strategien kombiniert. Tensor-Parallelismus teilt einzelne Schichten über Geräte auf, während Pipeline-Parallelismus das Modell vertikal in Stufen unterteilt. Datenparallelität repliziert die Pipeline über Gruppen von Geräten. Dieser hybride Ansatz ist Standard beim Training großer Modelle, wie er in Systemen wie Megatron-Turing NLG und OpenAI's GPT-4 verwendet wird.
Die Wahl der Pipeline-Tiefe (Anzahl der Stufen) und der Mikro-Batch-Größe beeinflusst die Leistung erheblich. Tiefere Pipelines reduzieren den Speicherbedarf pro Stufe, erhöhen jedoch den Kommunikations- und Blasen-Overhead. Größere Mikro-Batches verbessern die Auslastung, benötigen aber mehr Speicher für Aktivierungen. Die optimale Konfiguration hängt von der Modellgröße, dem Speicher der Geräte und der Bandbreite der Verbindungen ab.
Anwendungen und Auswirkungen
Pipeline-Parallelismus war maßgeblich am Training einiger der größten KI-Modelle beteiligt. Beispielsweise wurde GPT-3 (175B Parameter) unter Verwendung einer Kombination aus Modell- und Pipeline-Parallelismus trainiert. Auch Anthropic's Claude-Modelle und Google's PaLM nutzen ähnliche Techniken. Wichtige Frameworks unterstützen Pipeline-Parallelismus, darunter PyTorch (über das torch.distributed.pipeline-Modul), TensorFlow (über tf.distribute) und JAX (mit pjit).
Herausforderungen und Grenzen
Pipeline-Parallelismus bringt mehrere Herausforderungen mit sich. Kommunikations-Overhead entsteht durch das Senden von Aktivierungen und Gradienten zwischen den Stufen, was bei langsamen Verbindungen zum Engpass werden kann. Lastungleichgewichte treten auf, wenn Stufen unterschiedliche Rechenzeiten haben, was die Effizienz verringert. Der Speicherdruck durch das Speichern von Aktivierungen für Rückwärtsdurchläufe kann durch Aktivierungs-Neuberechnung gemildert werden, was jedoch zusätzliche Rechenzeit kostet. Darüber hinaus können asynchrone Schemata unter Konvergenzproblemen aufgrund veralteter Gewichte leiden, während synchrone Schemata Blasen erzeugen.
Zukünftige Richtungen
Die Forschung verbessert den Pipeline-Parallelismus weiterhin mit dem Ziel, Blasen zu reduzieren, den Speicherbedarf zu minimieren und sich an heterogene Hardware anzupassen. Techniken wie automatische Stufenpartitionierung, dynamisches Lastmanagement und die Integration mit Mixture-of-Experts-Modellen sind aktive Forschungsbereiche. Mit dem Aufkommen immer größerer Modelle und dem Bedarf an effizientem Training in Cloud-Umgebungen wird Pipeline-Parallelismus voraussichtlich ein zentraler Bestandteil verteilter Trainingssysteme bleiben.