Pipeline Parallelität: Effizienzsteigerung beim Training riesiger KI-Modelle

Pipelining-Parallelismus ist ein Schlüsselkonzept für das Training großer, verteilter KI-Modelle. Angesichts der kontinuierlichen Zunahme der Modellgröße – von OpenAIs GPT-2 mit 1,5 Milliarden Parametern bis hin zu GPT-3 mit 175 Milliarden Parametern – erweisen sich herkömmliche Trainingsmethoden wie das verteilte Datenparallele Training (DDP) als unzureichend, da die Modelle nicht auf einzelnen GPU-Geräten Platz finden. Um solche umfangreichen Modelle zu trainieren, muss auf Techniken des Modellparallelismus zurückgegriffen werden. In der Industrie wurden verschiedene Ansätze wie GShard und Mesh-Tensorflow von Google, ZeRO von Microsoft und Megatron-LM von Nvidia eingesetzt, um dieser Herausforderung zu begegnen.

Die Pipeline-Parallelität, die ursprünglich im GPipe-Papier beschrieben wurde, stellt eine dieser Techniken dar. Sie ermöglicht es, die Schichten eines Modells sequenziell auf mehrere GPUs zu verteilen und eine Minibatch von Daten in mehrere Mikrobatches aufzuteilen, die dann in einer Pipeline über mehrere Geräte hinweg ausgeführt werden. Obwohl PyTorch die grundlegenden Bausteine für eine solche pipelinierte Ausbildung bietet, bringt die Umsetzung von Pipeline-Parallelität eine gewisse Komplexität mit sich. Daher wäre es äußerst nützlich, wenn den Nutzern ein benutzerfreundlicher Rahmen wie DDP zur Verfügung gestellt werden könnte.

Es gibt zwei breite Kategorien von Pipelining-Ansätzen: synchrones Pipelining, wie bei GPipe/torchgpipe, das eine bessere Trainingsgenauigkeit bietet und die vorhandene Trainingsstruktur nicht wesentlich verändert, jedoch eine geringe GPU-Auslastung und niedrige Speichernutzung hat, und asynchrones Pipelining, wie bei PipeDream/XPipe, das eine höhere GPU-Auslastung und somit einen besseren Trainingsdurchsatz ermöglicht, aber möglicherweise die Modellqualität aufgrund der Parameterstagnation beeinträchtigt und Anpassungen an der Natur des kontinuierlichen Datenflusses erfordert. Asynchrone Systeme speichern in der Regel mehrere Versionen der Gewichte (Gewichtsversteckung/Vorhersage), was zu einem erhöhten Speicherbedarf führt.

In der Praxis bedeutet dies, dass bei der Verwendung von Pipeline-Parallelität die Effizienz unter Pipeline-Blasen leidet, die als unvermeidlich galten. In dieser Hinsicht stellt die jüngste Forschungsarbeit von @_akhaliq einen Durchbruch dar: Es wurde eine Planungsstrategie vorgestellt, die angeblich die ersten erfolgreichen Versuche unternimmt, Null-Pipeline-Blasen unter synchronen Trainingsschemata zu erreichen. Der Schlüssel zu dieser Verbesserung liegt in der Aufteilung der Rückwärtsberechnung in zwei Teile: einen, der den Gradienten für die Eingabe berechnet, und einen weiteren, der ihn für die Parameter berechnet. Basierend auf dieser Idee wurden neue Pipeline-Zeitpläne entwickelt, die die Basismethoden deutlich übertreffen.

Ein Algorithmus wurde entwickelt, der automatisch einen optimalen Zeitplan basierend auf einer spezifischen Modellkonfiguration und Speicherbegrenzung findet. Um eine wirklich Null-Pipeline-Blase zu erreichen, wurde eine neuartige Technik eingeführt, die Synchronisierungen während des Optimierungsschritts umgeht. Experimentelle Bewertungen zeigen, dass die Methode den 1F1B-Zeitplan um bis zu 23% im Durchsatz unter einer ähnlichen Speicherbegrenzung übertrifft. Diese Zahl kann weiter auf 31% gesteigert werden, wenn die Speicherbeschränkung gelockert wird. Die Ergebnisse stellen einen wichtigen Schritt nach vorne dar, um das wahre Potenzial der Pipeline-Parallelität auszuschöpfen.

Das von Huggingface entwickelte BLOOM-Modell ist ein solches Beispiel, bei dem Pipeline-Parallelismus verwendet wurde, um ein 175 Milliarden Parameter großes Transformer-Modell zu trainieren, das die Gewichte als bfloat16 speichert und 350 GB benötigt – mehr, als die zum Training verwendeten GPUs mit 80 GB Speicher bieten könnten. Das Training wurde schließlich auf 384 GPUs verteilt. Dies wurde durch die Zuweisung verschiedener Schichten des Modells zu verschiedenen GPUs möglich, ein Prozess, der als Modellpartitionierung bezeichnet wird. Eine naiv umgesetzte Modellpartitionierung führt jedoch zu einer geringen GPU-Auslastung.

Die vorgestellten Erkenntnisse und technischen Fortschritte im Bereich des Pipeline-Parallelismus sind ein Indiz dafür, wie schnell sich das Feld der verteilten KI-Trainingsmethoden entwickelt. Die Fähigkeit, große Modelle effizient zu trainieren, ist für die Zukunft der künstlichen Intelligenz von entscheidender Bedeutung, und die neuesten Forschungsergebnisse, wie die hier diskutierten, bieten wertvolle Einblicke, wie diese Herausforderungen bewältigt werden können. Mit solchen Fortschritten wird es möglich, komplexe Modelle schneller zu trainieren und sie für eine breitere Palette von Anwendungen zugänglich zu machen, was wiederum die Grenzen dessen, was KI leisten kann, weiter verschiebt.