[論文レビュー] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Samyam Rajbhandari, Jeff Rasley+2 more|arXiv (Cornell University)|Oct 4, 2019

Advanced Neural Network Applications参考文献 21被引用数 72

ひとこと要約

ZeRO は最適化オンメモリの状態、勾配、パラメータのメモリ効率的な分割（残差メモリ最適化を含む）を導入し、データ並列で trillion パラメータ級のモデルを訓練可能にし、巨大スピードアップと大規模モデル訓練の民主化を実現します。

ABSTRACT

Large deep learning models offer significant accuracy gains, but training billions to trillions of parameters is challenging. Existing solutions such as data and model parallelisms exhibit fundamental limitations to fit these models into limited device memory, while obtaining computation, communication and development efficiency. We develop a novel solution, Zero Redundancy Optimizer (ZeRO), to optimize memory, vastly improving training speed while increasing the model size that can be efficiently trained. ZeRO eliminates memory redundancies in data- and model-parallel training while retaining low communication volume and high computational granularity, allowing us to scale the model size proportional to the number of devices with sustained high efficiency. Our analysis on memory requirements and communication volume demonstrates: ZeRO has the potential to scale beyond 1 Trillion parameters using today's hardware. We implement and evaluate ZeRO: it trains large models of over 100B parameter with super-linear speedup on 400 GPUs, achieving throughput of 15 Petaflops. This represents an 8x increase in model size and 10x increase in achievable performance over state-of-the-art. In terms of usability, ZeRO can train large models of up to 13B parameters (e.g., larger than Megatron GPT 8.3B and T5 11B) without requiring model parallelism which is harder for scientists to apply. Last but not the least, researchers have used the system breakthroughs of ZeRO to create the world's largest language model (Turing-NLG, 17B parameters) with record breaking accuracy.

研究の動機と目的

• トリリオンパラメータ級モデルを訓練する必要性を動機づけ、既存のデータ/モデル/パイプライン並列性におけるメモリのボトルネックを特定する。
• モデル状態と残差のメモリ冗長性を排除しつつ効率を保持するために ZeRO を提案する。
• 利用可能なハードウェアとともにモデルサイズを拡張するための3つの分割段階（オプティマイザ状態、勾配、パラメータ）と残差メモリ最適化を開発・分析する。
• 実践的な訓練の改善を実証し、MP との統合を示し、大規模モデル訓練を民主化するオープンソースツールを提供する。

提案手法

• データ並列プロセス間でオプティマイザ状態を分割し（P_os）、各ランクが全状態の 1/N_d を保持する。
• 各パラメータ分割に対応する勾配を分割（P_g）し、分割境界でバケット化された reduct-scatter を実行する。
• データ並列プロセス間でパラメータを分割（P_p）し、前方/後方伝播時の非局所パラメータをオンデマンドでブロードキャストする。
• ZeRO-DP を ZeRO-R と組み合わせてアクティベーションも分割し、残差メモリ（チェックポイント、バッファ、断片化）を管理し、任意の CPU オフロードを含む（P_a、P_a+cpu）。
• 動的な通信スケジュールを用いて通信を最小化し、メモリ節約と計算の分離を実現する。通信オーバーヘッドを分析（P_p の場合最大1.5x）。
• 有用な場合にはモデル並列性（MP）との組み合わせをサポートするが、実装例では ZeRO-DP が多くのケースで MP に匹敵・上回り、広範なモデルのリファクタリングなしに訓練を可能にする。

実験結果

リサーチクエスチョン

• RQ1 ZeRO はモデル状態を分割して DP 効率を損なうことなくデバイスごとのメモリを大幅に削減できるか？
• RQ2 データ並列度 N_d の関数として、オプティマイザ状態、勾配、パラメータのメモリ節約量はどれくらいか？
• RQ3 残差メモリ（アクティベーション、バッファ、断片化）をどのように管理して非常に大きなモデルをさらに可能にできるか？
• RQ4 ZeRO-DP と ZeRO-R を適用したときの通信量とスケーラビリティへの影響はどうなるか？
• RQ5 ZeRO は数百億〜トリリオん級のモデル訓練をどのように実現し、実機ハードウェア上で現実的なスピードアップはどの程度か？

主な発見

DP度	7.5B モデル (GB) P_os	7.5B モデル (GB) P_os+g	7.5B モデル (GB) P_os+g+p	128B モデル (GB) P_os	128B モデル (GB) P_os+g	128B モデル (GB) P_os+g+p	1T モデル (GB) P_os	1T モデル (GB) P_os+g	1T モデル (GB) P_os+g+p
1	120	120	120	2048	2048	2048	16000	16000	16000
4	52.5	41.3	30	896	704	512	7000	5500	4000
16	35.6	21.6	7.5	608	368	128	4750	2875	1000
64	31.4	16.6	1.88	536	284	32	4187	2218	250
256	30.4	15.4	0.47	518	263	8	4046	2054	62.5
1024	30.1	15.1	0.12	513	257	2	4011	2013	15.6

• ZeRO は従来の DP/MP の制限を超えるモデルのデータ並列訓練を可能にし、パラメータと勾配に対しては N_d にほぼ比例するメモリ削減、オプティマイザ状態には段階に応じて約 4x–8x の削減を実現する。
• 3 段階（P_os、P_g、P_p）によりデバイスごとのメモリフットプリントを大幅に削減でき、巨大な GPU クラスター上で trillion-parameter モデルを実現可能。
• ZeRO-DP を 64-way DP と組み合わせると 7.5B モデルの訓練が可能となり、P_os+g+p では 7.5B モデルのメモリが概ね 1.88 GB 程度に縮小され、より大規模なスケールを可能にする。
• ZeRO-R は残差アクティベーションメモリをさらに削減し、効率のために一定サイズのバッファを使用し、断片化によるOOM を緩和するための即時メモリデフラグメンテーションを実行する。
• 実装結果として、400 台の Nvidia V100 GPU で 100B パラメータ級モデルを訓練し、1 GPU あたり 38 TFlops を超え、総合性能が 15 ペタフロップ超を達成。示されたデータは、当時の最先端より最大で最大モデルサイズが約 8 倍、実現可能な性能が約 10 倍向上することを示唆。
• ZeRO は DeepSpeed ライブラリの一部として公開され、大規模モデル訓練を民主化し、さらに分割（P_os+g+p）により1兆パラメータへとスケールすることを目指す。