テンソル並列（Tensor Parallelism）とは、ニューラルネットワークの個々のレイヤー内のテンソル（行列）演算を複数のGPUに分割して並列に実行する手法です。モデル並列の一種で、モデルを横方向（幅方向）に分割します。NVIDIAのMegatron-LMで提案された手法として知られています。
テンソル並列の仕組み
Transformerの全結合層を例にすると、重み行列を列方向または行方向に分割し、各GPUが分割された部分行列の演算を担当します。各GPUの計算結果をAllReduceまたはAllGatherで集約することで、分割しない場合と同じ結果を得られます。1つのレイヤーの計算を複数GPUで分担するため、レイヤー内の並列化と言えます。
| ひとことで言うと | 1つの層の行列計算そのものを分割して、複数GPUで同時に計算する方式。 |
|---|---|
| 何がうれしいか | 層単位の分割より待ち時間が少なく、GPUを有効に使える。 |
| 注意点 | 計算のたびに通信が発生する。NVLinkなど超高速接続が必須。 |
通信の特性
テンソル並列は各レイヤーの順伝播・逆伝播のたびに通信が発生するため、通信頻度が高い手法です。そのため、NVLinkのような高帯域・低レイテンシのインターコネクトが不可欠で、通常は同一ノード内のGPU間でのみ適用されます。ノード間ではパイプライン並列やデータ並列が使われます。
1巨大な行列計算1層の中の処理
▶
2行列を縦に分割各GPUが一部を担当
▶
3同時に計算並行して処理
▶
4結果を結合通信して1つにまとめる
シーケンス並列との組み合わせ
テンソル並列と組み合わせてシーケンス並列（Sequence Parallelism）も使用されます。LayerNormやDropoutなどテンソル並列で分割されない演算について、シーケンス（トークン列）方向に分割することで、メモリ使用量をさらに削減します。Megatron-LM v3でこの組み合わせが提案されています。