Microsoftは、クラスタサイズに対してスーパーリニアにスケールする、分散ディープラーニング最適化アルゴリズムZero Redundancy Optimizerのバージョン2(ZeRO-2)をオープンソースとして公開した。同社はZeRO-2を使用して、1,000億パラメータの自然言語処理(NLP)モデルを、従来の分散ラーニング技術の10倍の速度でトレーニングすることに成功している。
ブログ記事の中では、プログラムマネージャのRangan Majumder氏と特別エンジニア(distinguished engineer)のJunhua Wang氏が、アルゴリズムの内容と同社が行った実験について書いている。ZeRO-2は、MicrosoftのオープンソースであるDeepSpeedライブラリの一部で、ディープラーニングにおけるトレーニングの最適化を目的とする。ZeRO-2はトレーニング中のメモリ消費を最適化することで、1,700億パラメータまでのモデルの分散トレーニングを可能にしている。分散クラスタ内のワーカノード間の通信も削減されており、スーパーリニアな並列速度の向上を実現し、トレーニング時間を10分の1まで削減する。1,024GPUのクラスタ上でZeRO-2を使用したDeepSpeedチームは、BERT自然言語モデルのトレーニングで44分という、NVIDIAの記録を30パーセント上回る新記録を達成した。
NLP研究最近の傾向として、大規模なデータセットでトレーニングされた大規模モデルによる精度の改善がある。OpenAIの提案する"スケーリング法則"では、モデルの精度はモデルのサイズとべき乗の関係にあることが示されている。この考え方は先日、17,500億のパラメータを持つGPT-3モデルの作成によってテストされた。このようなモデルは、単一GPUのメモリに収めるには大き過ぎるという単純な理由から、トレーニング時のマシンクラスタと、パラメータをクラスタ全体に分散するモデルパラレルの技術が不可欠である。GPipeやNVIDIAのMegatronなど、効率的なモデル並列性を実装したオープンソースのフレームワークはいくつか公開されているが、クラスタノード間の通信によるオーバーヘッドのために完全にリニアな速度向上が期待できなかったり、フレームワークを使用するためにモデルのリファクタリングが必要になったりするなどの問題がある。
ZeRO-2は、モデル状態メモリの必要量の削減、レイヤアクティベーションのCPUへのオフロード、メモリフラグメンテーションの低減、という3つの方法を使って、トレーニングに必要なメモリを削減する。モデル状態メモリの削減に関しては、並列プロセス全体の勾配(gradient)とパラメータをパーティション化することで、メモリ要求量を8分の1にまで削減した。レイヤアクティベーション値はフォワードトレーニングパスから保存され、後のバックワードパス内で使用されるものだが、ZeRO-2ではそれを一時的に、GPUのメモリからホストCPUのメモリに移動させている。最後の問題は、十分なメモリが利用可能であっても、それが連続的でない時には、アロケーションが失敗する場合のあることだ。ZeRO-2では、アクティベーションや勾配などで一時的に使用するために、連続したメモリチャンクを事前にアロケートしておくことで、フラグメンテーションの発生を抑制している。
このようなメモリ最適化がモデル並列性の必要性、さらにはノード間通信のオーバーヘッドの必然性を低減することにより、データ並列トレーニングと併用されることでスーパーリニアな速度向上を実現するのだ。さらには、トレーニングに使用するGPUの数を増やすことにより、GPU単位のテラフロップスで測定した全体のスループットが向上する、ということも分かった。大規模なNLPモデルを使用した実験では、"20億から130億の範囲のパラメータを持つモデルサイズにおいて"、V100 GPUあたり37テラフロップスの平均スループットを確認している。開発チームは、1,024のV100 GPUを使ってBERTモデルをトレーニングし、1,472のV100 GPUを使って47分というNVIDIAのそれまでの記録を更新した。
チームメンバのJeff Rasley氏はHacker Newsでの議論に参加して、コミュニティからの質問に答えている。氏によると、DeepSpeedにはMicrosoftに"数百人の社内ユーザ"があり、実際に運用されているモデルのトレーニングに使用されている。TPUのサポートに関する質問には次のように答えている。
ZeROテクノロジは、TPUを含むあらゆるクラスタ設定のアクセラレータに対応可能ですが、TPUで試験したことはありません。TPUでDeepSpeedが動くようにするには、多少のリファクタリングが必要だと思われます。今のところサポートする計画はありませんが、コミュニティからのコントリビューションはもちろん大歓迎します。
DeepSpeedライブラリはGitHubで公開されている。