GSPMD：ニューラルネットワークの規模拡大を可能にする汎用的な並列化手法(1/2)

１．GSPMD：ニューラルネットワークの規模拡大を可能にする汎用的な並列化手法(1/2)まとめ

・機械学習アプリケーションではネットワーク規模の拡大が重要であることが判明している
・既存の並列化アルゴリズムは実現が大変だったり実装にも汎用性がなかったりする
・GSPMDはほとんどの深層学習ネットワークを規模拡大可能な汎用性が高い手法

２．GSPMDとは？

以下、ai.googleblog.comより「General and Scalable Parallelization for Neural Networks」の意訳です。元記事は2021年12月8日、Yuanzhong XuさんとYanping Huangさんによる投稿です。

大規模な並列感を出したかったアイキャッチ画像のクレジットはPhoto by Tarang Dave on Unsplash

コンピュータビジョン、言語理解、ニューラル機械翻訳など、多くの実世界の機械学習アプリケーションにおいて、モデルの品質を向上させるためには、使用する学習データの量、モデルサイズ、利用する計算機資源のいずれの観点からであっても、ニューラルネットワークの規模を拡大する事が重要であることが分かっています。

そのため、最近の研究はニューラルモデルを規模拡大するために重要な役割を果たす要因が何かを精査する事に意欲的です。

モデルの容量を増やすことは、モデルの品質を向上させるための健全なアプローチですが、そうすることで克服しなければならないシステムおよびソフトウェア工学上の課題が数多く発生します。

たとえば、アクセラレータのメモリ容量を超えるような大規模なモデルを学習するためには、複数のアクセラレータに渡ってモデルの重みと計算を分割することが必要になります。このような並列化の過程では、ネットワーク通信のオーバーヘッドが増加し、デバイスの利用率が低下する可能性があります。

さらに並列化アルゴリズムを実装するためには、通常、多大なエンジニアリング努力が必要ですが、そのような努力を行っても、モデルアーキテクチャが異なるアクセラレータ上では動作しない場合があります．

これらの規模拡大(scaling)に関する課題に対して、私たちは「GSPMD: General and Scalable Parallelization for ML Computation Graphs」を発表し、XLAコンパイラーに基づくオープンソースの自動並列化システムについて説明しました。

GSPMDは、ほとんどの深層学習ネットワークアーキテクチャを規模拡大することができ、すでにGShard-M4, LaMDA, BigSSL, ViT, MetNet-2などの多くの深層学習モデルに適用され、いくつかの領域で最先端の結果に繋がっています。また、GSPMDは、TensorFlowやJAXなど、XLAを共有コンパイラとして使用する複数のMLフレームワークにも統合されています。

概要

GSPMDは、MLモデルのプログラミングを並列化処理から切り離します。GSPMDを用いると、モデル開発者は、あたかも非常に大きなメモリと計算能力を持つ単一のデバイス上で実行されているかのようにプログラムを書くことができます。

ユーザは、モデルコード中の重要なテンソルの一部に、テンソルの分割方法を示す注釈コードを数行追加するだけで良いのです。

例えば、大規模なモデル並列Transformerを学習する場合、10個以下のテンソル(計算グラフ全体の1％以下)に注釈を付け、テンソルごとに1行のコードを追加するだけで済むかもしれません。

その後、GSPMDはグラフ全体の並列化計画を決定するコンパイラパスを実行し、各デバイスで実行可能な、数学的に等価な並列化計算へと変換します。これにより、ユーザーは並列化の実装ではなく、モデル構築に集中することができ、既存の単一デバイスのプログラムをより大規模に実行できるように簡単に移植することができます。

また、モデルプログラミングと並列処理を分離することで、開発者はコードの重複を最小限に抑えることができます。GSPMDでは、開発者はモデルを再実装することなく、異なる用例に対して異なる並列化アルゴリズムを使用することができます。

例えば、GShard-M4とLaMDAのモデルコードでは、同じモデル実装で異なるモデルやクラスタサイズに適した様々な並列化戦略を適用することが可能です。同様に、GSPMDを適用することで、BigSSLの大規模音声モデルは、以前の小規模モデルと同じ実装を共有することができます。

汎用性と柔軟性

モデル・アーキテクチャによって適した並列化手法が異なるため、GSPMDは様々な用例に適した並列化アルゴリズムをサポートするように設計されています。

例えば、1台のアクセラレータのメモリに収まるような小さなモデルでは、デバイスが異なる入力データを使って同じモデルを学習させるデータ並列が好まれます。一方、単一のアクセラレータのメモリ容量よりも大きなモデルでは、モデルを複数の逐次ステージに分割するパイプラインアルゴリズム(GPipeが採用しているアルゴリズム)や、モデル内の個々の計算演算子をより小さな並列演算子に分割する演算子レベルの並列化(Mesh-TensorFlowなど)がより適しています。

GSPMDは、上記のすべての並列化アルゴリズムを、統一された抽象化と実装でサポートしています。さらに、GSPMDは入れ子構造にした並列化パターンをサポートしています。例えば、モデルを個々のパイプライン・ステージに分割し、各ステージは演算子レベルの並列性を用いてさらに分割することなどができます。

また、GSPMDは、パフォーマンスチューニングの専門家が、多くのクロスデバイス通信を伴う実装ではなく、ハードウェアを最もよく活用するアルゴリズムに集中できるようにすることで、並列化アルゴリズムに関する技術革新を促進します。

例えば、大規模なTransformerモデルに対して、デバイスの2Dメッシュ上で多次元のテンソルを分割する新しい演算子レベルの並列化アルゴリズムを発見しました。これは、多次元にバランスよくデータを分配することにより、アクセラレータの計算効率を高く保ちつつ、アクセラレータのピークメモリ使用量を学習デバイスの数に対して線形に削減するものです。