[論文レビュー] Mesh-TensorFlow: Deep Learning for Supercomputers

Noam Shazeer, Youlong Cheng+10 more|arXiv (Cornell University)|Nov 5, 2018

Advanced Neural Network Applications参考文献 21被引用数 52

ひとこと要約

Mesh-TensorFlow は、マルチ次元のプロセッサメッシュ全体で分散テンソル計算を指定する言語を導入し、TPU 上で Transformers のような大規模モデルのスケーラブルなモデル並列およびデータ並列トレーニングを可能にし、最先端の結果を達成します。

ABSTRACT

Batch-splitting (data-parallelism) is the dominant distributed Deep Neural Network (DNN) training strategy, due to its universal applicability and its amenability to Single-Program-Multiple-Data (SPMD) programming. However, batch-splitting suffers from problems including the inability to train very large models (due to memory constraints), high latency, and inefficiency at small batch sizes. All of these can be solved by more general distribution strategies (model-parallelism). Unfortunately, efficient model-parallel algorithms tend to be complicated to discover, describe, and to implement, particularly on large clusters. We introduce Mesh-TensorFlow, a language for specifying a general class of distributed tensor computations. Where data-parallelism can be viewed as splitting tensors and operations along the "batch" dimension, in Mesh-TensorFlow, the user can specify any tensor-dimensions to be split across any dimensions of a multi-dimensional mesh of processors. A Mesh-TensorFlow graph compiles into a SPMD program consisting of parallel operations coupled with collective communication primitives such as Allreduce. We use Mesh-TensorFlow to implement an efficient data-parallel, model-parallel version of the Transformer sequence-to-sequence model. Using TPU meshes of up to 512 cores, we train Transformer models with up to 5 billion parameters, surpassing state of the art results on WMT'14 English-to-French translation task and the one-billion-word language modeling benchmark. Mesh-Tensorflow is available at https://github.com/tensorflow/mesh .

研究の動機と目的

• 大規模DNNのメモリ帯域と待機時間のボトルネックに対処するため、純粋なデータ並列性を超えたスケーラブルなトレーニングを動機づける。
• Mesh-TensorFlow を、マルチディメンションのプロセッサメッシュ全体にわたる分散テンソル計算を指定する言語として導入する。
• Mesh-TensorFlow グラフを、集合通信を伴う SPMD プログラムにコンパイルする方法を示す。
• TPU クラスター上で数十億パラメータの Transformer モデルをトレーニングすることで実用的な利点を示す。

提案手法

• 名前付きテンソル次元と、多次元のプロセッサメッシュを定義する。
• テンソル次元をメッシュ次元へマッピングするグローバルな計算レイアウトを指定する。
• 各テンソルをプロセッサごとのスライスとして表現し、操作をローカル計算として実装し、必要に応じて集約（Allreduce）を行う。
• einsum 風の演算（Einsum）と縮約を用いて、分散シャード間の行列積と縮約を表現する。
• データ並列、モデル並列、および混合のレイアウトを提供し、計算、通信、メモリの観点での性能トレードオフを分析する。

実験結果

リサーチクエスチョン

• RQ1 Mesh-TensorFlow はデータ並列性を超える広いクラスの分散テンソル計算を表現し、効率的に実行できるだろうか？
• RQ2 異なる分配レイアウト（データ並列、モデル並列、ハイブリッド）が、巨大な TPU メッシュでの通信、メモリ、スケーラビリティにどのような影響を与えるか？
• RQ3 大規模クラスター上で Transformer ライクなアーキテクチャに Mesh-TensorFlow を適用することで、どのような性能とモデルサイズの利点が得られるか？

主な発見

• A Mesh-TensorFlow グラフは、並列演算と MPI ライクな集約を備えた SPMD プログラムにコンパイルされる。
• データ並列、モデル並列、ハイブリッドレイアウトは、TPU メッシュ上で数十億パラメータに達する Transformer モデルのトレーニングを可能にする。
• 最大 5 billion parameters の Transformer モデルを最大 512 コアでトレーニングし、WMT’14 En–Fr 翻訳と One Billion Word 言語モデリング ベンチマークで最先端の結果を達成した。
• 多次元メッシュ（例：2D 512-core TPUs）を使用することで、モデルサイズとアテンションヘッドを拡大しても、計算効率は実質的に維持され、ピークの50％超を維持した。
• この手法はデータ並列とモデル並列を組み合わせて、プロセッサ数に比例してバッチサイズとモデル次元の両方をスケールさせることを可能にする。