マイクロソフトのDeepSpeedは、少ないGPUで数兆パラメータのAIモデルをトレーニングできる

マイクロソフトはDeepSpeedライブラリのアップデート版をリリースした。このライブラリでは、何兆個ものパラメータを含むAIモデルを学習するための新しいアプローチが導入されている。同社は、3D並列化と呼ばれるこの技術は、ワークロード要件の多様なニーズに適応し、スケーリング効率のバランスを取りながら、非常に大きなモデルに電力を供給すると主張している。

数十億個のパラメータを持つ単一の大規模AIモデルは、さまざまな困難な領域で大きな進歩を遂げてきた。研究によると、言語、文法、知識、概念、文脈のニュアンスを吸収し、スピーチを要約したり、ライブゲームのチャットの内容を精査したり、複雑な法律文書を解析したり、GitHubを精査してコードを生成したりすることができるため、AIモデルの性能が高いことがわかっている。しかし、モデルを訓練するには膨大な計算リソースが必要だ。2018年のOpenAIの分析によると、2012年から2018年にかけて、最大のAIのトレーニングで使用される計算量は、3.5カ月間で30万倍以上に増加し、ムーアの法則のペースをはるかに超えている。

強化されたDeepSpeedは、データ並列訓練、モデル並列訓練、パイプライン並列訓練の3つの技術を活用して、「兆単位」のモデル訓練を可能にするとうたっている。1兆パラメータのモデルをトレーニングするには、少なくとも400個のNvidia A100 GPU（それぞれ40GBのメモリを持つ）のメモリを組み合わせる必要があり、マイクロソフトの試算では、50%の効率で動作する4,000個のA100がトレーニングを完了するのに約100日かかるという。これは、マイクロソフトがOpenAIと共同設計した1万枚以上のグラフィックスカードを搭載したAIスーパーコンピュータには敵わないが、その規模では高い計算効率を得ることは難しい傾向にある。

DeepSpeedは大規模モデルを4つのパイプラインステージで小さなコンポーネント（レイヤー）に分割する。各パイプラインステージ内のレイヤーは、さらに4つの「ワーカー」に分割され、実際の学習を実行する。各パイプラインは2つのデータ並列インスタンスに複製され、ワーカーはマルチGPUシステムにマッピングされる。これらとその他のパフォーマンス向上のおかげで、1兆パラメータのモデルを800台のNvidia V100 GPUでスケーリングすることができるとマイクロソフトは述べている。

DeepSpeedの最新リリースには、GPUとホストCPUの両方の計算リソースとメモリリソースを利用して、1台のV100で最大130億パラメータのモデルを学習できる技術であるZeRO-オフロードも搭載されている。マイクロソフトは、これは最先端技術の10倍の規模であり、少ない計算リソースでもデータ科学者がトレーニングにアクセスできるようになると主張している。

「これらのDeepSpeedの新技術は、計算、メモリ、通信の効率を極限まで高め、数十億から数兆のパラメータのモデルトレーニングを可能にする」とマイクロソフトはブログ記事で書いている。「この技術はまた、非常に長い入力シーケンスを可能にし、単一のGPUを搭載したハードウェアシステム、数千のGPUを搭載したハイエンドクラスタ、または非常に低速なイーサネットネットワークを搭載したローエンドクラスタでもパワーを発揮します...我々は、ディープラーニングトレーニングのスピードとスケールの限界を押し広げ、高速で革新を続けている」。

Photo by Microsoft