Null-Blasen-Pipeline-Parallelismus

Zero Bubble Pipeline Parallelism

November 30, 2023

Autoren: Penghui Qi, Xinyi Wan, Guangxing Huang, Min Lin

cs.AI

Zusammenfassung

Pipeline-Parallelismus ist eine der Schlüsselkomponenten für das groß angelegte verteilte Training, doch seine Effizienz leidet unter Pipeline-Blasen, die bisher als unvermeidbar galten. In dieser Arbeit stellen wir eine Scheduling-Strategie vor, die unseres Wissens nach die erste ist, die erfolgreich null Pipeline-Blasen unter synchronen Trainingssemantiken erreicht. Die zentrale Idee hinter dieser Verbesserung besteht darin, die Rückwärtsberechnung in zwei Teile zu unterteilen: einen, der den Gradienten für die Eingabe berechnet, und einen anderen, der den Gradienten für die Parameter berechnet. Basierend auf dieser Idee haben wir neuartige Pipeline-Schedules entwickelt, die die Baseline-Methoden deutlich übertreffen. Darüber hinaus entwickeln wir einen Algorithmus, der automatisch einen optimalen Schedule basierend auf der spezifischen Modellkonfiguration und der Speicherbegrenzung findet. Zusätzlich führen wir eine neuartige Technik ein, um Synchronisationen während des Optimierungsschritts zu umgehen und so wirklich null Blasen zu erreichen. Experimentelle Auswertungen zeigen, dass unsere Methode den 1F1B-Schedule bei einem ähnlichen Speicherlimit um bis zu 23% im Durchsatz übertrifft. Diese Zahl kann auf 31% gesteigert werden, wenn die Speicherbeschränkung gelockert wird. Wir glauben, dass unsere Ergebnisse einen bedeutenden Schritt nach vorne darstellen, um das wahre Potenzial des Pipeline-Parallelismus auszuschöpfen. Wir haben unsere Implementierung basierend auf dem beliebten Megatron-LM-Repository auf https://github.com/sail-sg/zero-bubble-pipeline-parallelism quelloffen veröffentlicht.

English

Pipeline parallelism is one of the key components for large-scale distributed training, yet its efficiency suffers from pipeline bubbles which were deemed inevitable. In this work, we introduce a scheduling strategy that, to our knowledge, is the first to successfully achieve zero pipeline bubbles under synchronous training semantics. The key idea behind this improvement is to split the backward computation into two parts, one that computes gradient for the input and another that computes for the parameters. Based on this idea, we handcraft novel pipeline schedules that significantly outperform the baseline methods. We further develop an algorithm that automatically finds an optimal schedule based on specific model configuration and memory limit. Additionally, to truly achieve zero bubble, we introduce a novel technique to bypass synchronizations during the optimizer step. Experimental evaluations show that our method outperforms the 1F1B schedule up to 23% in throughput under a similar memory limit. This number can be further pushed to 31% when the memory constraint is relaxed. We believe our results mark a major step forward in harnessing the true potential of pipeline parallelism. We open sourced our implementation based on the popular Megatron-LM repository on https://github.com/sail-sg/zero-bubble-pipeline-parallelism.