[論文レビュー] PipeFisher: Efficient Training of Large Language Models Using Pipelining and Fisher Information Matrices

Kazuki Osawa, Shigang Li+1 more|arXiv (Cornell University)|Nov 25, 2022

Topic Modeling被引用数 4

ひとこと要約

PipeFisher は、パイプライン並列 LLM 学習における空き時間（バブル）を活用して、フィッシャー情報行列に基づく 2 階微分最適化手法 K-FAC を実行する、革新的な学習スキームを提案する。K-FAC の曲率計算と行列逆行列計算をこれらの空き時間にオフロードすることで、GPU 利用率を 59.8% から 97.6% まで向上させ、BERT-Large のシミュレートされたフェーズ 1 の事前学習時間を 75.7% 減少させるとともに、より良い最適化により収束性も向上させる。

ABSTRACT

Pipeline parallelism enables efficient training of Large Language Models (LLMs) on large-scale distributed accelerator clusters. Yet, pipeline bubbles during startup and tear-down reduce the utilization of accelerators. Although efficient pipeline schemes with micro-batching and bidirectional pipelines have been proposed to maximize utilization, a significant number of bubbles cannot be filled using synchronous forward and backward passes. To address this problem, we suggest that extra work be assigned to the bubbles to gain auxiliary benefits in LLM training. As an example in this direction, we propose PipeFisher, which assigns the work of K-FAC, a second-order optimization method based on the Fisher information matrix, to the bubbles to accelerate convergence. In Phase 1 pretraining of BERT-Base and -Large models, PipeFisher reduces the (simulated) training time to 50-75% compared to training with a first-order optimizer by greatly improving the accelerator utilization and benefiting from the improved convergence by K-FAC.

研究の動機と目的

• パイプライン並列 LLM 学習におけるアクセラレータ利用率の低さを是正すること。これは、起動および終了時の「バブル」と呼ばれる空き時間に起因する。
• これらの空き時間を活用し、単なる利用率向上を超えた付加的利点をもたらす意味のある計算作業を割り当てる。
• K-FAC を用いた 2 階微分最適化が、通信コストを増加させることなくパイプラインスケジューリングに効率的に統合可能であることを示すこと。
• GPU 利用率の向上と K-FAC による収束速度の向上という相乗効果が、学習時間の顕著な短縮に寄与することを示すこと。

提案手法

• PipeFisher は、GPipe や Chimera などのパイプライン並列学習スケジュールにおける空き時間（バブル）に、K-FAC の曲率計算と行列逆行列計算を統合する。
• パイプラインのバブル中に、フィッシャー情報行列（A_l, B_l）の計算とその逆行列計算という K-FAC の作業をスケジューリングし、フォワードおよびバックワードパスのブロッキングを回避する。
• 初期学習段階での収束安定性を確保するため、直前のステップからの古くなった逆行列を最初の前処理に使用する。
• 既存のパイプラインスケジューリングフレームワークを活用し、標準的なパイプライン学習に加えて最小限の計算オーバーヘッド（ステップあたり約 6.5%）で実現可能である。
• 本手法はあらゆるパイプライン並列方式と互換性があり、Transformer および非Transformerアーキテクチャの両方へ適用可能であるが、非一様なレイヤーでは負荷分散がより複雑になる。
• より短い学習サイクルにより、さまざまな学習率スケジュールを評価するコストが低下するため、ハイパーパramータサーチが高速化される。

実験結果

リサーチクエスチョン

• RQ1 パイプライン並列 LLM 学習における空き時間（バブル）は、計算的に高負荷な最適化手法の実行に効果的に活用可能か？
• RQ2 K-FAC の計算をパイプラインのバブルにオフロードすることで、全体の学習効率と収束速度が向上するか？
• RQ3 通信コストを増加させることなく、バブルに補助的作業を割り当てることで、GPU 利用率をどの程度向上できるか？
• RQ4 PipeFisher を用いた K-FAC の統合は、BERT 事前学習における学習時間とモデル性能にどのような影響を及ぼすか？

主な発見

• PipeFisher は、Chimera パイプラインスケジューリングを用いた BERT-Large 事前学習において、GPU 利用率を 59.8% から 97.6% まで向上させ、空き時間の顕著な削減を実現した。
• BERT-Large のフェーズ 1 のシミュレートされた学習時間は、275.1 分から 208.3 分に短縮され、75.7% の改善が達成された。これは、高い利用率と K-FAC による収束速度の向上のおかげである。
• BERT-Base においても、PipeFisher はベースラインの 50% の学習時間にまで短縮した。これは、モデルサイズの異なる状況でも一貫した向上効果を示している。
• PipeFisher の計算オーバーヘッドは最小限（ステップあたり約 6.5%）であり、前処理が標準的なパイプライン学習に追加される唯一のコストである。
• 本手法により、さまざまな学習率スケジュールを評価するための時間が短縮され、ハイパーパramータサーチが高速化された。
• PipeFisher は、Shampoo や Sharpness-Aware Minimization (SAM) などの他の 2 階微分最適化手法や補助的タスクへも一般化可能で、拡張が可能である。