[論文レビュー] ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

Samyam Rajbhandari, Olatunji Ruwase+3 more|arXiv (Cornell University)|Apr 16, 2021

Advanced Neural Network Applications参考文献 30被引用数 57

ひとこと要約

ZeRO-Infinity はモデル状態と活性化を CPU/NVMe にオフロードし、メモリ中心のタイル処理で限られた GPU リソース上で兆パラメータのモデルを訓練する。高いスループットとアクセス性を実現する。

ABSTRACT

In the last three years, the largest dense deep learning models have grown over 1000x to reach hundreds of billions of parameters, while the GPU memory has only grown by 5x (16 GB to 80 GB). Therefore, the growth in model scale has been supported primarily though system innovations that allow large models to fit in the aggregate GPU memory of multiple GPUs. However, we are getting close to the GPU memory wall. It requires 800 NVIDIA V100 GPUs just to fit a trillion parameter model for training, and such clusters are simply out of reach for most data scientists. In addition, training models at that scale requires complex combinations of parallelism techniques that puts a big burden on the data scientists to refactor their model. In this paper we present ZeRO-Infinity, a novel heterogeneous system technology that leverages GPU, CPU, and NVMe memory to allow for unprecedented model scale on limited resources without requiring model code refactoring. At the same time it achieves excellent training throughput and scalability, unencumbered by the limited CPU or NVMe bandwidth. ZeRO-Infinity can fit models with tens and even hundreds of trillions of parameters for training on current generation GPU clusters. It can be used to fine-tune trillion parameter models on a single NVIDIA DGX-2 node, making large models more accessible. In terms of training throughput and scalability, it sustains over 25 petaflops on 512 NVIDIA V100 GPUs(40% of peak), while also demonstrating super linear scalability. An open source implementation of ZeRO-Infinity is available through DeepSpeed, a deep learning optimization library that makes distributed training easy, efficient, and effective.

研究の動機と目的

• 極端規模のディープラーニングにおける GPU メモリの壁の課題を説明し、リファクタリングなしで到達可能なモデルサイズを増やすシステムの動機付けを行う。
• 既存のクラスター上に巨大なモデルを適合させるため、ZeRO-Infinity の異種メモリオフロードとタイル処理技術を紹介する。
• 兆パラメータ規模を含む非常に大規模なモデルでの訓練効率とスケーラビリティを示す。
• モデルのリファクタリングや複雑な並列処理の必要性を減らすまたは排除するための使いやすさの改善を評価する。
• オープンソース実装の詳細を提供し、将来のハードウェア設計への影響について議論する。

提案手法

• 同時に GPU、CPU、NVMe メモリを活用する infinity オフロードエンジンを提案する。
• 作業メモリを削減するために大規模オペレータをタイルに分割するメモリ中心のタイル処理を導入する。
• デバイス間の総メモリ帯域幅を活用する帯域幅中心のパーティショニングを定義する。
• 異種メモリ間での通信と計算を重ね合わせて隠すオーバーラップ中心の設計を開発する。
• 任意のアーキテクチャを訓練しつつモデルのリファクタリングを避けるための容易さを重視した実装を提供する。
• Transformerベースのモデルを用いた大規模モデル訓練のためのメモリと帯域幅要件を特徴づける。

実験結果

リサーチクエスチョン

• RQ1 実用的な GPU クラスター上で数百兆パラメータのモデルを訓練するために、GPU メモリの壁をどう打破できるか。
• RQ2 CPUおよびNVMeメモリを活用して、巨大なGPUクラスターを持たない研究者にも大規模モデルを利用可能にできるか。
• RQ3 モデルコードのリファクタリングや複雑な3D並列化を導入せずに兆パラメータモデルを訓練できるか。
• RQ4 異種メモリを用いた安定した効率的な訓練を維持するために必要なメモリと帯域幅の特性は何か。
• RQ5 現在のハードウェア上で、スループット、スケーラビリティ、アクセス性の点で ZeRO-Infinity はどのように機能するか？

主な発見

• ZeRO-Infinity は 32 trillion parameters のモデルを 32 NVIDIA DGX-2 ノード（512 GPUs）で訓練できる。
• 同じハードウェア上でスループットが 25 ペタフロップスを超えることを維持する（ピークの 40%）。
• 兆パラメータモデルで超線形スケーラビリティを示す。
• 単一の DGX-2 ノードでモデル並列化やリファクタリングなしに、兆パラメータまでのファインチューニングを可能にする。
• メモリ中心のタイル処理は演算子ごとの作業メモリを削減し、モデル並列化なしで極めて大きな層を適合させる。
• ZeRO-Infinity のオープンソース実装は DeepSpeed にあります。