[論文レビュー] GSPMD: General and Scalable Parallelization for ML Computation Graphs

Yuanzhong Xu, HyoukJoong Lee+14 more|arXiv (Cornell University)|May 10, 2021

Parallel Computing and Optimization Techniques参考文献 24被引用数 37

ひとこと要約

GSPMDは、単純なテンソル分割注釈を用いてML計算グラフを自動的に分割する自動化されたコンパイラベースのシステムで、デバイス間でのデータ/モデル/パイプライン並列性を拡張可能にします。

ABSTRACT

We present GSPMD, an automatic, compiler-based parallelization system for common machine learning computations. It allows users to write programs in the same way as for a single device, then give hints through a few annotations on how to distribute tensors, based on which GSPMD will parallelize the computation. Its representation of partitioning is simple yet general, allowing it to express different or mixed paradigms of parallelism on a wide variety of models. GSPMD infers the partitioning for every operator based on limited user annotations, making it convenient to scale existing single-device programs. It solves several technical challenges for production usage, allowing GSPMD to achieve 50% to 62% compute utilization on up to 2048 Cloud TPUv3 cores for models with up to one trillion parameters.

研究の動機と目的

• 単一デバイスの実行を超える大規模MLモデルのスケーラブルな並列化を動機づける。
• 複数の並列性パラダイムをサポートする、一般的でコンパイラ駆動の分割機構を提供する。
• ユーザーが単一デバイス上のようにモデルを記述し、分割注釈を付けて分散実行を駆動できるようにする。
• 統一されたフレームワーク内で、ネストされた混合並列パターン（データ、モデル、空間、オプティマイザ状態）を可能にする。

提案手法

• XLAベースのコンパイラバックエンドを拡張して、一般的なテンソル分割表現（replicated、tiled、partially tiled）を実装する。
• テンソル次元をデバイスメッシュにマッピングして分割注釈を生成する mesh_split API を導入する。
• シャーディングの完了と各演算子の分割パスを開発し、演算子間で shardings を伝搬・統合する。
• Single Program Multiple Data (SPMD) partitioning をサポートして、数千のパーティションに拡張し、パディング/マスキングを用いて静的形状を管理する。
• マイクロバッチ処理パイプラインをステージ分割テンソルへ翻訳するラッパーを介して、パイプライン並列をテンソル分割へ還元する。
• 複雑な演算子（例：Einsum、Convolution）やランクポリモーフィック演算を扱うために、ネスト型・再帰的分割を取り入れる。

実験結果

リサーチクエスチョン

• RQ1 GSPMDは、1つのフレームワークで、多様な並列パターン（データ、モデル、空間、オプティマイザ状態、パイプライン）を表現し、組み合わせることができるのか。
• RQ2 限られたユーザー注釈が、ML計算グラフ全体の完全なシャーディングをどれだけ効果的に推進できるか。
• RQ3 SPMDパーティショニング（静的形状、ハロー交換、通信）の実務上の課題は何か、それらにどう対処できるか。
• RQ4 大規模な TPU展開で、計算利用率とメモリスケーリングの観点でGSPMDはどのように性能を発揮するか。
• RQ5 GSPMDはマルチモーダルまたは大規模モデルに対して、ネストされた不均一な並列パターンを実現できるか。

主な発見

• モデルが1兆パラメータまでで、最大2048 Cloud TPUv3コアで50%〜62%の計算利用率を達成。
• 不均一な/静的形状をサポートし、演算子間のシャーディングを統合して、広範な手書きルールなしでネスト並列を可能にする。
• データ、モデル、空間、オプティマイザ状態のシャーディングを、単純なデバイスメッシュAPI（mesh_split）を持つ統一的なテンソル分割モデルで表現する。
• パイプライン並列を層ごとのシャーディングへ還元し、軽量ラッパーでパイプライニングを有効にして、別個のパイプライン基盤を回避する。
• すべてのパーティションの単一プログラム（SPMD）を維持してコンパイルのスケーラビリティを保ちつつ、複雑な演算子の意味論（例：Convolution、Einsum）を扱う。
• XLA/TensorFlow/JAXバックエンドとの統合を示し、プロダクション風のコンパイルと通信プリミティブ（AllReduce, AllGather, ReduceScatter, CollectivePermute）をサポートする。