技術要約：PACIによるバブルの打破

問題提起

パイプライン並列化は、モデルをデバイス間で分割し、計算をオーバーラップさせることで、大規模なニューラルネットワーク（特にTransformer）を訓練するための重要な技術です。しかし、既存のスケジューリング戦略は、ハードウェア利用率、メモリ効率、および最適化の一貫性の間で、永続的なトレードオフに直面しています。

• 同期スケジュールの例（1F1B-Flushなど）: これらはフォワード/バックワードの重みの整合性（あるマイクロバッチのバックワード・パスが、そのフォワード・パスと同じパラメータを使用することを保証すること）を維持しますが、「パイプライン・バブル」（アイドル時間）が発生し、スループットが低下します。
• 非同期スケジュール: これらは完全なハードウェア利用率を実現するためにバブルを排除しますが、フォワード/バックワードの重みバージョンの不整合を導入します。これらのスケジュールでは、マイクロバッチのバックワード・パスが、フォワード・パスの実行時から大幅に進化してしまったパラメータ・バージョンを使用して実行される可能性があります。
• 既存の緩和策: 現在の非同期手法は、重みのスタッシング（複数のパラメータ・コピーの保存）、予測（将来の重みの近似）、またはグローバルな同期を通じて、この不整合を解決しようとしています。これらのアプローチは、大幅なメモリ・オーバーヘッド、追加の計算、システム複雑性、または非標準的な最適化セマンティクスを伴います。

著者らは、フォワード/バックワードの不整合を完全に排除することなく、十分に小さく抑えることができれば、既存の緩和策のコストを回避しつつ、非同期実行の効率性を維持できるのではないか、という問いを投げかけています。

手法：PACI

本論文では、重みのスタッシング、予測、またはグローバルな同期を行うことなく、バージョンドリフトを制限する非同期1F1Bパイプライン・スケジュールであるPACI (Pipeline Asynchronous training with Controlled Inconsistency) を導入しています。

コアメカニズム：バージョン制御としての勾配累積

PACIは、不整合が非同期性のみによって決まるのではなく、マイクロバッチのフォワード・パスとバックワード・パスの間に発生するオプティマイザ・アップデートの数によって支配されるという観察に基づいています。

1. アップデートとバックワード・パスのデカップリング: 標準的な非同期1F1Bでは、毎回のバックワード・パスの後にパラメータが更新されます。PACIはこれをデカップリングします。各ステージは入力の到着直後にバックワード演算を実行しますが、単一のオプティマイザ・ステップを適用する前に、$a$ 回のバックワード・パスにわたってローカルに勾配を累積します。
2. バージョンドリフトの制限: パイプライン遅延に対してパラメータ・バージョンの進化を遅らせることで、PACIはマイクロバッチが通過するアップデートの数を制限します。最大不整合（$\Delta_{max}$）は、パイプラインの深さ（$N$）と累積係数（$a$）によって以下のように制限されます：
   $$\Delta_{max} \le \left\lfloor \frac{N-1}{a} \right\rfloor$$
3. ローカル・フロー制御: 非一様なステージ時間を扱うために、PACIはローカルなバックプレッシャー・ルールを採用しています。各ステージは未解決のフォワード・パスの数（$q_i$）を追跡し、$q_i < N + 1 - i$ である場合にのみ新しいフォワードを許可します。これにより、グローバルな同期バリアを導入することなく、アップストリームのステージがダウンストリームのステージよりも際限なく先行して進むことを防ぎます。
4. リソース効率: PACIは追加の重みメモリを必要としません。単一のパラメータ・コピーを各ステージで使用し、オプティマイザに既に存在する標準的な勾配バッファを再利用します。重みのスタッシング、予測バッファ、またはグローバルな同期は不要です。

マイクロバッチングとの違い

同期スケジュールにおいて、マイクロバッチ数（$m$）を増やす主な目的は、バブルのオーバーヘッドを償却することです。PACIにおいて、マイクロバッチング（累積係数 $a$ による）は、主にバージョンドリフトを制御するための**一貫性のノブ（調整弁）**として使用されます。これにより、PACIは高い利用率というシステムの目標と、低い不整合という最適化の目標を分離することができます。

主な貢献

1. スタッシングなしでの限定された不整合: 本論文は、ローカルな勾配累積が、追加のメモリ、予測、または同期なしに、フォワード/バックワードの不整合を制御するチューナブルなバージョン制御メカニズムとして機能できることを示しています。
2. 改善されたTime-to-Accuracy: PACIは、同期的な1F1B-flushと同等の最終パープレキシティを維持しながら、安定した事前学習を実現し、エンドツーエンドの実行時間を（バッチサイズ128で）最大1.69倍、（バッチサイズ256で）1.41倍（最速のflushベースラインと比較して）短縮します。
3. 理論に合致したスループット分析: 著者らは、同期的な1F1B-flushとPACIのスループットの差が、完全にパイプライン・バブルの効率性によって説明されることを示しています。PACIは、同期的なフラッシングと同じピークメモリ・フットプリントを維持しながら、パイプラインの理論上のフル利用スループットを達成します。

実験結果

著者らは、8つのNVIDIA RTX PRO 6000 GPUを用いた8ステージのパイプライン上で、GPT-2 MediumのOpenWebTextによるスクラッチからの事前学習を用いてPACIを評価しました。

• 安定性 (Q1): 限定された不整合 $\Delta_{max} \le 2$ （累積係数 $a=4$ または $a=8$ によって達成）の場合、PACIは同期的な1F1B-flushの損失ダイナミクスに密接に一致します。最終的な損失は同等であり、特筆すべきは、実行ごとの変動性（訓練損失のRMS標準偏差）が低いことです。
• Time-to-Accuracy (Q2): PACIは、同期的なベースラインよりも大幅に早くターゲットとなるパープレキシティレベルに到達します。
  • バッチサイズ128において、PACIは最適なチューニングが施された1F1B-flushベースラインよりも1.69倍高速です。
  • 一致したマイクロバッチ構成と比較して、PACIは最大2.04倍のスピードアップを達成します。
• スループットとメモリ (Q3):
  • 1F1B-flushのスループットはバブル効率（マイクロバッチが増えるにつれて改善するが、完全な利用には達しない）に依存しますが、PACIはフル利用のスループットを達成します。
  • PACIは1F1B-flushと同じピークメモリ・フットプリントを実現します。一方で、他の非同期手法（PipeDreamやPipeMareなど）は、重みのスタッシングや補助的な状態のために、より多くのメモリを必要とします。
  • 最大28.3Bパラメータまでの大規模な外挿により、PACIはZero-Bubble (ZB-2p) スケジュールのスループットに到達しつつ、1F1B-flushのメモリ効率を維持できることが示唆されています。

意義と主張

本論文は、フォワード/バックワードの不整合を排除する必要はないと主張しています。代わりに、勾配累積を通じて明示的に制限することで、それは大幅な効率向上のためのトレードオフとして安全に利用できるとしています。

• 動作点: PACIは、パイプライン並列化における新しい動作点を確立します：ゼロのパイプライン・バブル、ゼロの追加重みメモリ、そして低い限定された不整合です。
• パラダイムシフト: 本研究は、非同期訓練において重みのスタッシングや予測が必要であるという考えに異を唱えています。パラメータの進化の「速度」をパイプライン遅延に対して制御することが、最適化の安定性を維持するのに十分であることを示唆しています。
• 実用的な影響: 「バブル」のオーバーヘッドを、メモリのペナルティを負うことなく取り除くことで、PACIは、特に同期スケジュールがアイドル時間によって制限される領域において、大規模言語モデルの訓練時間を短縮するための実用的な道筋を提供します。

著者らは、現在の評価がGPT-2 Medium、8ステージのパイプライン、および特定の訓練構成に限定されていることを認め、限界についても謙虚な姿勢を保っています。アクティベーション・チェックポインティングの統合や、ロールバックによる無効な勾配（NaN）の処理などは、今後の課題として残されています。