torch-tensor-parallelism

テンソル並列処理の実装ガイド

このスキルでは、PyTorch でテンソル並列処理のパターンを実装する方法を説明します。特に、複数のデバイスに計算を分散する ColumnParallelLinear レイヤーと RowParallelLinear レイヤーを対象とします。

基本概念

テンソル並列処理の概要

テンソル並列処理では、個々のレイヤーを複数のデバイスに分割し、単一の順伝播／逆伝播パス内で計算を並列化します。主なパターンは次の 2 つです。

1. ColumnParallelLinear：重みを出力次元（列）に沿って分割します。各デバイスは出力特徴量の一部を計算し、その後、all-gather によって結果を連結します。

2. RowParallelLinear：重みを入力次元（行）に沿って分割します。各デバイスは入力のシャードを使って部分出力を計算し、その後、all-reduce によって結果を合算します。

重要な実装要件

テンソル並列処理を実装する際（特にシミュレーションやテスト環境では）、順伝播ではローカルのシャードを計算するだけでなく、集合通信も実際に実行する必要があります。

• ColumnParallelLinear：すべてのランクの出力を連結する必要があります（all-gather のセマンティクス）。
• RowParallelLinear：すべてのランクの出力を合算する必要があります（all-reduce のセマンティクス）。

よくある誤りは、ローカルのシャードだけを返し、集合通信は外部フレームワークが処理すると想定することです。明示的に指定されていない限り、実装では最終的な完全な出力を生成する必要があります。

実装方法

ステップ 1：並列処理パターンを把握する

実装する前に、次の点を明確にしてください。

1. どの次元を分割するか（入力特徴量と出力特徴量のどちらか）
2. 結果を結合する集合通信操作は何か（all-gather と all-reduce のどちらか）
3. 実装で分散実行をシミュレートするのか、実際の分散実行に備えるのか
4. 並列処理のコンテキストでバイアスをどのように扱うか

ステップ 2：重みの分割

形状が (out_features, in_features) の重み行列 W の場合：

ColumnParallelLinear：

• W を dim=0（出力特徴量）に沿って分割します。
• 各ランクは、形状が (out_features // world_size, in_features) の W_shard を取得します。
• ランクごとの出力形状：(batch, out_features // world_size)

RowParallelLinear：

• W を dim=1（入力特徴量）に沿って分割します。
• 各ランクは、形状が (out_features, in_features // world_size) の W_shard を取得します。
• 各ランクへの入力には、対応する入力のシャードを渡します。
• ランクごとの出力形状：(batch, out_features) - 部分和

ステップ 3：順伝播の実装

ColumnParallelLinear の順伝播：

1. Compute local output: y_local = x @ W_shard.T + bias_shard (if bias per shard)
2. All-gather to concatenate: y = concat([y_0, y_1, ..., y_n], dim=-1)
3. Return complete output of shape (batch, out_features)

RowParallelLinear の順伝播：

1. Get input shard: x_shard = x[..., start:end] for this rank
2. Compute partial output: y_partial = x_shard @ W_shard.T
3. All-reduce to sum: y = sum([y_0, y_1, ..., y_n])
4. Add bias (only once, not per-rank): y = y + bias
5. Return complete output of shape (batch, out_features)

ステップ 4：バイアスの処理

ColumnParallelLinear：

• バイアスは出力特徴量とともに分割できます。
• 各ランクは、出力のシャードに自身のバイアスのシャードを加算します。
• all-gather の後、バイアス全体が適用された状態になります。

RowParallelLinear：

• バイアスを分割したり、ランクごとに加算したりしてはいけません（バイアスが N 倍になります）。
• all-reduce の後に一度だけバイアスを加算します。
• 通常はランク 0 のみがバイアスを加算するか、合算後にバイアスを加算します。

検証方法

数学的な検証

ローカルでテストできない場合は、数学的な分析を通じて実装の正しさを検証します。

1. 簡単な例：world_size=2 の 2x4 の重み行列を使います。
2. 計算を追跡する：各ランクが生成する値を手作業で計算します。
3. 結合を検証する：all-gather／all-reduce によって正しい最終出力が得られることを確認します。
4. ベースラインと比較する：並列処理の出力が、並列化していない計算と一致することを確認します。

形状の検証チェックリスト

• 入力形状が想定どおりである（batch, in_features）
• 重みのシャード形状が想定される分割方法と一致する
• ローカル出力形状が並列処理の種類に対して正しい
• 最終出力形状が (batch, out_features) と一致する（分割された次元ではありません）

検討すべきテストケース

1. world_size=1：単純なケースです。並列化されていない実装と完全に一致する必要があります
2. world_size=2,4,8：一般的な並列構成です
3. 割り切れない次元：out_features % world_size != 0 の場合、どうなるでしょうか？
4. 異なるバッチサイズ：バッチ次元が正しく処理されることを確認します
5. バイアスあり／なし：両方の構成をテストします

よくある落とし穴

落とし穴 1：ローカルシャードのみを返す

症状：出力テンソルのサイズが (out_features) ではなく (out_features / world_size) になります

原因：実装がローカルシャードを計算しますが、all-gather を実行しません

対処法：すべてのランクの結果を結合する集合通信操作を実装します

落とし穴 2：RowParallelLinear でのバイアス処理の誤り

症状：出力値が予想より N 倍大きくなります（N は world_size）

原因：各ランクがバイアス全体を加算した後、値が合算されます

対処法：ランクごとに加算するのではなく、all-reduce の後にバイアスを一度だけ加算します

落とし穴 3：「シミュレーション」の要件を誤解する

症状：world_size=1 では動作しますが、より大きな world_size では失敗します

原因：外部フレームワークが集合通信操作を処理すると想定しています

対処法：要件を注意深く読みます。「all_gather を使用するかのように」とは、その操作を実装するという意味です

落とし穴 4：ファイル書き込みの切り詰め

症状：実装に構文エラーがある、またはコードが欠けています

原因：ファイル書き込み操作が途中で切り詰められています

対処法：整合性を確認するため、書き込み後に必ずファイル全体を読み戻します

落とし穴 5：シャーディングする次元の誤り

症状：行列乗算中に形状不一致エラーが発生します

原因：誤った次元に沿ってシャーディングしています（行と列を取り違えています）

対処法：ColumnParallel は出力特徴量（重みの dim=0）をシャーディングし、RowParallel は入力特徴量（重みの dim=1）をシャーディングします

実装前チェックリスト

コードを書く前に、以下を理解していることを確認します。

• 必要な集合通信操作（all-gather か all-reduce か）
• 最終的な出力形状
• シミュレーションで集合通信操作を実際に実行するか、実行を委ねるか
• この並列化方式でのバイアスの処理方法
• エッジケース（world_size=1、割り切れない次元）の動作

実装後チェックリスト

コードを書いた後に、以下を確認します。

• 実装ファイル全体を読み戻し、途中で切り詰められていないことを確認します
• すべての world_size で出力形状が想定される次元と一致することを確認します
• 簡単な例を手動でたどり、正しいことを確認します
• 単純なケース（world_size=1）が並列化されていないベースラインと一致することをテストします
• 少なくとも1つの単純でないケース（world_size=2 または 4）をテストします