[論文レビュー] ZeRO-Offload: Democratizing Billion-Scale Model Training

Jie Ren, Samyam Rajbhandari+6 more|arXiv (Cornell University)|Jan 18, 2021

Advanced Neural Network Applications参考文献 27被引用数 61

ひとこと要約

ZeRO-Offload は勾配、オプティマイザ状態、そしてオプティマイザ計算を CPU にオフロードして、1 台の GPU で最大 13B パラメータのモデルを訓練可能にし、スケーラブルな性能と PyTorch とのシームレスな統合を実現します。

ABSTRACT

Large-scale model training has been a playing ground for a limited few requiring complex model refactoring and access to prohibitively expensive GPU clusters. ZeRO-Offload changes the large model training landscape by making large model training accessible to nearly everyone. It can train models with over 13 billion parameters on a single GPU, a 10x increase in size compared to popular framework such as PyTorch, and it does so without requiring any model change from the data scientists or sacrificing computational efficiency. ZeRO-Offload enables large model training by offloading data and compute to CPU. To preserve compute efficiency, it is designed to minimize the data movement to/from GPU, and reduce CPU compute time while maximizing memory savings on GPU. As a result, ZeRO-Offload can achieve 40 TFlops/GPU on a single NVIDIA V100 GPU for 10B parameter model compared to 30TF using PyTorch alone for a 1.4B parameter model, the largest that can be trained without running out of memory. ZeRO-Offload is also designed to scale on multiple-GPUs when available, offering near linear speedup on up to 128 GPUs. Additionally, it can work together with model parallelism to train models with over 70 billion parameters on a single DGX-2 box, a 4.5x increase in model size compared to using model parallelism alone. By combining compute and memory efficiency with ease-of-use, ZeRO-Offload democratizes large-scale model training making it accessible to even data scientists with access to just a single GPU.

研究の動機と目的

• アクセス可能な十億パラメータ級モデル訓練の必要性を動機づけ、ハードウェアの障壁を低減する。
• GPU メモリ削減を最大化しつつ CPU 計算と CPU-GPU 通信を最小化するユニークなオフロード戦略を提案する。
• ZeRO-powered データ並列性と可能なモデル並列性を用いて、単一 GPU から最大 128 GPU までのスケーラビリティをデモンストレーションする。
• 高速な CPU 実行（高速 Adam）とスループットと精度を維持するスケジューリング技法を提供する。

提案手法

• モデル訓練はデータフローグラフとして表現され、計算・通信・メモリ節約を最適化するために CPU と GPU の間で分割される。
• オフロード戦略では FP16 パラメータを GPU に保持しつつ、FP16 勾配と FP32 オプティマイザ状態を CPU にオフロードし、パラメータ更新を CPU で行う。
• 分割分析は最小 CPU 計算量と最小 CPU-GPU 通信を示し、最適なオフロード決定につながる。
• オフロードを ZeRO-stage-2 データ並列性と組み合わせて、最大 128 GPU までのほぼ線形スケーリングと、より大きなモデル向けのモデル並列性との互換性を実現する。
• SIMD、ループ展開、マルチスレッドを用いた最適化された CPU Adam 実装を使用し、CPU と GPU 作業を重ね合わせるための 1 ステップ遅延パラメータ更新を追加する。
• 遅延パラメータ更新（DPU）オプションは CPU 更新と GPU 計算の重ね合わせを可能にし、精度低下なしにスループットを向上させる。

実験結果

リサーチクエスチョン

• RQ1 ZeRO-Offload は、1GPU のメモリに収まらないモデルの訓練を、効率を損なうことなく可能にするか。
• RQ2 CPU へのオフロードは ZeRO データ並列性とどのように相互作用して多数の GPU へスケールさせるのか。
• RQ3 十億パラメータ級モデルを単一/複数の GPU で訓練する際の実用的なスループットとメモリ節約はどの程度か。
• RQ4 最適化された CPU Adam と DPU は、性能を向上させつつモデルの収束性を維持できるか。

主な発見

• 単一の NVIDIA V100 GPU 40 TFLOPS で 13B パラメータのモデル訓練が可能になり、オフロードなしだと適合するのは 1.2B。
• ZeRO-powered データ並列性と組み合わせると、最大 128 GPU でほぼ直線的なスケーラビリティ。
• DGX-2 ノード上で、モデル並列性と組み合わせることで最大で 70B パラメータの訓練が可能。
• CPU-Adam の最適化は標準の PyTorch Adam 実装に対して 6 倍以上の速度アップを生み、遅延パラメータ更新を用いるとエンドツーエンドのスループットが最大で 1.5 倍向上。
• このアプローチは、モデルリファクタリングを必要とせず、最小限の CPU 計算オーバーヘッドと制御された通信で、10 倍大きいモデルの訓練を可能にする。
• オフロード戦略は、CPU 計算と CPU-GPU 通信を最小化しつつ GPU メモリ節約を最大化するという定義された目的の下で、唯一無二かつ最適である。