Microsoftは最近、ZeRO-3オフロードをオープンソース化した。これは、非常に大規模な深層学習モデルをトレーニングしつつも、メモリ効率を改善する、DeepSpeed AIトレーニングライブラリの拡張である。ZeRO-3オフロードを使用すると、ユーザは、単一のGPUで最大400億のパラメーター、512のGPUで2兆を超えるパラメーターを使用してモデルをトレーニングできる。
DeepSpeedチームは、最近のブログ投稿でリリースの機能と利点の概要を提供した。ZeRO-3オフロードは、PyTorchフレームワーク上に構築された深層学習モデルの分散トレーニングのメモリ効率を高めている。そして、複数のGPUにまたがる超線形スケールを提供する。一部のデータのストレージをGPUからCPUにオフロードすることで、GPUごとにより大きなサイズのモデルをトレーニングできる。これによって、単一のGPUで最大400億のパラメーターのモデルサイズを実現できる。トレーニングにDeepSpeedフレームワークを採用する場合、モデルコードのリファクタリングは最小限で済む。そのため、現在のユーザは構成ファイルを変更することで新機能を利用できる。DeepSpeedチームによると、このリリースは「効率的な大規模DLトレーニングをすべての人が利用できるようにすることでAIを民主化するという私たちの継続的な目標に向けてのものである」。
ディープラーニングモデルの最近の傾向は、モデルパラメーターの数が指数関数的に増加していることである。2018年には数億であったパラメーターが2020年には数百兆になる。より大きなモデルではパフォーマンスが向上するが、モデルのトレーニングには複数のGPUアクセラレーターが必要となる。そして、モデルのパラメーターのさまざまなサブセットに個別のGPUを使用するモデル並列処理などの分散トレーニング手法が必要となる。ただし、PyTorchフレームワークはこの手法をサポートしているが、通常はモデルトレーニングコードを変更する必要がある。
Microsoftは、2020年初頭にDeepSpeedライブラリとZero Redundancy Optimizer（ZeRO）を最初にリリースした。MicrosoftのProject Turingは、ライブラリを使用してTuring Natural Language Generation（T-NLG）モデルをトレーニングした。これは、170億のパラメータで当時最大の既知の言語モデルであった。ZeROは、最適化の3つの段階を定義している。
- ZeRO-1: オプティマイザー状態のパーティションニング
- ZeRO-2: ZeRO-1とグラデーションパーティショニング
- ZeRO-3: ZeRO-2とパラメータパーティショニング
DeepSpeedの最初のリリースには、最初のステージであるZeRO-1のみが含まれていた。その後のリリースには、ZeRO-2とZeRO-Offloadが含まれていた。これは、トレーニングマシンのGPUからCPUにデータと計算を「オフロード」するためのスキームである。これにより、GPUメモリが解放され、単一のGPUでより大きなモデルを管理できるようになる。
新しいリリースでは、ZeROの最適化とオフロードに関する3つの段階すべてを組み合わせたZeRO-3オフロードが実装されている。これにより、ZeRO-2でサポートされているモデルの3倍の大きさのモデルのシングルGPUトレーニングが可能になる。これは、Microsoftによると、以前の最先端技術であった。より大きなモデルの場合、2テラのパラメーターモデルを512のGPUのクラスターでトレーニングできる。これは、以前の最先端の要件である1600GPUと比較した場合である。このシステムは、より高いスループットもサポートできる。標準のPyTorchトレーニングでは30テラflopsであるのに対し、GPUあたり50テラflopsである。
大規模なモデルが一般的になるにつれて、分散トレーニングが研究開発の主要な焦点になった。PyTorchの主なライバルであるGoogleのTensorFlowフレームワークには、大規模なモデルをトレーニングするためのMesh TensorFlowと呼ばれる別のライブラリがある。PyTorchの作成者であるFacebookは、FairScaleと呼ばれるPyTorchの分散トレーニング拡張機能を作成した。事前トレーニングされたAIモデルの人気のあるソースであるHuggingfaceは、大規模モデルのトレーニングにDeepSpeedとFairScaleの両方を採用している。Hacker Newsでのディスカッションで、あるユーザは、FairScaleがDeepSpeedのアイデアのいくつかも実装していると述べました。別の人は、ZeROの最適化がすべての種類のAIモデルでうまく機能するわけではない可能性があることを指摘した。
これは主にトランスフォーマーモデルに適用できます。コンピュータービジョンに取り組んでいる場合、この論文のアイデアは異質です。
DeepSpeedのソースコードはGitHubで入手できる。