## Memo ## Memo with LLM ### 論文情報 - **タイトル**: Zero Bubble (Almost) Pipeline Parallelism - **著者と所属**: - Penghui Qi (Sea AI Lab) - Xinyi Wan (Sea AI Lab) - Guangxing Huang (Sea AI Lab) - Min Lin (Sea AI Lab) - **カンファレンス**: ICLR 2024 - **発表年**: 2024年 ### 論文概要 本論文は、大規模分散学習における重要な要素であるパイプライン並列化において、従来避けられないとされていたパイプラインバブル（アイドル時間）をほぼゼロにする新しいスケジューリング戦略を提案しています。バックワード計算を入力勾配計算とパラメータ勾配計算の2つに分割することで、1F1Bスケジュールと比較して最大23%のスループット向上を実現し、メモリ制約を緩和した場合は最大31%の改善が得られることを実証しています。 ### 詳細解説 #### 問題設定 **入力と出力**: - 入力: パイプライン段数 p、マイクロバッチ数 m、アクティベーションメモリ制限 M_limit、各演算（F, B, W）の実行時間推定値 T_F, T_B, T_W - 出力: 最適化されたパイプラインスケジュール（各段での演算の実行順序） **データの種類**: - 大規模ニューラルネットワークモデル（特にTransformerアーキテクチャ） - GPUクラスタ上での分散学習 - 実験では GPT-3 類似のモデル（1.5B～28.3Bパラメータ）を使用 **問題の背景**: パイプライン並列化（PP）では、モデルを複数の段に分割し、各段を異なるデバイスに配置します。従来の1F1B（1-Forward-1-Backward）スケジュールでは、依存関係により避けられないアイドル時間（バブル）が発生し、計算資源の利用効率が低下します。 #### 提案手法 **コア・アイデア: バックワードパスの分割** 従来、バックワードパスは一つの関数として実装されていましたが、本研究では以下の2つに分割します： 1. **B（入力勾配計算）**: $\frac{\partial L}{\partial x} = f_B(x, \theta, \frac{\partial L}{\partial y})$ 2. **W（パラメータ勾配計算）**: $\frac{\partial L}{\partial \theta} = f_W(x, \theta, \frac{\partial L}{\partial y})$ この分割により、B は次の段の B にのみ依存し、W は柔軟に配置できるようになります。 **手作りスケジュール** 論文では2つの手作りスケジュールを提示しています： 1. **ZB-H1（メモリ効率的スケジュール）**: - 1F1Bと同じメモリ制約下で動作 - バブルサイズを1F1Bの約1/3に削減 - ピークアクティベーションメモリ: $(p + m - 1) \cdot M_B$ 2. **ZB-H2（ゼロバブルスケジュール）**: - より大きなメモリを使用してゼロバブルを実現 - ウォームアップ段階で追加のFパスを実行 - Wパスを並べ替えて平行四辺形の形状を作成 - ピークアクティベーションメモリ: $2(p - 1) \cdot M_B + M_W$ **自動スケジューリングアルゴリズム** 実際の環境では $T_F = T_B = T_W$ という仮定は成立しないため、ヒューリスティックアルゴリズムを開発： 1. ウォームアップ段階: メモリ制限内で可能な限り多くのFパスをスケジュール 2. 定常状態: 1F-1B-1Wパターンに従い、$T_W$ より大きいギャップにWを挿入してバブルを埋める 3. 終了段階: 残りのWパスを順次スケジュール さらに、整数線形計画法（ILP）を用いた最適解の探索も可能です。 **オプティマイザ同期のバイパス** ゼロバブルを真に達成するため、従来のオプティマイザステップ前の同期を事後検証に置き換えます： - 各段は部分的に削減されたグローバル状態を受け取り、オプティマイザステップを実行 - 次のイテレーションのウォームアップ段階で完全に削減されたグローバル状態を伝播 - 必要に応じてロールバックと再実行を行う この手法により、勾配クリッピングやNAN/INFチェックのための同期がパイプラインを破壊することを防ぎます。 **ZB-V（メモリ効率的ゼロバブルスケジュール）** 1F1Bと同じメモリ制約下でゼロバブルに近い性能を実現する追加の手法： - モデルを2p個のチャンクに均等に分割 - 各ワーカーに正確に2つのチャンクを割り当て - チャンクの依存関係が「V」字形のパターンに従う - ピークアクティベーションメモリ: $(p + m - 1) \cdot M_B$ （1F1Bと同等） #### 新規性 **先行研究との比較**: 1. **GPipe (Huang et al., 2019)**: - マイクロバッチ数を増やしてバブルを削減 - 中間アクティベーションを破棄し再計算（約20%の計算オーバーヘッド） 2. **PipeDream (Harlap et al., 2018)**: - 非同期PPでバブルフリーを実現 - 厳密な最適化セマンティクスを犠牲にする 3. **1F1B (Fan et al., 2021; Narayanan et al., 2021)**: - GPipeより早期にバックワードパスを実行してメモリを解放 - 同じマイクロバッチ数で同様のバブル率、但しピークメモリは低い 4. **Interleaved 1F1B (Narayanan et al., 2021)**: - 複数の段を同一デバイスに割り当て - バブルサイズを削減するが、通信とメモリが増加 **本研究の独自性**: - **同期セマンティクスを保持しながらゼロバブルを達成**: これまでの研究では、完全なゼロバブルは非同期設定でのみ可能でした - **より細かい粒度での最適化**: バックワードパスを2つに分割することで、従来のレイヤーレベルよりも細かい制御が可能 - **メモリとスループットのトレードオフを制御**: ZB-1p, ZB-2p, ZB-Vなど、異なるメモリ制約に対応する複数のスケジュールを提供 - **自動スケジューリング**: 手作業による調整ではなく、モデル構成とメモリ制限に基づいて最適スケジュールを自動的に探索 #### 実験設定 **使用したデータセット**: - GPT-3類似のモデルアーキテクチャを使用 - 4つのモデルサイズ: 1.5B, 6.2B, 14.6B, 28.3Bパラメータ **モデル構成**: | モデル | レイヤー数 | ヘッド数 | 隠れ層次元 | シーケンス長 | パイプライン段数 | |--------|------------|----------|------------|--------------|------------------| | 1.5B | 22 | 24 | 2304 | 1024 | 8 | | 6.2B | 30 | 32 | 4096 | 1024 | 8 | | 14.6B | 46 | 40 | 5120 | 1024 | 16 | | 28.3B | 62 | 48 | 6144 | 1024 | 32 | **ハードウェア**: - 最大32個のNVIDIA A100 SXM 80G GPU - 4ノード構成 - RoCE RDMA ネットワークで相互接続 **評価指標**: 1. **スループット（サンプル/秒）**: 各イテレーションの実行時間を測定し、ウォームアップ後に記録 2. **バブル率**: $\text{bubble rate} = \frac{\text{cost} - \text{optimal cost}}{\text{cost}}$ - cost: 全段の中で最大の実行時間 - optimal cost: 全通信が計算とオーバーラップした場合の最適実行時間 3. **ピークメモリ使用量（GB）**: アクティベーションメモリの最大使用量 **比較手法**: - 1F1B: 標準的なパイプライン並列化手法 - 1F1B-I: インターリーブド1F1B - ZB-1p: 1F1Bと同じメモリ制約の自動探索スケジュール - ZB-2p: メモリを2倍にしてゼロバブルを実現する自動探索スケジュール - ZB-V: 1F1Bと同じメモリでゼロバブルに近い性能を実現 **正確性の検証**: - Megatron-LMの再現性を利用 - 固定乱数シードでモデルを初期化 - 各イテレーション後の損失値がビット単位で一致することを確認 #### 実験結果 **主要な実験結果（表形式）**: 1. **スループット比較（1.5Bモデル、8GPU、24マイクロバッチ）**: - 1F1B: 11.8 サンプル/秒 - 1F1B-I: 13.1 サンプル/秒 - ZB-1p: 12.9 サンプル/秒（1F1Bより9.3%向上） - ZB-2p: 14.5 サンプル/秒（1F1Bより22.9%向上） 2. **スループット比較（6.2Bモデル、8GPU、24マイクロバッチ）**: - 1F1B: 3.50 サンプル/秒 - 1F1B-I: 4.01 サンプル/秒 - ZB-1p: 3.88 サンプル/秒（1F1Bより10.9%向上） - ZB-2p: 4.32 サンプル/秒（1F1Bより23.4%向上） 3. **バブル率の比較（14.6Bモデル、16段、48マイクロバッチ）**: - 1F1B: 25.52% - 1F1B-I: 11.04% - ZB-H1: 13.97% - ZB-H2: 6.72% - ZB-1p: 13.97% - ZB-2p: 0.66%（ほぼゼロ） 4. **ZB-Vの性能（6.2Bモデル、16GPU、64マイクロバッチ）**: - ZB-V: 4.21 サンプル/秒、メモリ 64GB - ZB-2p*: 4.37 サンプル/秒、メモリ 64GB - ZB-1p: 4.00 サンプル/秒、メモリ 62GB - 1F1B: 3.57 サンプル/秒、メモリ 61GB **重要な観察結果**: 1. **ZB-2pの一貫した優位性**: 全ての設定でZB-2pが最高のスループットを達成し、バブル率は1%未満 2. **マイクロバッチ数への依存性**: - 1F1B、1F1B-I、ZB-1pはマイクロバッチ数と強い正の相関 - ZB-2pはマイクロバッチ数が少なくても効率を維持（バブル率がほぼゼロのため） 3. **メモリ効率**: - ZB-1pは1F1Bと同等のメモリで、特にマルチノード設定で1F1B-Iを上回る - ZB-2pは約2倍のメモリを使用するが、最高のスループットを実現 4. **自動スケジューリングの効果**: - ZB-2pは手作りのZB-H2を一貫して上回る - より正確な T_F, T_B, T_W の推定値により、現実的なシナリオに適応 5. **オプティマイザ同期バイパスの効果**: - 事後検証により約8%のスループット向上 - 同期版と比較して顕著な性能差 6. **ZB-Vのトレードオフ**: - 1F1Bと同じメモリ制約でZB-2p*に匹敵する性能 - マイクロバッチサイズが大きい場合（14.6Bおよび28.3Bモデル）で8%の改善 - バブル率は1F1Bの約1/4 **メモリ限界の分析**: メモリ限界 M_limit とバブル率の関係を調査した結果： - M_limit を増やすと、バブル率はほぼ線形に減少 - M_limit ≈ 2 × (p×M_B) でゼロバブル率に近づく閾値に到達 - この閾値を超えると、コストが利得を上回る **実環境での実行プロファイル**: 16GPU上でZB-2pの理論的スケジュールと実際の実行をプロットした結果、ほぼ完璧な一致を確認。わずかなバブルは存在するものの、全体的なアライメントは良好で、真のゼロバブルスケジュールであることを視覚的に証明しています。 ## Abstract パイプライン並列化は大規模分散学習の重要な要素の一つですが、その効率は避けられないとされてきたパイプラインバブルによって損なわれています。本研究では、我々の知る限り、同期学習セマンティクスの下でゼロパイプラインバブルを初めて成功裏に達成するスケジューリング戦略を紹介します。この改善の背後にある重要なアイデアは、バックワード計算を2つの部分に分割することです。一つは入力の勾配を計算し、もう一つはパラメータの勾配を計算します。このアイデアに基づき、ベースライン手法を大幅に上回る新しいパイプラインスケジュールを手作りしました。さらに、特定のモデル構成とメモリ制限に基づいて最適なスケジュールを自動的に見つけるアルゴリズムを開発しました。さらに、真にゼロバブルを達成するために、オプティマイザステップ中の同期をバイパスする新しい技術を導入しています。実験的評価により、我々の手法は同様のメモリ制限下で1F1Bスケジュールを最大23%上回るスループットを示しています。この数値は、メモリ制約が緩和された場合、31%まで押し上げることができます。我々は、この結果がパイプライン並列化の真の潜在能力を引き出す上で大きな前進を示すものと信じています。Megatron-LMに基づく実装をhttps://github.com/sail-sg/zero-bubble-pipeline-parallelismでオープンソース化しました。