130 億個のパラメータを持つモデルをトレーニングするには GPU がいくつ必要ですか?マイクロソフト: 1つで十分

今日のモデルには数千億、あるいは数兆ものパラメータがあります。一般の人がモデルをトレーニングできない場合、どうすればいいのでしょうか?

少し前に、Google は 1.6 兆個のパラメータを持つ言語モデルである Switch Transformer をリリースし、GPT-3 が記録したパラメータ記録 (1,750 億) を新たな高さに押し上げました。こうした大規模モデルの出現により、一般の研究者はますます絶望的になっています。「お金」と大量の GPU がなければ AI 研究を行うことは不可能なのでしょうか?

このような状況で、一部の研究者は、これらの大規模モデルのトレーニングをより現実的にするにはどうすればよいかと考え始めました。言い換えれば、より少ないカードを使用してより大きなモデルをトレーニングするにはどうすればよいでしょうか?

この問題を解決するために、マイクロソフトとカリフォルニア大学マーセド校の研究者らは、「ZeRO-Offload」と呼ばれる異種ディープラーニングトレーニング技術を提案した。これは、単一のGPU上で130億のパラメータを持つディープラーニングモデルをトレーニングすることができ、一般の研究者も大規模モデルをトレーニングできるようになる。 Pytorch などの一般的なフレームワークと比較すると、 ZeRO-Offload では、データ サイエンティストがモデルに変更を加えたり計算効率を犠牲にしたりすることなく、トレーニング可能なモデルのサイズが 10 倍に増加します。

論文リンク: https://arxiv.org/pdf/2101.06840.pdf

ZeRO-Offload は、データと計算を CPU にオフロードすることで、大規模なモデルのトレーニングを可能にします。計算効率を低下させないために、GPU との間のデータ転送を最小限に抑え、GPU メモリを可能な限り節約しながら CPU 計算時間を短縮するように設計されています。したがって、100 億のパラメータを持つモデルの場合、ZeRO-Offload は単一の NVIDIA V100 GPU で 40 TFlops/GPU を達成できます。対照的に、PyTorch を使用して 14 億のパラメータを持つモデルをトレーニングすると、30TFlops しか達成されません。これは、メモリ不足に陥ることなくトレーニングできる最大のモデルです。 ZeRO-Offload は、マルチ GPU セットアップに拡張して線形の高速化を実現することもできます。最大 128 個の GPU でほぼ線形の高速化が実現されます。

さらに、ZeRO-Offload をモデル並列処理と組み合わせて使用​​することで、単一の DGX-2 ボックス AI サーバー上で 700 億を超えるパラメータを持つモデルをトレーニングできます。この量のパラメータにより、モデルの並列処理のみを使用する場合と比較して 4.5 倍のスケール改善が実現します。

以下では、Medium ブロガーの LORENZ KUHN 氏のブログ投稿と併せて、この論文を詳しく見ていきます。

ZeRO-Offloadとは何ですか?

ZeRO-Offload は、データと計算を GPU から CPU にオフロードすることで、ニューラル ネットワークのトレーニング中に GPU メモリの使用量を削減する方法です。これにより、トレーニングのスループットが向上し、CPU でデータを移動して計算を実行することによる速度低下を回避できます。

ZeRO-offload を使用すると、単一の GPU でも同じハードウェアを使用して 10 倍の規模のモデルをトレーニングできます。たとえば、32GB RAM を搭載した V100 GPU で数百億のパラメータを使用して GPT-2 をトレーニングします。

さらに、ZeRO-offload は、マルチ GPU セットアップでほぼ線形のスケーラビリティを実現します。

研究者にとって、ZeRO-offload は次のような場合に適しています。

ZeRO-offload ではより大きなバッチ サイズでトレーニングできるため、より大きなモデルをトレーニングしたい場合、または現在のモデルをより速くトレーニングしたい場合。

PyTorch を使用しており、Microsoft の DeepSpeed ライブラリを使用する意思/能力がある (ZeRO-offload の他の実装はまだ利用できません)、または公式の実装に基づいて独自に適応させることを試みることができます。

現在のバージョンの ZeRO-Offload は Adam の混合精度トレーニングで使用する必要があるなど、いくつかのモデリング制限を受け入れる必要があります。

使い方は？

ZeRO-Offload は Microsoft の DeepSpeed ライブラリに実装されており、公式の実装アドレスは https://github.com/microsoft/DeepSpeed/blob/6e65c2cc084ecfc393c67a2f64639e8d08d325f6/deepspeed/runtime/zero/stage2.py です。

DeepSpeed で設定してしまえば、ZeRO-Offload を使用するのに特別な作業はほとんど必要なく、フラグと構成ファイルをいくつか変更するだけです。

現在、Hugging Face のトランスフォーマー ライブラリは、DeepSpeed と試験的に統合されています。使用方法とベンチマーク結果については、https://huggingface.co/blog/zero-deepspeed-fairscale を参照してください。

Facebook Research の fairscale には ZeRO の部分的な実装があり、ZeRO-Offload は ZeRO のマルチ GPU メモリ最適化手法に基づいて構築されています。 CPU オフロードは現在サポートされていません。

ZeRO-Offloadの仕組み

ZeRO-Offload は、Zero Redundancy Optimizer (ZeRO) 上に構築されています。 ZeRO は、データ並列およびモデル並列トレーニングにおけるメモリ最適化のために 2020 年 2 月に Microsoft が提案した 1 兆規模のモデル パラメータ トレーニング手法であり、勾配、パラメータ、およびオプティマイザーの状態が冗長性なしでマルチ GPU メモリに分散されます。これにより、GPU 間の通信オーバーヘッドが比較的低くなります。

標準のデータ並列ベースラインと比較した、3 つのフェーズにおける ZeRO のメモリと通信の節約。

ZeROを振り返ってみましょう:

データ並列性とモデル並列性の問題を解決するために、ZeRO は、オプティマイザー状態分割、勾配分割、およびパラメータ分割という 3 段階の最適化手法を提供します。3 つの段階は順番に実装されます。

オプティマイザ分割状態では、ZeRO はメモリ使用量を 3/4 削減し、通信量はデータ並列処理と同じになります。

勾配分割を追加: 通信量とデータ並列性は同じまま、メモリ使用量が 7/8 削減されました。

パラメータ分割を追加します。メモリ削減はデータの並列処理に比例します。たとえば、64 個の GPU に分割すると、メモリを 1/64 に削減できます。通信量が50%増加します。

昨年 9 月のブログ投稿で、Microsoft は ZeRO-Offload を次のように紹介しました。

ZeRO-Offload は、ZeRO-2 の最適化状態と勾配分割を継承します。ただし、ZeRO-2 とは異なり、ZeRO-Offload は各 GPU でオプティマイザーの状態と勾配分割を維持せず、両方をホスト CPU メモリにオフロードします。トレーニング フェーズ全体を通じて、オプティマイザーの状態は CPU メモリに保持されます。勾配は、バックプロパゲーション中に GPU 上で Reduce-Scatter を使用して計算および平均化され、その後、各データ並列スレッドは、その分割に属する勾配の平均を CPU メモリにオフロードし (下の図の g offload を参照)、残りを破棄します。勾配が CPU に到達すると、各データ並列スレッドは CPU 上で直接並列に分割されたオプティマイザー状態を更新します (下の図の p 更新を参照)。

その後、パラメータ分割は GPU に戻され、更新されたすべてのパラメータを収集するために GPU 上で all-gather 操作が実行されます (下の図の g swap を参照)。 ZeRO-Offload は、個別の CUDA ストリームも利用して、通信と計算の重複を排除し、トレーニングの効率を最大化します。

ZeRO-Offload の概要。

ZeRO-Offload は、Adam を使用した混合精度トレーニング用に設計されていることに注意してください。つまり、ZeRO-Offload の現在のバージョンでは、Adam の最適化バージョンである DeepCPUAdam が使用されています。主な理由は、CPU 計算がプロセス全体のボトルネックになるのを避けるためです。 DeepCPUAdam は Adam PyTorch 実装よりも 6 倍高速です。

実験結果

最後に、ZeRO-Offload の論文に記載されているいくつかの実験結果を見てみましょう。

下の図 7 は、ZeRO-Offload テクノロジーを使用して 1、4、または 16 個の GPU (1 個の DGX-2) でトレーニングできる最大のモデルを示しています。

下の図 11 は、GPU の数が増えるにつれて GPU あたりのスループットがどのように変化するかを示しています。 GPU の数が徐々に 128 まで増加すると、ZeRO-Offload はスループットのほぼ線形加速を実現できることがわかります。

下の図 8 は、PyTorch、L2L、ZeRO-Offload を使用して達成された GPU あたりのスループットの違いを示しています。 ZeRO-Offload で達成される GPU あたりのスループットは、L2L で達成されるスループットよりも平均で 14% 高く (最大で 22% 高い) ことがわかります。