ゼロバブルパイプラインパラレリズム

Zero Bubble Pipeline Parallelism

November 30, 2023

著者: Penghui Qi, Xinyi Wan, Guangxing Huang, Min Lin

cs.AI

要旨

パイプライン並列処理は大規模分散トレーニングにおける重要な要素の一つであるが、その効率は避けられないとされていたパイプラインバブルによって低下する。本研究では、我々の知る限り、同期トレーニングセマンティクスの下で初めてパイプラインバブルをゼロにすることに成功したスケジューリング戦略を提案する。この改善の鍵となるアイデアは、逆伝播計算を入力に対する勾配を計算する部分とパラメータに対する勾配を計算する部分に分割することである。このアイデアに基づき、我々はベースライン手法を大幅に上回る新しいパイプラインスケジュールを手作業で設計した。さらに、特定のモデル構成とメモリ制限に基づいて最適なスケジュールを自動的に見つけるアルゴリズムを開発した。加えて、真にバブルをゼロにするために、オプティマイザステップ中の同期をバイパスする新しい技術を導入した。実験的評価により、我々の手法は類似のメモリ制限下で1F1Bスケジュールを最大23%スループットで上回ることが示された。この数値は、メモリ制約が緩和された場合、さらに31%まで向上する。我々は、これらの結果がパイプライン並列処理の真の可能性を引き出すための大きな前進を示すものと信じている。我々は、人気のあるMegatron-LMリポジトリに基づく実装をhttps://github.com/sail-sg/zero-bubble-pipeline-parallelismでオープンソースとして公開した。

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.