## 概要 Mesh-TensorFlow は [[Google Brain]] が開発した、分散テンソル計算を記述するための言語であり、[[@2018__NeurIPS__Mesh-TensorFlow - Deep Learning for Supercomputers]] で提案された。[[TensorFlow]] とほぼ同一の構文(グラフ・テンソル・演算・変数・自動微分)を持ちつつ、テンソルの各次元に名前を付けられる点が特徴で、この名前付き次元を、多次元プロセッサメッシュの任意の次元へ分割する「計算レイアウト」を大域的に一つ指定するだけで、データ並列とモデル並列を任意の粒度で組み合わせた分散実行を記述できる。Mesh-TensorFlow のグラフは、Allreduce のような集団通信プリミティブと結合した並列演算からなる SPMD プログラムへコンパイルされる。 ## アーキテクチャ - **メッシュ**: n 次元のプロセッサ配列。物理ネットワーク位相を意味しない命名上の抽象で、同一の物理プロセッサ集合に対して複数のメッシュ形状(1 次元・2 次元・3 次元)を定義できる。 - **名前付きテンソル次元**: 同一テンソル内で重複しない名前を各次元に持たせることで、「batch」のような論理次元を複数のテンソル・演算間で一貫して分割できる。 - **計算レイアウト**: テンソル次元からメッシュ次元への部分写像。データ並列(`[("batch", "all_processors")]`)からモデル並列、両者の混合、3 次元以上への分割までを同じ枠組みで表現する。 - 初期実装は Python ライブラリで、TPU 向け SPMD TensorFlow コード、または複数 CPU/GPU 向け MIMD コードを生成できる。 ## 実績 論文発表時点で、Mesh-TensorFlow を用いて Transformer のモデル並列実装を構築し、最大 512 コアの TPUv2 メッシュ上で最大 50 億パラメータのモデルを学習した。WMT'14 英仏翻訳(BLEU 43.9)と 10 億語言語モデリングベンチマーク(word-perplexity 24.0)で当時の最高性能を達成した。(Source: [[@2018__NeurIPS__Mesh-TensorFlow - Deep Learning for Supercomputers]]) コードは https://github.com/tensorflow/mesh で公開されている。 ## 関連 - [[@2018__NeurIPS__Mesh-TensorFlow - Deep Learning for Supercomputers]] - [[メッシュ並列]] - [[TensorFlow]] - [[Noam Shazeer]] - [[Google Brain]] ## 出典 - [[@2018__NeurIPS__Mesh-TensorFlow - Deep Learning for Supercomputers]]