マイクロソフトのZeRO-2がAIトレーニングを10倍に高速化

マイクロソフトは、クラスタサイズに応じて超線形にスケーリングする分散型ディープラーニング最適化アルゴリズムであるZero Redundancy Optimizer version 2 (ZeRO-2) をオープンソース化した。ZeRO-2を使用することによって、Microsoftは1,000億パラメータの自然言語処理（NLP）モデルを従来の分散学習技術の10倍の速さで学習した。

同社のプログラムマネージャーのRangan MajumderとエンジニアであるJunhua Wangはブログ記事で、このアルゴリズムと実験について説明している。ZeRO-2は、ディープラーニング学習最適化のためのMicrosoftのオープンソースDeepSpeedライブラリの一部。ZeRO-2は、トレーニング中のメモリ消費を最適化し、1700億個ものパラメータを持つモデルの分散トレーニングを可能にする。また、このアルゴリズムは分散クラスタ内のワーカーノード間の通信を削減し、トレーニング時間を最大10倍に短縮する超並列高速化を実現している。1,024個のGPUを搭載したクラスタ上でZeRO-2を使用したDeepSpeedチームは、BERT自然言語モデルのトレーニングに44分という記録的な時間を達成し、NVIDIAの結果と比較して30%以上改善した。

最近のNLP研究の傾向として、大規模なデータセットで訓練された大規模なモデルの精度が向上している。OpenAIは、モデルの精度がモデルサイズと力法則の関係を持つことを示す一連の「スケーリングの法則」を提案しており、最近では1750億個のパラメータを持つGPT-3モデルを作成してこの考えをテストした。これらのモデルは単純に大きすぎて1つのGPUのメモリに収まりきらないため，これらのモデルを学習するには，マシンのクラスタと，クラスタ全体にパラメータを分散させるモデル並列学習技術が必要となる。GPipeやNVIDIAのMegatronなど、効率的なモデル並列化を実装したオープンソースのフレームワークがいくつかあるが、これらはクラスタノード間の通信のオーバーヘッドのためにサブリニアな高速化が必要であり、フレームワークを使用するにはモデルのリファクタリングが必要になることがよくある。

ZeRO-2は3つの戦略を使用してトレーニングに必要なメモリを削減する。メモリ要件の削減、活性化関数のレイヤーのCPUへのオフロード、メモリの断片化の削減だ。ZeRO-2は、並列プロセス間でグラデーションとパラメータを分割することで、メモリ要件を最大8倍まで削減できる。活性化関数の値は前方のトレーニングパスから保存され、後で後方のパスで使用されるが、ZeRO-2は一時的にGPUのメモリからホストCPUのメモリに移動する。最後に、メモリが使用可能な場合でも、使用可能なメモリが連続していないと、メモリ割り当てに失敗することがある。ZeRO-2は、アクティベーションやグラデーションなどの一時的な用途のために、連続したメモリチャンクを事前に割り当てておくことで、断片化を減らすことができる。

これらのメモリ最適化により、必要なモデルの並列性、ひいては必要なノード間通信のオーバーヘッドを削減することができ、データ並列トレーニングと併用することで、超直線的な高速化を実現します。DeepSpeedチームは、トレーニングに使用するGPUの数を増やすことで、GPUあたりのテラフロップ数で測定される全体的なスループットが向上することを発見した。大規模なNLPモデルの実験では、「20億から130億のパラメータを持つモデルサイズの場合」、V100 GPUあたり37テラフロップスの平均スループットが観測された。特定のモデルサイズとクラスタに対して、ZeRO-2はベースラインのMegatronアプローチよりも最大10倍高速にモデルを訓練することができたとチームは指摘している。1,024台のV100 GPUを使用して、BERTモデルを44分で学習し、NVIDIAが1,472台のV100 GPUを使用して47分というそれまでの記録を更新した。

DeepSpeedチームのメンバーであるJeff Rasleyは、Hacker Newsのディスカッションに参加し、コミュニティからの質問に答えた。Rasleyは、DeepSpeedにはMicrosoftに「何百人もの社内ユーザー」がおり、本番で稼働しているモデルのトレーニングにDeepSpeedを使用していると指摘した。TPUのサポートについて質問されると、Rasleyは答えた。

ZeRO技術は、TPUやクラスタ設定の任意のアクセラレータと互換性がありますが、TPUではテストしていません。TPUでDeepSpeedを動作させるためには、小さなリファクタリングが必要になるでしょう。我々はまだそれらをサポートするための内部計画を持っていませんが、もちろんコミュニティからの貢献に対して完全にオープンです。

DeepSpeedのライブラリはGitHubで公開されている。

Photo by Microsoft