Distributed Training(分散学習)

article technology high #distributed-training#data-parallel#tensor-parallel#pipeline-parallel#sequence-parallel#expert-parallel#zero#fsdp#deepspeed#megatron-lm#ring-attention#nccl#allreduce#llm-training#gpu-cluster
Created: 2026-05-26 Updated:

LLM 訓練の 5 並列軸 (DP / TP / PP / SP-CP / EP) を ZeRO・FSDP・Megatron-LM・Ring Attention・DeepEP で整理、3D 合成・メモリ計算 (70B=1.12 TB)・落とし穴・2026 景観を集約。

Distributed Training(分散学習)

70B fp16 + Adam のモデル状態は静的だけで約 1.12 TB、H100 SXM (80 GB) や B200 (192 GB) 単一枚では収まらない。Chinchilla (Hoffmann 2022, arXiv:2203.15556) の最適計算量 6ND\approx 6ND から 70B × 1.4T tokens で 5.88 × 10²³ FLOP、H100 単独 1,300 万時間超。メモリ・計算ウォール同時突破のため分散学習は不可避、データ並列 (DP) / テンソル並列 (TP) / パイプライン並列 (PP) / シーケンス・コンテキスト並列 (SP・CP) / エキスパート並列 (EP) の 5 軸を 3D〜5D で合成する設計問題に帰着する。Transformer は tech-35、MoE は tech-36、Attention は tech-38、シャッフル/grad accumulation は tech-39 を参照。

モデル状態のメモリ収支

fp16 混合精度 + Adam で 16 B/param (Rajbhandari 2020, arXiv:1910.02054 §3): Params (2) + Grads (2) + Adam master/m/v fp32 (4+4+4)。70B → 1.12 TB、13B でも 208 GB。アクティベーション・通信バッファは別途上積み。イントラノード帯域 (NVLink 4 = 900 GB/s, 5 = 1.8 TB/s) とノード間 (IB NDR 400 Gbps ≈ 50 GB/s 片方向) は常に 約 1 桁差 — これが並列軸の物理配置を決める最大の制約。

5 つの並列化軸

データ並列 (DP / DDP)

各ワーカーがモデル完全コピー、ミニバッチ分割、backward で AllReduce 集約。PyTorch DDP はリング AllReduce + 勾配バケッティング (25 MB) で計算と overlap。長所: 実装単純・1 phase 同期。短所: モデルが 1 GPU 収容前提、通信 2M(P1)/P\approx 2M(P-1)/PPP 増大に脆弱。適用: 〜7B または fp8 量化済み中規模。

ZeRO と FSDP

Rajbhandari 2020 (arXiv:1910.02054) は DP のメモリ冗長性を 3 段で削減: ZeRO-1 (Optimizer state, 最大 8×) → ZeRO-2 (+Gradients) → ZeRO-3 (+Parameters, 最大 64×)。ZeRO-3 は forward/backward で AllGather → 使用後即破棄、通信量は ZeRO-1/2 で DP 同等・ZeRO-3 で微増。ZeRO-Infinity (arXiv:2104.07857) は NVMe/CPU offload で兆パラ級だが帯域律速。

PyTorch FSDP (Zhao 2023, arXiv:2304.11277) は ZeRO-3 相当のネイティブ実装。FSDP2 は per-parameter DTensor シャーディングに刷新、torch.compile + activation checkpointing と完全 composable (fully_shard() API)。2025〜2026 PyTorch 標準、13B〜100B のゴールデンゾーン。

テンソル並列 (TP)

Megatron-LM (Shoeybi 2019, arXiv:1909.08053)。MLP / Attention の GEMM を列・行分割: Y=GeLU(XA)BY = \text{GeLU}(XA)B を A 列分割 → B 行分割 → AllReduce、Attention は head 単位で同様。各 layer につき forward/backward 各 2 回 AllReduce、メッセージは batch×seq×hidden\text{batch} \times \text{seq} \times \text{hidden} 規模。

TP は NVLink 前提。IB NDR を跨ぐと AllReduce レイテンシが 10〜20× 跳ね上がりスループットが壊滅。TP は必ずイントラノードに閉じる — 通常 TP=8 (DGX 1 ノード)、GB200 NVL72 のみ 16〜72 に拡張可能。

パイプライン並列 (PP)

モデルを depth 方向に pp ステージ分割、バブル率 =(p1)/(p1+m)= (p-1)/(p-1+m) (mm=マイクロバッチ数)。GPipe (Huang 2018, arXiv:1811.06965) は F-then-B でアクティベーション O(pm)O(pm)1F1B / PipeDream (Narayanan 2019, arXiv:1806.03377) は 1 forward 1 backward で O(p)O(p)Interleaved 1F1B (Narayanan 2021, arXiv:2104.04473) は仮想パイプライン vv チャンクでバブル率 ×1/v\times 1/vZero Bubble Pipeline (Qi 2024, arXiv:2401.10241) は backward を WW (weight grad) と BB (activation grad) に分離して理論バブル 0 (Megatron-Core 実装)。実践則: m4pm \ge 4p (例: p=8,m=4p=8, m=4 でバブル 46%)。

シーケンス並列 / コンテキスト並列 (SP / CP)

長文脈訓練を支える。Megatron SP (Korthikanti 2022, arXiv:2205.05198) は LayerNorm/Dropout アクティベーションをシーケンス方向に分散し TP と組合せて 1/TP1/\text{TP} に削減、TP の拡張として使う。Ring Attention (Liu 2023, arXiv:2310.01889) はリング状 GPU で QQ スライス保持しつつ K/VK/V をリング伝搬、1M+ トークン訓練の主軸。DeepSpeed Ulysses (Jacobs 2023, arXiv:2309.14509) は head 分割 + All-to-All で QKV 転置、Ring より通信回数少ないが head 数 worker 数で割り切れる制約あり。

エキスパート並列 (EP) — MoE 専用

tech-36 の Expert を GPU 分散、トークン dispatch を All-to-All で実現。GShard (Lepikhin 2020, arXiv:2006.16668) が capacity factor を定式化、DeepEP (DeepSeek 2025) が IB 上の高効率 All-to-All で DeepSeek-V3 (671B MoE) 訓練に投入。超大規模は各 Expert 内 TP + Expert 間 EP の 2 段 (EP × TP) が標準。

3D / 4D 並列の合成

Megatron-Turing NLG 530B (Smith 2022, arXiv:2201.11990): TP=8 × PP=35 × DP≈10、計 2,240 A100。2024〜2026 標準レシピは TP=8 をイントラノード (NVLink)、PP をラック内ノード間 (IB NDR)、DP/ZeRO をラック間に配置、長文脈は + SP/CP、MoE は + EP。制約: TP×PP\text{TP} \times \text{PP} が 1 ノード超えない構成が最効率、SP は TP ペア固定、EP は他軸と直交だが All-to-All 帯域律速。統合実装は DeepSpeed Megatron-DeepSpeedNVIDIA NeMo、Meta torchtitan

通信プリミティブ

プリミティブ用途通信量
AllReduceDP 勾配集約2M(P1)/P2M(P-1)/P
AllGatherZeRO-3 param, TPM(P1)/PM(P-1)/P
ReduceScatterZeRO-2/3 gradM(P1)/PM(P-1)/P
All-to-AllTP 転置, EP dispatchM(P1)/PM(P-1)/P
Send/Recv (P2P)PP アクティベーションMM

AllReduce は ReduceScatter+AllGather に分解 (ring-allreduce)、NCCL は実測ピーク帯域の 70〜85%。NCCL 2.18+ SHARP で Quantum-2 IB スイッチのインネットワーク演算が可能。GB200 NVL72 は 72 枚を単一 NVLink ドメインに収容し TP=72 をラック内完結化 — 並列設計を 2026 年に書き換えうる転換点。

メモリ計算 — Llama-3 70B

静的状態: Params 140 + Grads 140 + Adam (master+m+v) 840 = 1,120 GB。ZeRO-3 で DP=32 に分散すれば 35 GB/GPU で H100 80GB 収容可。

アクティベーション (Korthikanti 2022 式): per-layer ≈ seq × batch × hidden × (34 + 5 × heads × seq / hidden)。Llama-3 70B (hidden=8192, heads=64, layers=80, seq=4096, batch=1) で per-layer ≈ 6.51 GB、80 層で 520 GB。Selective Recomputation (Korthikanti 2022) で Attention softmax/dropout のみ再計算 → 30〜40% 保持・再計算 +5〜10% FLOP、標準採用。Full Activation Checkpointing (Chen 2016) は全破棄+再計算で O(L)O(\sqrt{L}) メモリ・計算 +33%。実践例: ZeRO-3 + TP=8 + Selective Recomputation で 32 枚 H100 (4 ノード) に seq=4096・batch=32。global batch 設計は tech-39 参照。Llama 3.1 405B は 16,384 H100 で 4D 並列訓練。

落とし穴 — 頻発する 7 つの失敗モード

  1. TP のノード間延伸 — AllReduce が IB 律速、スループット 1/10〜1/20。対策: TP はイントラノード厳守、例外は NVL72 のみ。
  2. PP マイクロバッチ不足p=8,m=4p=8, m=4 でバブル 46%。対策: m4pm \ge 4p または Interleaved / Zero Bubble PP。
  3. DistributedSampler.set_epoch(epoch) 忘れ — 全エポック同シャッフル、データ多様性ゼロ。詳細 tech-39
  4. ZeRO-3 + no_sync() + AC 組合せバグ — grad-accum バウンダリで backward が no_sync() 外に漏れ、初期マイクロバッチ勾配 sync が非同期化。対策: framework 組み込み grad-accum API を使う。
  5. 混合精度 LayerNorm + SP の grad clip 漏れ — LayerNorm fp32 下で SP 分散すると clip_grad_norm_ がローカルノルム打切で過少 clip。対策: グローバル AllReduce 後にノルム計算 (FSDP2/DeepSpeed 自動)。
  6. 異種 GPU 混在のストラグラー — TP/PP リング内の遅い GPU が律速 (H100/A100 差 3×)。対策: 同種を同ノード、異種は DP のみ混在。
  7. fp8 訓練の loss spike — scaling factor 不適切、特に勾配 E5M2 の動的レンジ狭。対策: DelayedScaling で scaling history window 広く、fp8↔fp16 cast を Attention/MLP の外側へ。

フレームワーク景観 (2025〜2026)

フレームワーク主機能
PyTorch FSDP2ZeRO-3 + per-param DTensor + torch.compile、PyTorch 標準
DeepSpeedZeRO-1/2/3 + ZeRO-Infinity (NVMe/CPU offload)
Megatron-LM / Megatron-CoreTP+PP+SP+EP リファレンス (NVIDIA)
NVIDIA NeMoMegatron-Core + FSDP + PEFT/SFT/RLHF、エンタープライズ標準
torchtitanTP+PP+FSDP2+CP のクリーン PyTorch 実装 (Meta)
JAX/Pax (GSPMD), MaxTextXLA sharding、TPU ネイティブ (Google)
HAI-LLM (DeepSeek)EP 特化 (DeepEP) + FP8、DeepSeek-V3 671B を 2.788M H800-hours

注目事例: DeepSeek-V3 (DeepSeek-AI 2024, arXiv:2412.19437) は HAI-LLM + DeepEP + FP8 で 671B MoE を同規模他社の 1/3〜1/5 コストで訓練。Llama 3.1 405B (Meta 2024) は 16,384 H100 で TP × PP × DP の 4D 並列。GPT-NeoX / Pythia (Biderman 2023, arXiv:2304.01373) は 1F1B PP + ZeRO-1 + TP のオープン教科書。

バッチサイズと LR スケーリング

Global Batch = micro_batch × grad_accum × DPCritical Batch Size (McCandlish 2018, arXiv:1812.06162) 超で統計効率低下 (GPT-3 で約 3.2M tokens/step)。Linear LR Scaling Rule (Goyal 2017, arXiv:1706.02677): batch k× → LR k×、Linear Warmup 併用。2024〜2025 は WSD (Warmup-Stable-Decay) が主流。Grad Clipping はグローバルノルム必須 — ZeRO-3/FSDP のローカルノルム打切は過少 clip、DeepSpeed と FSDP2 は自動グローバル化。

2026 フロンティア

  • 4D/5D 自動並列: Alpa (Zheng 2022, arXiv:2201.12023)・GSPMD を経て Auto-parallel (FSDP2 + torch.compile) の実用化。
  • Ring Attention 1M+ トークン訓練: 長文脈 SFT/CPT 拡大、通信 vs 計算比の最適化が課題。
  • DeepEP / MoE 通信特化: Expert load imbalance を通信スケジュールで吸収、EP × TP 合成の標準化。
  • NVLink Fusion ラック (GB200 NVL72): TP=72 ラック内完結、IB は DP/PP 専用に — 並列軸の物理配置が一変。
  • DiLoCo (Douillard 2023, arXiv:2311.08105): 局所 HH ステップ → outer optimizer で疑似勾配集約、通信 1/H1/H で地理分散も可能化。
  • FP8 / FP4 訓練: Transformer Engine + DeepSeek 実績を経て 2026 標準化、Blackwell NV_FP4 で fp4 も実験段階。

See also

Local graph