分散深層学習、特にMegatron-LMのようなフレームワークを使用する際には、各パラメータが独立しているわけではなく、互いに関係し合い、いくつかの制約条件を満たす必要があります。主要なパラメータ間の関係性の制約を以下に列挙します。 **1. GPU数と並列度の関係:** * **`world_size = tensor_parallel_size * pipeline_parallel_size * data_parallel_size`** * `world_size`: 学習に使用するGPUの総数。 * `tensor_parallel_size` (TP): テンソル並列の分割数。 * `pipeline_parallel_size` (PP): パイプライン並列のステージ数。 * `data_parallel_size` (DP): データ並列の分割数。 * **制約**: GPUの総数は、TP、PP、DPの各並列度の積と完全に一致しなければなりません。これは最も基本的な制約です。 * **含意**: TP次数とPP次数を決めると、残りのGPUがデータ並列 (DP) に割り当てられます (`data_parallel_size = world_size / (tensor_parallel_size * pipeline_parallel_size)`)。したがって、`world_size` は `tensor_parallel_size * pipeline_parallel_size` で割り切れなければなりません。 **2. バッチサイズの関係:** * **`global_batch_size = micro_batch_size * data_parallel_size * gradient_accumulation_steps`** * `global_batch_size`: 1回のパラメータ更新（Optimizerステップ）で考慮される実効的な総バッチサイズ。学習のハイパーパラメータとして重要です。 * `micro_batch_size`: 1つのGPU（またはTP/PPで構成される論理的なGPUグループ）が1回のForward/Backwardパスで処理する最小単位のバッチサイズ。GPUメモリ使用量（特にアクティベーション）に直接影響します。 * `gradient_accumulation_steps`: 何回マイクロバッチのForward/Backwardを実行してから勾配を平均化し、1回のOptimizerステップを実行するかを示す数。 * **制約**: これらのパラメータはこの等式を満たすように設定する必要があります。 * **含意**: `global_batch_size` を固定したい場合、`micro_batch_size` や `gradient_accumulation_steps` を変更すると、他のパラメータ（特に `data_parallel_size` がTP/PP/world\_sizeによって決まるため、`micro_batch_size` か `gradient_accumulation_steps`）も調整する必要があります。例えば、メモリ制約から `micro_batch_size` を小さくした場合、同じ `global_batch_size` を維持するには `gradient_accumulation_steps` を増やす必要があります。 **3. テンソル並列 (TP) の制約:** * **`hidden_size` は `tensor_parallel_size` で割り切れなければならない。** * TransformerのLinear層などは、隠れ層次元 (`hidden_size`) で分割されるため、均等に分割できる必要があります。 * **`num_attention_heads` は `tensor_parallel_size` で割り切れなければならない。** * Multi-Head Attentionのヘッドは、TPランク間で分割されるため、均等に分割できる必要があります。 * **(推奨) `tensor_parallel_size` は通常、1ノード内のGPU数以下に設定される。** * TPは高頻度・低レイテンシの通信（All-Reduceなど）を必要とするため、高速なNVLink/NVSwitchで接続されたノード内で完結させるのが一般的です。ノードを跨ぐTPは、InfiniBand/Ethernet経由となりパフォーマンスが大幅に低下します。 **4. パイプライン並列 (PP) の制約:** * **(推奨) `num_layers` (モデルの総レイヤー数) は `pipeline_parallel_size` で割り切れることが望ましい。** * 各パイプラインステージに割り当てられるレイヤー数を均等にするためです。完全に割り切れなくても動作はしますが、ステージ間の計算負荷が不均衡になり、「最も遅いステージ」が全体のボトルネックとなりやすくなります（パイプラインバブル以外の非効率性を生む）。 * **`virtual_pipeline_stage` を使用する場合、その値は `pipeline_parallel_size` の倍数であることが一般的。** * これは物理的なPPステージをさらに細かくスケジューリングするためのもので、負荷分散を改善する目的で使われます。 **5. シーケンス並列 (SP) の制約:** * **SPは通常、TPと組み合わせて使用される。** (`tensor_parallel_size > 1` が必要)。 * Transformerのレイヤー内でシーケンス次元に沿って分割と通信を行うため、TPの通信機構と連携します。 **6. メモリ関連の制約 (間接的な制約):** * **GPUメモリ容量 >= (パラメータ + Optimizer状態 + アクティベーション) のメモリ使用量** * これは直接的なパラメータ間の等式ではありませんが、パラメータ選択における最も重要な**物理的制約**です。 * 各メモリ要素は、TP次数、PP次数、DP次数（ZeRO利用時）、`micro_batch_size`、Activation Checkpointingの有無/方法、混合精度の使用などに依存して変化します。これらのパラメータを調整することで、モデルをGPUメモリに収める必要があります。 これらの制約を理解し、パラメータを設定する際に確認することが、分散学習ジョブを正しく起動し、期待通りのパフォーマンスを得るための第一歩となります。特に、GPU数、並列度、バッチサイズの間の関係は基本であり、間違いやすいポイントです。 > AssertionError: world size is not divisible by tensor parallel size (784) times pipeline parallel size (16) data_parallel_size = 784 / (8 * 4) =