[論文レビュー] Zero Bubble Pipeline Parallelism

Penghui Qi, Xinyi Wan+2 more|arXiv (Cornell University)|Nov 30, 2023

Parallel Computing and Optimization Techniques被引用数 4

ひとこと要約

本稿では、バックワード計算を入力勾配（B）とパrameter勾配（W）に分解することで、同期学習におけるほぼゼロのパイプラインバブルを達成する新しいスケジューリング戦略、Zero Bubble Pipeline Parallelism（ZB）を提案する。より細かい粒度での計算グラフの再編成と、新しいスケジュール（ZB-V）の導入により、1F1Bと比較してバブル率を最大75%まで低減し、同程度のメモリ制限下ではスループットを23%向上させ、メモリ制限を緩和した場合には最大31%向上する。

ABSTRACT

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.

研究の動機と目的

• 同期学習の文脈において、パイプラインバブルが原因で生じる持続的な非効率性を解消すること。
• 1F1Bのような既存のスケジューリング戦略ですら非ゼロのバブル率を抱えるという限界を乗り越えること。
• バブル率を最小限に抑えつつ、メモリ制約を尊重し、大規模モデル学習を効率的に行えるスケジューリングフレームワークを設計すること。
• 最適化の意味論を損なわず、ベースライン手法を超えるメモリ使用量を増やさずにゼロバブル学習を実現すること。
• モデル構成とメモリ制限に基づいて最適なスケジュールを自動生成するアルゴリズムを開発すること。

提案手法

• バックワードパスを、入力（B）に関する勾配計算とパラメータ（W）に関する勾配計算という2つの独立したコンponentsに分解することで、従来のモノリシックなバックワード関数を解除する。
• パイプラインバブルの原因となる依存関係を解消するため、BとWの計算順序を再編成した新しいパイプラインスケジュール「ZB-V」を導入し、フォワード、バックワード、パラメータ勾配の実行時間が等しい場合には完全にバブルをゼロにできる。
• 最適化ステップ中に、同期をバイパスする新規技術を設計することで、無駄な待機時間を排除し、さらにバブルを低減する。
• メモリ制約とモデル制約の下で最適なスケジュールを求める自動スケジューリングアルゴリズムを開発し、バブル率を最小化する。
• レイヤー単位の抽象化ではなく、関数レベル（F, B, W）での計算グラフ表現を採用することで、より細かい粒度のスケジューリングと最適化を可能にする。
• Megatron-LM を基盤として実装・オープンソース化し、既存の大規模学習フレームワークへの統合を可能にする。

実験結果

リサーチクエスチョン

• RQ1 バックワード計算の再設計により、同期パイプライン並列学習におけるパイプラインバブルを完全に排除できるか？
• RQ2 メモリ効率と標準的な学習セマンティクスとの互換性を保ちながら、可能な限り低いバブル率を達成するスケジューリング戦略は何か？
• RQ3 バックワードパスを入力勾配とパrameter勾配に分解することで、パイプライン効率とメモリ使用量にどのような影響を与えるか？
• RQ4 自動スケジューリングが、多様なモデルサイズとメモリ予算において、どの程度バブル率を低減できるか？
• RQ5 マイクロバッチサイズとバブル率のトレードオフは何か？最大スループットを達成するために最適化可能か？

主な発見

• 28.3Bパラメータのモデル（32段階、256マイクロバッチ）において、ZB-Vはバブル率0.0274を達成した。一方1F1Bは0.2676であり、89.8%の低減が確認された。
• 同程度のメモリ制限下ではZB-Vが1F1B比で最大23%のスループット向上を達成し、メモリ制限を緩和した場合には最大31%の向上を達成した。
• 全評価対象のモデルサイズおよび構成において、ZB-Vは1F1Bと比較して最大75%のバブル率低減を実現した。
• ZB-Vは1F1Bと同等のメモリフットプリントを維持しながら、ほぼゼロのバブルを達成した。一方、ZB-H2は同等の性能を得るためにはメモリを2倍に必要とする。
• 自動スケジューリングアルゴリズムは、特にメモリ制限が厳しい環境下でも、ヒューリスティックベースラインを上回る最適なスケジュールを効果的に生成した。
• マイクロバッチサイズの増加により、大規模モデルでは最大8%のスループット向上が得られたが、小規模モデルではその恩恵が薄れ、マイクロバッチサイズとバブル率のトレードオフが顕著に現れた。