Microsoftは最近、ZeRO-Infinityを発表した。これは、オープンソースのDeepSpeed AIトレーニングライブラリへの追加である。このライブラリは、非常に大規模なディープラーニングモデルをトレーニングするためのメモリ使用を最適化するものである。MicrosoftはZeRO-Infinityを使って、32のGPUのクラスター上で32兆のパラメーターを持つモデルをトレーニングし、1つのGPUで1兆のパラメーターモデルの微調整のデモを実施した。
DeepSpeedチームは、最近のブログ投稿で新機能について説明した。ZeRO-Infinityは、メモリ最適化手法のZero Redundancy Optimizer (ZeRO) ファミリーの最新版である。ZeRO-Infinityでは、大規模な深層学習モデルをトレーニングする際のメモリと帯域幅の制約に対処するためのいくつかの新しい戦略を導入している。例えば、CPUとNon-Volatile Memory Express (NVMe) メモリを活用するための新しいオフロードエンジン、モデルの並列化なしで大規模なオペレーターを処理するためのメモリ中心のタイリング、帯域幅コストを削減するための帯域幅中心のパーティショニング、データ通信をスケジュールするためのオーバーラップ中心の設計などである。DeepSpeedチームは次のように述べている。
ZeRO-Infinityの改善によって、システムは、GPUのメモリの壁を超えて、数十兆のパラメーターでモデルをトレーニングする能力を提供します。これは、最先端のシステムがサポートできるよりも桁違いに大きいものです。また、100兆パラメーターのモデルをトレーニングするための有望な道のりも示されています。
深層学習研究の最近の傾向は、より大きなモデルをより多くのデータでトレーニングすることである。最大のモデルのいくつかは特定のタスクで超人的なパフォーマンスを達成している。ただし、これらのモデルをトレーニングするには、GPUの大規模で高価なクラスターが必要となる。多くの場合、モデル開発者は転移学習を使って、事前トレーニングに必要だったコンピューティングリソースのほんの一部を使用して、事前トレーニング済みの大規模モデルを微調整できる。それでも、GPT-3などの非常に大きなモデルは大きすぎて、1台のマシンで微調整できない。どちらのシナリオでも、分散トレーニングフレームワークを活用するためにコードのリファクタリングが必要になることが多い。
これらの問題に対処するために、Microsoftは2020年の初めに、AI at Scaleプログラムの一環として、DeepSpeedライブラリとZero Redundancy Optimizer (ZeRO) を初リリースした。ZeROは3ステージで改善され、各ステージでモデル状態のパーティショニングが追加され、GPUからトレーニングマシンのCPUにデータと計算を「オフロード」する機能が追加された。ステージ3は今年の初めにリリースされ、1台のマシンで最大400億個のパラメーター、512個のGPUのクラスターで2兆個を超えるパラメーターでモデルをトレーニングできるようになった。
ZeROの最新版であるZeRO-Infinityでは、大規模モデルのトレーニングにおける2つのボトルネック (メモリサイズとメモリ帯域幅) に対処するための新しいスキームを取り入れている。インフィニティオフロードエンジンによって、CPUとNVMeメモリを使って、モデルパラメーターとアクティベーションの保存に使用できるメモリの量を増やすことができる。前世代のZeROとは異なり、インフィニティエンジンはモデル全体をこれらの場所にオフロードできる。メモリセントリックタイリングはもう 1 つの新しい手法である。これによって、大きなモデル レイヤーを順次実行できる小さな「タイル」に分割することで、メモリフットプリントを削減できる。これにより、モデルの並列性を必要とせずに大規模なモデルをトレーニングできる。帯域幅の問題を処理するために、ZeRO-Infinityは帯域幅中心のパーティショニングとオーバーラップエンジンを導入している。このパーティショニングは、複数のデータ並列プロセス間でモデルパラメーターを分割するものである。また、オーバーラップエンジンは、 NVMe-to-CPU、CPU-to-GPU、GPU-to-GPUの通信を同時に実行するものである。
チームは、ZeRO-Infinityのスケーリング能力を検証するためにいくつかの実験を行った。さまざまなサイズの「GPTのような」Transformer モデルをトレーニングした。最先端の3D並列化フレームワークと比較して、ZeRO-Infinity は同じコンピューティングハードウェアを使用して40倍大きいモデルを処理した。ZeROの以前のバージョンと比較して、新しいバージョンは64GPUクラスターで2倍のスピードアップを達成した。1Tパラメータモデルをトレーニングする場合、ZeRO-Infinityは64GPUから512GPUまでのさまざまなクラスターサイズで超線形的にスケーリングした。
メモリ最適化のZeROファミリを含むDeepSpeedライブラリは、PyTorch深層学習フレームワーク用に作成されており、他のいくつかのPyTorchベースのプロジェクトで採用されている。事前トレーニング済みAIモデルの人気のあるソースであるHuggingFaceは、新しいZeRO-Infinityリリースと統合されている。PyTorchの分散トレーニングラッパーであるPyTorch LightningもDeepSpeedとZeRoの最初の3ステージを採用している。大規模なPyTorchモデルをトレーニングするためのFacebookのFairScaleライブラリにも、いくつかのZeROテクノロジが含まれている。
Redditでのディスカッションで、あるコメント投稿者は、DeepSpeedライブラリを「信じられないほど価値がある」と述べた。別の人は次のように指摘した。
ただし、これらの手法 (ゼロオフロード、ゼロインフィニティ) は、通常、より大きなモデルのトレーニングにはあまり役に立ちません。非常に大きなモデルを最初からトレーニングする場合、通常、メモリはボトルネックではありません。計算がボトルネックとなります。それでも、これらの技術は微調整に非常に役立ちます。
DeepSpeedライブラリコードはGitHubで入手できる。