> [!abstract] 概要(arXiv abstract の日本語訳) > バッチ分割(データ並列)は、その汎用的な適用可能性と Single-Program-Multiple-Data(SPMD)プログラミングへの適合性のため、分散深層ニューラルネットワーク(DNN)学習における主流の戦略である。しかしバッチ分割は、(メモリ制約による)非常に大規模なモデルを学習できないこと、高レイテンシ、小バッチサイズにおける非効率といった問題に悩まされる。これらすべては、より汎用的な分散戦略(モデル並列)によって解決できる。不運なことに、効率的なモデル並列アルゴリズムは発見・記述・実装が複雑になりがちであり、特に大規模クラスタ上ではそうである。我々は Mesh-TensorFlow を導入する。これは分散テンソル計算の一般的なクラスを記述するための言語である。データ並列は、テンソルと演算を「バッチ」次元に沿って分割するものと見なせるが、Mesh-TensorFlow ではユーザーは任意のテンソル次元を、多次元プロセッサメッシュの任意の次元へ分割するよう指定できる。Mesh-TensorFlow のグラフは、Allreduce のような集団通信プリミティブと結合した並列演算からなる SPMD プログラムへコンパイルされる。我々は Mesh-TensorFlow を用いて、Transformer[21] のシーケンス・ツー・シーケンスモデルの効率的なデータ並列・モデル並列版を実装した。最大 512 コアの TPU メッシュを用いて、最大 50 億パラメータの Transformer モデルを学習し、WMT'14 の英仏翻訳タスクと 10 億語の言語モデリングベンチマークにおいて当時最高性能を上回った。Mesh-Tensorflow は https://github.com/tensorflow/mesh で公開されている。 ## 論文情報 - タイトル: Mesh-TensorFlow: Deep Learning for Supercomputers - 著者: Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, Blake Hechtman(全員 [[Google Brain]] 所属) - 媒体: 32nd Conference on Neural Information Processing Systems(NeurIPS 2018)、Montréal, Canada - arXiv ID: [1811.02084](https://arxiv.org/abs/1811.02084)(2018年11月5日投稿) - コード: https://github.com/tensorflow/mesh ## 概要 Mesh-TensorFlow は、データ並列(バッチ次元の分割)を一般化し、任意のテンソル次元をプロセッサの多次元メッシュの任意の次元へ分割することを記述できる言語である。生成された計算グラフは、Allreduce 等の集団通信プリミティブと組み合わさった SPMD プログラムへコンパイルされる。著者らはこれを用いて Transformer のモデル並列実装を構築し、512 コアの TPU メッシュ上で 50 億パラメータのモデルを学習して、機械翻訳と言語モデリングの両ベンチマークで当時最高性能を達成した。 ## 問題設定 データ並列は各プロセッサが全パラメータの完全なコピーを保持し、バッチをサブバッチに分割して各プロセッサが独立に順伝播・逆伝播を計算し、勾配を Allreduce で集約する(Algorithm 1)。この方式はどのテンソル・演算も「バッチ」次元でのみ分割されるか、さもなければ完全に複製されるかのいずれかであると見なせる。しかし大規模モデルの学習では、パラメータや活性化を保持するためのメモリと、パラメータを同期するための時間が、純粋なデータ並列アルゴリズムを不可能または非効率にする。モデル並列[9]がこれを解決しうるが、既存の MIMD 実装は分散戦略の指定が複雑で、生成されるプログラムも巨大になり最適化が難しい。 ## 提案手法 - **ハードウェア前提**: 各プロセッサがローカルメモリを持つ、同一で信頼できるプロセッサのクラスタを対象とする。**メッシュ**を、そうしたプロセッサの n 次元配列として定義する。メッシュは物理ネットワーク位相を意味しない命名上の抽象であり、同じ物理プロセッサ集合上に異なるメッシュを定義できる(例: 512 コアの TPU クラスタは 16×16×2 のトーラス相互接続を持つが、これを 3 次元メッシュ [16, 16, 2]、2 次元メッシュ [32, 16]、1 次元メッシュ [512] のいずれにも表現できる)。 - **SPMD バッチ分割からの一般化**: データ並列アルゴリズムは、計算中の各テンソルが(バッチ次元を持つ場合)全プロセッサに分割されるか、(持たない場合)全プロセッサに複製されるかのいずれかであり、バッチ次元を縮約する演算だけが追加の Allreduce を要する、という形で記述できる。Mesh-TensorFlow はこのアイデアを、任意の次元での分割へ一般化する。 - **名前付き次元とメッシュ**: テンソルの各次元に名前を持たせることで、異なるテンソル・演算間で同じ分割方針を適用できる論理次元(「batch」等)を表現できる。同一テンソルが同じ名前の次元を 2 つ持つことは許されない。プロセッサ側もメッシュとして n 次元の名前付き次元を持つ。 - **計算レイアウト**: テンソル次元からメッシュ次元への部分写像である「計算レイアウト(computation layout)」を大域的に一つ指定する。例えばデータ並列は、1 次元メッシュ `all_processors` を用い、レイアウト `[("batch", "all_processors")]` として表現される。これは「batch」次元を持つ全テンソルがその次元に沿って全プロセッサへ分割され、他の全テンソルは完全に複製されることを意味する。 - **テンソル表現**: テンソルのレイアウトは、そのテンソルの次元からメッシュ次元への単射的な部分写像であり、大域レイアウトを当該テンソルの次元へ制限したものとして計算される。同一テンソルの 2 つの次元が同じメッシュ次元へ写像されることは許されない。レイアウトが空なら各プロセッサに完全複製される。現行実装では、テンソル次元のサイズがメッシュ次元のサイズで割り切れることが要求される。 - **演算の実装**: 各演算は各プロセッサでの並列計算と、必要に応じた集団通信によって実装される。要素ごとの演算は入出力の形状(レイアウト)が同一であるため各プロセッサの計算だけで済む。reduce 系(`reduce_sum()` 等)は出力次元が入力次元の部分集合になる演算で、各スライスのローカル reduce の後、縮約された次元に対応するメッシュ次元について Allreduce を行う。Einstein 求和(einsum、行列積を含む)は、全入力を和集合の形状へブロードキャストして要素積を取り、出力に無い次元を reduce する演算として定義でき、Mesh-TensorFlow の名前付き次元がこの記法を特に便利にする。 - **reshape の通信パターン**: reshape は非分散の場合は単純だが、Mesh-TensorFlow では出力テンソルのレイアウトが入力と異なる場合に通信を要する。入力側で分割・出力側で非分割ならメッシュ次元をまたぐ MPI-allgather、出力側で分割・入力側で非分割なら通信なしのスライス処理、入力と出力で異なる次元が同じメッシュ次元に分割される場合は MPI-alltoall を要する(データ並列とモデル並列を層ごとに切り替える場合[20]に相当)。 - **構文**: Mesh-TensorFlow 言語は TensorFlow[16] とほぼ同一で、グラフ・テンソル・演算・変数・デバイス(=メッシュ)・自動微分という概念を持つ。主な違いは、テンソル次元が名前とサイズの両方を持つことである。初期実装は Python ライブラリであり、TPU 向け SPMD TensorFlow コード、または複数 CPU/GPU 向け MIMD コード(デバイス配置による)を生成できる。 ### 例: 全結合 2 層のレイアウト比較 $y = \mathrm{Relu}(xw + \mathrm{bias})v$ という単純な 2 層全結合(入出力層 `io` 次元、隠れ層 `hidden` 次元、バッチ `batch` 次元)を用いて、著者らは 4 種類の計算レイアウトを比較する。 - **データ並列レイアウト**: `mesh_shape = [("all", n)]`、`computation_layout = [("batch", "all")]`。パラメータ `w, v, bias` は全プロセッサに複製され、活性化 `x, h, y` はバッチ次元で分割される。順伝播に通信は不要だが、パラメータ勾配の計算がバッチ次元を縮約するため Allreduce を要する。 **Figure 2: データ並列レイアウトの図解(n=2 プロセッサ)** ![[_attachments/arxiv-1811.02084/fig02-data-parallel-layout.png]] (Figure 2. `x, h, y` はバッチ次元(`b/n`)で 2 分割され、行列にはそのスライスを保持するプロセッサの番号が青字で示される。`w` と `v` は両方のランク「0, 1」を持ち、全プロセッサへ完全複製されていることを表す。) - **モデル並列レイアウト**: `mesh_shape = [("all", n)]`、`computation_layout = [("hidden", "all")]`。バッチではなく隠れ層の単位を分割する。入出力層 `x, y` は複製され、隠れ活性化 `h` とパラメータ `w, v, bias` は隠れ次元で分割される。`y` を計算する際に分割された隠れ次元を縮約するため Allreduce が生じる。 **Figure 3: モデル並列レイアウトの図解(n=2 プロセッサ)** ![[_attachments/arxiv-1811.02084/fig03-model-parallel-layout.png]] (Figure 3. 隠れ層の次元(`d_h/n`)がプロセッサ 0 と 1 に分割され、入出力層 `x, y` は両方が「0, 1」のランクを持ち複製されている。) - **データ並列 × モデル並列(2 次元)**: r×c のプロセッサメッシュ上で `mesh_shape = [("rows", r), ("cols", c)]`、`computation_layout = [("batch", "rows"), ("hidden", "cols")]` とすると、行方向にバッチを、列方向に隠れ層を分割できる。プロセッサ数を二乗のオーダーで増やしつつ、バッチサイズと隠れ層サイズを線形にしか増やさなくても効率を保てる。 **Figure 4: データ並列とモデル並列を混合した 2×2 メッシュのレイアウト** ![[_attachments/arxiv-1811.02084/fig04-mixed-2x2-layout.png]] (Figure 4. 2×2 のプロセッサメッシュに 4 プロセッサが直列ランク 0〜3 で配置され、`x` はバッチ次元(`b/r`)で行ごとに、`w, v` は隠れ次元(`d_h/c`)で列ごとに分割される。) - **3 次元メッシュ**: `mesh_shape = [("rows", r), ("cols", c), ("planes", p)]`、`computation_layout = [("batch", "rows"), ("hidden", "cols"), ("io", "planes")]` とすると、`batch`・`hidden`・`io` の 3 次元すべてを分割でき、プロセッサ数を三乗のオーダーで増やしても各次元サイズの増加は線形で済む。 **Figure 5: 2×2×2 の 3 次元メッシュに拡張したレイアウト** ![[_attachments/arxiv-1811.02084/fig05-mixed-2x2x2-layout.png]] (Figure 5. 8 プロセッサが 2 行×2 列×2 面のメッシュに直列ランク 0〜7 で配置され、`io` 次元(`d_io/p`)も追加でプロセッサ面(planes)方向に分割される。) 不十分なレイアウト(空のレイアウト。全演算を全プロセッサで複製するため時間もメモリも節約しない)と不正なレイアウト(同一テンソルの 2 次元が同一メッシュ次元へ写像される)についても論文は具体例で説明している。 ### 性能比較(表) 各レイアウトの計算時間・通信時間・通信/計算比・プロセッサあたりメモリを比較した表(定数因子・低次項は省略)。 | Layout | Comp. Time | Comm. Time | communication/computation | Memory/Processor | |---|---|---|---|---| | `[]`(空) | $bd_{io}d_h$ | 0 | 0 | $bd_{io} + bd_h + d_{io}d_h$ | | `[("batch", "all")]` | $bd_{io}d_h/n$ | $d_{io}d_h$ | $b/n$ | $\frac{b}{n}d_{io} + \frac{b}{n}d_h + d_{io}d_h$ | | `[("hidden", "all")]` | $bd_{io}d_h/n$ | $d_{io}b$ | $n/d_h$ | $bd_{io} + b\frac{d_h}{n} + d_{io}\frac{d_h}{n}$ | | `[("batch","rows"),("hidden","cols")]` | $bd_{io}d_h/rc$ | $d_{io}(\frac{b}{r}+\frac{d_h}{c})$ | $\frac{c}{d_h}+\frac{r}{b}$ | $\frac{b}{r}d_{io}+\frac{b}{r}\frac{d_h}{c}+d_{io}\frac{d_h}{c}$ | | `[("batch","rows"),("hidden","cols"),("io","planes")]` | $bd_{io}d_h/rcp$ | $\frac{b}{r}\frac{d_{io}}{p}+\frac{b}{r}\frac{d_h}{c}+\frac{d_{io}}{p}\frac{d_h}{c}$ | $\frac{c}{d_h}+\frac{p}{d_{io}}+\frac{r}{b}$ | 同上(通信量列と同形) | (Table 1. データ並列レイアウトの通信/計算比は $b/n$、すなわちプロセッサあたりバッチサイズの逆数になり、バッチが小さすぎると性能が劣化する。モデル並列レイアウトでは比が $n/d_h$ で、隠れ層を細かく分割しすぎると劣化する。2 次元・3 次元の混合レイアウトでは、プロセッサ数を二乗・三乗のオーダーで増やしても、バッチサイズと層サイズを線形にしか増やさなくても効率を保てることを示す。) ## 新規性 既存のモデル並列研究([9]の MIMD 実装など)は、分散戦略の指定が複雑で生成プログラムが巨大になりがちだった。Mesh-TensorFlow は、ユーザーが分割したいテンソル次元と、それをどのメッシュ次元へ写すかだけを名指しで指定すれば、SPMD プログラムへの自動コンパイルとレイアウト探索の容易化を両立させる。データ並列(バッチ次元の分割)を「メッシュの 1 次元への分割」という特殊ケースとして包含し、同じ枠組みで任意の次元・任意次元数のメッシュへの分割を表現できることが核心的な貢献である。 ## 実験設定 - **モデル**: Transformer[21] のモデル並列レイアウト。`vocab`・`d_ff`(フィードフォワード隠れ層サイズ)・`heads`(注意ヘッド数)の 3 次元をそれぞれ全プロセッサへ分割するレイアウト `computation_layout = [("vocab", "all"), ("d_ff", "all"), ("heads", "all")]` を採用した。これらの次元をプロセッサ数に比例して増やすことで、プロセッサあたりの通信量・メモリ使用量を一定に保てる(語彙サイズは増やしていない)。さらに大規模化のため、2 次元 TPU メッシュの一方の次元にバッチを、他方に上記のモデル次元を割り当てるデータ並列・モデル並列混合レイアウトも用いた: `computation_layout = [("batch", "rows"), ("vocab", "cols"), ("d_ff", "cols"), ("heads", "cols")]`。 - **ハードウェア**: TPUv2 の 2 次元メッシュ、最大 16×32=512 コア。 - **データセット**: 10 億語言語モデリングベンチマーク(billion-word benchmark)、Tensor2Tensor ライブラリの `languagemodel_wiki_noref_v128k_l1k`(Wikipedia、5 億語超)、WMT'14 英仏(En-Fr)・英独(En-De)翻訳タスク。 - **比較対象**: 10 億語ベンチマークでは先行研究のベスト DNN[20](6.5B パラメータ、word-perplexity 28.0)とベストアンサンブル[17]、WMT'14 では Transformer 原論文[21]の結果。 - **評価指標**: word-perplexity・subword-perplexity(言語モデリング)、BLEU(sacrebleu で評価。機械翻訳)。 ## 実験結果 - **フィードフォワード次元・ヘッド数のスケーリング**: 上記のレイアウトを用いて、フィードフォワード隠れ次元を最大 262144、注意ヘッド数を最大 256 まで増やし、最大 16×32=512 コアの TPUv2 メッシュで学習した。このとき最大のモデルで計算効率 50% 超(最大 11.5 PFLOP/s のうち 6 PFLOP/s)を維持した。 **Table 2: Transformer デコーダ言語モデル($d_{model}=1024$, $d_k=d_v=256$)** | $d_{ff}$ | heads | パラメータ数(億) | Billion-Word ppl | Wikipedia subword-ppl | |---|---|---|---|---| | 4096 | 4 | 1.4 | 35.0 | 8.74 | | 8192 | 8 | 2.2 | 31.7 | 8.03 | | 16384 | 16 | 3.7 | 28.9 | 7.44 | | 32768 | 32 | 6.7 | 26.8 | 6.99 | | 65516 | 64 | 12.8 | 25.1 | 6.55 | | 131072 | 128 | 24.8 | 24.1 | 6.24 | | 262144 | 256 | 49.0 | 24.0(0.9倍のロジットで評価すると23.5) | 6.01 | | Prev Best DNN[20] | | 65 | 28.0 | — | | Best DNN Ensemble[17] | | | 26.1 | — | | Best Ensemble(手法混在)[17] | | > 1000 | 23.7 | — | (Table 2. モデルが大きくなるほど perplexity が単調に改善し、最大の 49 億パラメータモデルは Billion-Word ベンチマークで word-perplexity 24.0(logits を 0.9 倍して評価すると 23.5。過学習が原因と推測される)を達成し、当時の公開結果として最良だった。Wikipedia の `languagemodel_wiki_noref_v128k_l1k` データセットは 5 億語超のテキストからなり、公開されたベースライン結果が存在しないため比較対象がない。) - **billion-word ベンチマークの学習コスト**: 10 エポック学習し、最大モデル(49 億パラメータ)は 512 コア TPUv2 クラスタで 13 時間かかった。バッチサイズは全モデルで 256 シーケンス × 256 トークンで、バッチはメッシュのサイズ 16 の次元に、モデル次元はサイズ 32 の次元に沿って分割された。 **Table 3: Transformer 機械翻訳結果($d_{model}=1024$, $d_k=d_v=128$)** | $d_{ff}$ | heads | $d_k, d_v$ | パラメータ数(億) | WMT14 EN-DE BLEU | WMT14 EN-FR BLEU | |---|---|---|---|---|---| | 2048 | 4 | 128 | 1.5 | 25.5 | 41.8 | | 4096 | 8 | 128 | 2.4 | 26.5 | 42.5 | | 8192 | 16 | 128 | 4.2 | 27.1 | 43.3 | | 16384 | 32 | 128 | 7.7 | 27.5 | 43.5 | | 32768 | 64 | 128 | 14.8 | 27.5 | 43.8 | | 65536 | 128 | 128 | 28.9 | 26.7 | 43.9 | | 4096 | 16 | 64 | 2.1 | 28.4 | 41.8(原論文[21]の設定) | (Table 3. WMT'14 En-Fr では最大モデル(29 億パラメータ)が BLEU 43.9 を達成し、当時の公開結果として最良だった。128-コア TPUv2 クラスタで 22 時間、3 エポック学習した。En-De では学習データが小さいため、モデルサイズによる改善幅がより小さかった。) ## 考察 - モデルサイズを増やすほど言語モデリング・機械翻訳の両タスクで品質が改善する、というスケーリングの傾向がこの時点で明確に示された。ただし En-De では学習データの小ささが改善幅を制約したことから、モデル容量の拡大がデータ量とのバランスを要することが読み取れる。 - Table 1 の性能比較は、レイアウト選択の指針として「通信/計算比」を挙げる。データ並列は per-processor バッチサイズが小さすぎると劣化し、モデル並列は隠れ層を細かく分割しすぎると劣化する。2 次元・3 次元の混合レイアウトはこの制約を緩和し、プロセッサ数を多項式的に増やしても線形のモデル・バッチサイズ増加で効率を保てる。 ## 強み / 弱点・課題 - **強み**: 名前付き次元によって、ユーザーは各テンソルごとに個別のデータレイアウトを指定する必要がなくなり、次元名を一つ指定するだけで関連する全テンソル・演算の分割方針を揃えられる。データ並列をメッシュの特殊ケースとして統一的に扱える点も、既存の owner-compute 系フレームワーク(11章 関連研究で触れる Jia et al. の研究)との差異である。 - **弱点・課題**: 現行実装はテンソル次元のサイズがメッシュ次元のサイズで割り切れることを要求する(§5)。最適な計算レイアウトの自動探索は今後の課題として残されており(§11)、SPMD 実装は TPU 向けが中心で、CPU/GPU クラスタでの SPMD プログラミングの実装も将来課題とされる。空間分割された畳み込みのような「halo」領域通信を要する演算の実装も今後の課題である。