# シーケンス並列化 ## 定義 シーケンス並列化(sequence parallelism)は、トランスフォーマーの訓練においてテンソル並列化が適用されない領域(LayerNorm・Dropout)の処理をシーケンス次元に沿って分割し、冗長複製を排除するメモリ最適化手法である。 ### 背景: テンソル並列化の「穴」 [[Megatron-LM]] 式のテンソル並列化(Shoeybi et al. 2019)は、アテンション/MLP ブロック内部のアクティベーションを t(テンソル並列度)台のデバイス間で分割する。しかし: - 各ブロックへの入力テンソル(Q/K/V 計算への入力、h→4h 線形層への入力)は分割されず、全デバイスで複製される - ブロック後段の LayerNorm と Dropout もテンソル並列の対象外で複製される これが 1 層あたりのアクティベーション式 Eq.2 で `10sbh` の項(テンソル並列度 t で除算されない部分)として現れる。 ### 解法: g と ḡ 演算の挿入 テンソル並列化の非適用領域ではシーケンス次元に沿った独立性が保たれているため、この次元での分割が可能である。テンソル並列化の境界に新しい通信演算 g と ḡ を挿入する: - **g**: 順伝播 = All-Gather(シーケンス軸を結合し全データを揃える)、逆伝播 = Reduce-Scatter - **ḡ**: 順伝播 = Reduce-Scatter(加算して再分割)、逆伝播 = All-Gather **通信帯域はゼロ追加**: テンソル並列化の既存 all-reduce は内部的に reduce-scatter + all-gather の 2 ステップであるため、g と ḡ を使う総通信量はテンソル並列化のみの場合と等しい。追加のネットワーク帯域は発生しない。 MLP ブロックの定式化(Eq.3): $[Y_1^s, Y_2^s] = \text{LayerNorm}([X_1^s, X_2^s])$ $Y = g(Y_1^s, Y_2^s)$ $[Z_1^h, Z_2^h] = [\text{GeLU}(Y A_1^c), \text{GeLU}(Y A_2^c)]$ $[W_1^s, W_2^s] = \bar{g}(W_1, W_2)$ ここで上付き添字はテンソルの分割次元(s: シーケンス、h: 隠れ次元)を示す。 ### 効果 テンソル+シーケンス並列化後のアクティベーション式(Eq.4): $\frac{sbh}{t}\left(34 + 5\frac{as}{h}\right)$ これは並列化なし式(Eq.1)をテンソル並列度 t で割ったものと一致し、アクティベーションがテンソル並列グループ全体で均等分散される。 逆伝播時に追加の All-Gather が 1 回必要になるが、後続の勾配計算と通信をオーバーラップさせることで、実測での順伝播速度向上(22B モデルで 6%: 7.7ms → 7.2ms)を実現する。 ## 横断的知見 - **先行シーケンス並列化(Li+ 2021a)との根本的差異**: Li+ 2021a の sequence parallelism は全デバイスにパラメータとオプティマイザ状態を複製するため、大規模モデル訓練には不適であった。Korthikanti+ 2023 の手法は [[Megatron-LM]] のテンソル並列化と相補的に設計され、パラメータ/オプティマイザ状態の複製を必要としない。 - **通信量等価性の保証**: `all-reduce = reduce-scatter + all-gather` という恒等式により、追加通信ゼロが保証される。これは実測での速度向上(シーケンス並列化単独で −3% オーバーヘッド)の根拠である。 - **[[並列化戦略]]の「通信局所性原則」との整合**: サーベイ/MegaScale が「TP をノード内 NVLink に閉じる」理由と同様、シーケンス並列化もテンソル並列グループ内(ノード内 NVLink)で reduce-scatter/all-gather を実行するため、高帯域ノード内通信の恩恵を受ける。 - **[[選択的活性化再計算]]との相乗**: シーケンス並列化が定数項(10sbh → 10sbh/t)を削減し、選択的再計算がシーケンス長二乗項(5as²b → 排除)を削減する。2 手法はアクティベーションメモリの異なる項を担当するため、組み合わせで 5 倍削減という超加法的効果が得られる。 ## 未解決の問い - シーケンス並列化は現行実装では「テンソル並列度と同じグループ」に閉じる。テンソル並列と独立したシーケンス並列グループを設定したとき、通信トポロジと性能はどう変わるか(DeepSpeed-Ulysses、Ring Self-Attention との比較)。 - 非常に長いシーケンス(1M トークン以上)では、アテンション自体を ring ベースで分散する Ring Self-Attention 系(DistFlashAttn・Striped Attention)との組み合わせが必要になる。Korthikanti+ 2023 の手法はアテンション外部(LayerNorm・Dropout)のシーケンス分散であり、アテンション内部のシーケンス分散(ring attention)と相補的か競合するかを明確にする必要がある。 - `reduce-scatter` と `all-gather` の分離実行が、NVLink/NVSwitch と InfiniBand の混在ネットワークでどのようなボトルネックを生むか。 ## 関連 - 親概念: [[並列化戦略]] / [[LLM分散学習]] - 姉妹概念: [[選択的活性化再計算]] / [[チェックポイント]] - ソース: [[@2023__MLSys__Reducing Activation Recomputation in Large Transformer Models]] - エンティティ: [[Megatron-LM]] / [[NVIDIA]] / [[Vijay Korthikanti]] / [[Mohammad Shoeybi]] / [[Bryan Catanzaro]] - 関連 MOC: [[分散深層学習 - MOC]] ## 出典 - [[@2023__MLSys__Reducing Activation Recomputation in Large Transformer Models]](Korthikanti+, MLSys 2023: g/ḡ 演算の導出・Eq.4 のアクティベーション式・通信帯域等価性の証明)