Skip to content
Published on

マルチGPU学習の四つの並列化 — 何を分割し、何を通信するか

シェア
Authors

はじめに — 80GBカード8枚でも7Bモデルが載らない理由

「7Bモデルならbf16で14GBだから、80GBカード1枚に十分収まる」という計算が成り立つのは推論だけです。同じモデルを学習しようとすると、カード1枚ではそもそも始まりもせず、8枚つないでも設定次第で依然として落ちます。

理由は、学習がパラメータのほかにさらに三つの塊を抱えていなければならないからです。勾配、オプティマイザ状態、そして活性化です。この四つの塊のうち何をどう分割するかが並列化戦略のすべてであり、分割した代償としてGPU同士が何をどれだけやり取りしなければならないかが性能のすべてです。

この記事はその二つの軸、何を分割するかどれだけ通信するかで四つの並列化を整理します。数字はすべて自分で検算できる形で示します。推論側のメモリ計算はLLM推論のVRAM計算ですでに扱ったので、ここでは学習だけを見ます。

メモリの内訳 — パラメータあたり16バイト、そしてそれより大きい活性化

状態メモリはパラメータあたり16バイト

混合精度学習でAdam系オプティマイザを使うとき、パラメータ一つが占めるバイト数は次のように内訳されます。

項目精度パラメータあたりバイト数
パラメータのコピーbf162
勾配bf162
マスター重みfp324
Adam 1次モーメントfp324
Adam 2次モーメントfp324
合計16

この16バイトはZeRO論文が示した内訳と同じです。7Bモデルなら112GB、70Bなら1,120GBです。80GBカード1枚では7Bさえ学習できないという結論はここから出てきます。

オプティマイザを変えるとこの内訳も変わります。SGD with momentumならfp32状態がマスター重みとモメンタムの二つだけなのでパラメータあたり12バイトになり、8ビットAdamやAdafactorを使うとさらに減ります。ただし事前学習の規模でオプティマイザを変えるのはメモリの問題ではなく収束の問題なので、低精度オプティマイザは主にファインチューニングで使われます。

活性化こそが実際にカードを溢れさせる

上の内訳には活性化が含まれていません。ところが事前学習でOOMを起こす主犯は大抵活性化です。KorthikantiらのReducing Activation Recomputation in Large Transformer Modelsは、並列化がない場合のトランスフォーマー層一つ分の活性化メモリを、シーケンス長s、バッチb、隠れ次元h、アテンションヘッド数aについてs·b·h·(34 + 5as/h)バイトとまとめています。

def act_gb(s: int, b: int, h: int, a: int, layers: int, mode: str) -> float:
    sbh = s * b * h
    per_layer = {
        "none": sbh * (34 + 5 * a * s / h),   # すべて保存
        "selective": sbh * 34,                # アテンション行列のみ再計算
        "full": sbh * 2,                      # 層の入力のみ保存し、すべて再計算
    }[mode]
    return per_layer * layers / 1e9


s, b, h, a, L = 4096, 1, 4096, 32, 32   # 7B級構成、サンプル1個
for mode in ("none", "selective", "full"):
    print(f"{mode:10} {act_gb(s, b, h, a, L, mode):7.2f} GB / サンプル")
# none        104.15 GB / サンプル
# selective    18.25 GB / サンプル
# full          1.07 GB / サンプル

サンプル一つで104GBです。マイクロバッチをわずか4に設定するだけで416GBになります。すべて再計算すれば1.07GBまで落ちますが、代わりに順伝播をもう一度回す分だけ計算量がおよそ3分の1増えます。

ここで34と一緒に付いている5as/h項に注目してください。この項はシーケンス長の二乗に比例します。sを4Kから32Kに増やすと34項は8倍になりますが、5as/h項は64倍になります。長いコンテキストの学習が急に難しくなるのはまさにこの地点であり、コンテキスト並列が必要になる理由もここにあります。

データ並列とZeRO — 内訳のどの項を分割するか

最も単純な並列化であるDDPは、すべてのGPUが上の16バイトをそのまま複製して持ちます。バッチだけを分割し、勾配をall-reduceしたあとそれぞれが同じ更新を行います。実装は最も単純ですが、メモリはまったく節約されません。

ZeROはこの複製を段階的に取り除いていきます。

def per_gpu_state_gb(params_b: float, world: int, stage: int) -> float:
    """Adam混合精度基準、GPU1枚が抱える状態メモリ(GB、10^9基準)。
    stage 0 = DDP、1 = オプティマイザ分割、2 = +勾配、3 = +パラメータ
    活性化は含みません。"""
    p = params_b * 1e9
    param, grad, opt = 2 * p, 2 * p, 12 * p
    if stage >= 1:
        opt /= world
    if stage >= 2:
        grad /= world
    if stage >= 3:
        param /= world
    return (param + grad + opt) / 1e9


for stage in (0, 1, 2, 3):
    row = [round(per_gpu_state_gb(n, 8, stage), 1) for n in (7, 13, 70)]
    print(f"stage {stage}  world=8  7B/13B/70B -> {row} GB")
# stage 0  world=8  7B/13B/70B -> [112.0, 208.0, 1120.0] GB
# stage 1  world=8  7B/13B/70B -> [38.5, 71.5, 385.0] GB
# stage 2  world=8  7B/13B/70B -> [26.2, 48.8, 262.5] GB
# stage 3  world=8  7B/13B/70B -> [14.0, 26.0, 140.0] GB

読み方はこうです。ステージ3は16バイト全体をGPU数で割るので、GPUを増やせば際限なく小さくなります。反対にステージ1はパラメータと勾配の4バイトが分割されずに残るため、GPUをいくら増やしても7Bモデルは28GBを下回りません。GPUを増やしたのにメモリが減らないなら、大抵ZeROのステージが低いのです

PyTorchではFSDPが事実上ZeROステージ3に相当します。2026年8月2日時点でPyTorch公式チュートリアルは"FSDP1 is deprecated"と明記し、torch.distributed.fully_shardベースのFSDP2を推奨しています。チュートリアルにはどのバージョンから廃止なのかは書かれていないので、この点は使用しているPyTorchバージョンのリリースノートで改めて確認する方が安全です。確認したPyTorch安定版バージョンは2.13.0で、リリース日は2026年7月8日です。

# FSDP2の最小形態 — モジュールをラップする代わりにその場でシャーディングする
import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

mesh = init_device_mesh("cuda", (torch.distributed.get_world_size(),), mesh_dim_names=("dp",))
policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)

# ルートだけでなくサブモジュールにも適用してこそ通信単位が細かく分割される
for block in model.layers:
    fully_shard(block, mesh=mesh, mp_policy=policy)
fully_shard(model, mesh=mesh, mp_policy=policy)

最後の行だけを書いてループを省略すると、モデル全体が一つの通信単位になり、順伝播の開始時点でパラメータ全体を一度に集めてしまいます。メモリ削減効果が消え、通信と計算の重なりもなくなります。「FSDPを有効にしたのになぜ減らないのか」という質問の多くは、まさにこのミスによるものです。

テンソル並列 — 行列を分割し、レイヤーごとに通信する

テンソル並列は一つの行列積を複数のGPUで分担して計算します。Megatron-LM論文の古典的な構成は、MLPの最初の行列を列方向に、二番目を行方向に分割して途中の通信なしにつなぎ合わせ、ブロックの末尾でだけall-reduceを一度行います。アテンションブロックも同じ方式でヘッド単位に分割します。

結果として、トランスフォーマー層一つあたり順伝播2回、逆伝播2回のall-reduceが発生します。通信の対象はパラメータではなく活性化テンソルであり、サイズはs·b·hです。

def tp_bytes_per_step(s, b, h, layers, tp, dtype_bytes=2):
    """テンソル並列でGPU1枚が1ステップに動かすバイト数(リングall-reduce近似)。"""
    tensor = s * b * h * dtype_bytes
    ring = 2 * (tp - 1) / tp          # リングall-reduceのGPUあたり移動量係数
    per_layer = 4 * tensor * ring     # 順伝播2 + 逆伝播2
    return per_layer * layers


print(tp_bytes_per_step(4096, 1, 4096, 32, tp=8) / 1e9, "GB/step")
# 7.516192768 GB/step

ステップあたり7.5GBを動かさなければなりません。NVLink 4を使うH100 SXMはGPUあたり双方向900GB/sを公称しているので、この通信は10ミリ秒前後で終わります。同じ通信を400Gb/sのInfiniBandポート一つで行うと50GB/sになるので、150ミリ秒になります。テンソル並列グループがノードの境界を越えた瞬間に学習速度が一桁落ちる理由が、この算数です

リンク公称帯域幅位置
NVLink 4 (H100 SXM)GPUあたり900 GB/s双方向ノード内
NVLink 5 (B200)GPUあたり1.8 TB/s双方向ノード内
PCIe Gen5 x16約128 GB/s双方向NVLinkのないノード内
InfiniBand NDRポート1つ400 Gb/s、約50 GB/sノード間

上の数値はメーカー公称値であり、実測の有効帯域幅はこれより低くなります。正確な値はクラスタでnccl-testsall_reduce_perfを直接実行して得る方がよいでしょう。実務のルールは単純です。テンソル並列の次数はノード内のGPU数を超えません。8-GPUノードならTPは8以下です。

パイプライン並列 — 通信は最も安く、バブルは最も高くつく

パイプライン並列はレイヤーをステージに切り分けて別々のGPUに載せます。ステージの境界でだけ活性化をやり取りするので、通信量は圧倒的に少なくなります。32レイヤーを4ステージに切ると境界は3つだけになり、各境界でs·b·hサイズのテンソル一つが点対点でやり取りされます。そのためパイプラインはノードの境界を越えてもかまいません。

代償はバブルです。最初のステージが最初のマイクロバッチを処理している間、残りのステージは遊んでおり、最後のマイクロバッチの逆伝播が終わるまで前のステージがまた遊びます。

def bubble_ratio(stages: int, micro_batches: int) -> float:
    return (stages - 1) / micro_batches

for m in (4, 8, 16, 32, 64):
    print(f"micro_batches={m:3}  8ステージバブル {bubble_ratio(8, m):.1%}")
# micro_batches=  4  8ステージバブル 175.0%
# micro_batches=  8  8ステージバブル 87.5%
# micro_batches= 16  8ステージバブル 43.8%
# micro_batches= 32  8ステージバブル 21.9%
# micro_batches= 64  8ステージバブル 10.9%

バブルを10パーセント台に抑えるには、マイクロバッチがステージ数の8倍以上必要です。ところがマイクロバッチを増やすと活性化メモリが増え、グローバルバッチサイズも大きくなって学習率スケジュールに影響します。この三つの値が互いに結びついているという点が、パイプラインのチューニングを難しくしています。

インターリーブドスケジュールや、DeepSeek-V3のDualPipeのようにバブルを減らすスケジュールがいくつも登場しています。DualPipeは順伝播と逆伝播の計算-通信区間を互いに重ねてパイプラインバブルを減らす方式で、DeepSeek-V3技術報告書に説明があります。PyTorch側ではtorch.distributed.pipeliningが標準APIとして提供されています。

シーケンス並列とコンテキスト並列 — 四つ目の軸

前の三つの軸はすべてパラメータとバッチを基準に分割します。シーケンス長を基準に分割する軸がもう一つあります。用語が混同されがちなので、区別して使うことにします。

  • シーケンス並列: Megatron系でテンソル並列と対になる技法です。テンソル並列が分割できずに残したLayerNormとdropout区間の活性化をシーケンス軸で分割し、重複保存をなくします。テンソル並列と同じグループの中で動作します。
  • コンテキスト並列: アテンション計算そのものをシーケンス軸で分割します。各GPUがシーケンスの一部だけを持ち、キーとバリューをリング状に回しながら部分アテンションを積算します。リングアテンション系がここに属します。DeepSpeedのUlyssesは似た目標をall-to-all通信で達成する別の実装です。

コンテキスト並列を使う基準は明確です。シーケンス長を伸ばし、マイクロバッチを1まで下げ、活性化をすべて再計算してもなおOOMになるなら、それがコンテキスト並列を有効にすべき時点です。Llama 3 405Bの学習が4D並列化(TP、CP、PP、DP)を使った理由も、128Kコンテキスト拡張段階のためです。

MoEモデルを学習するなら、ここにエキスパート並列がもう一つ加わります。エキスパートをGPUに分けて載せ、トークンをall-to-allでルーティングする方式です。通信パターンがall-reduceではなくall-to-allであるためネットワークトポロジーにはるかに敏感で、ロードバランシングが崩れると特定のGPUだけが働き、残りは待たされます。

組み合わせを選ぶ — サイズとGPU数で読む表

実務で決めるべきことは「どの並列化が優れているか」ではなく「この規模でどんな順序に積み上げるか」です。一般的に通用している順序は、通信コストの高い軸を内側に配置することです。TPをノード内に、CPをその次に、PPをノード間に、DPを一番外側に置きます。

モデルサイズGPU数初期組み合わせ根拠
1B〜8B1〜8ZeRO-2またはFSDP2パラメータがカードに収まる。TPは通信を増やすだけ
8B〜30B8〜64FSDP2 + アクティベーションチェックポインティング単一の軸なのでデバッグが容易
30B〜100B64〜512TP 2〜8(ノード内)× FSDP2TPをNVLink内に閉じ込める
100B以上512〜TP 8 × PP 2〜16 × DPPPでノード間通信量を下げる
コンテキスト32K以上任意上の組み合わせ + CP活性化のシーケンス二乗項のため
MoE任意上の組み合わせ + EPエキスパートを分散、all-to-allを許容

この表は出発点であって正解ではありません。実際には目標とするグローバルバッチサイズが組み合わせを強く制約します。グローバルバッチはマイクロバッチ×勾配累積×データ並列次数であり、DP次数は全体のGPU数をTP、PP、CPの次数で割った値に固定されます。TPを8から4に変えるとDPが二倍になり、グローバルバッチも二倍になって学習率を再調整しなければなりません。並列化設定の変更は純粋なインフラ作業ではなく、ハイパーパラメータの変更です

よくある失敗パターン

現場で繰り返し目にするものだけを挙げます。

  1. TPグループがノードをまたぐ。16-GPUのTPを8-GPUノード2台にまたがせると速度が崩壊します。nvidia-smi topo -mで実際のリンクを確認し、ランチャーがランクをどの順序でノードに配置するかを確認してください。
  2. FSDPをルートにだけ適用する。前で見たミスです。サブモジュール単位でラップしてこそ通信が計算と重なります。
  3. アクティベーションチェックポインティングを有効にしてもOOM。再計算の対象にならない埋め込みと出力層、そしてロジットテンソルが残ります。語彙サイズが大きいモデルではロジットがs·b·Vのサイズになるため、単独で数十GBになります。チャンク単位の損失計算に分けて処理する必要があります。
  4. マイクロバッチを減らしてOOMを避け、バブルを放置する。PP環境でマイクロバッチを減らすとメモリは解決しますが、バブルが大きくなってスループットが半分になります。二つは同じレバーの両端です。
  5. 勾配累積と正規化の不一致。累積ステップ数で損失を割る位置を間違えると、実効学習率が変わってしまいます。並列構成が変わるたびに、短いランで損失曲線が以前の構成と重なるかを確認してください。
  6. NCCLタイムアウトを延ばして問題を覆い隠す。タイムアウトが出るときは大抵、特定のランクが死んでいるか、遅いノードが一つあるかのどちらかです。NCCL_DEBUG=INFOでどのランクが止まったのかをまず見る必要があります。

おわりに — 並列化の選択は帯域幅予算を使うことです

四つの並列化は互いの代替物ではなく、それぞれ別の資源を消費する道具です。ZeROとFSDPはメモリを帯域幅に変え、テンソル並列はメモリをノード内帯域幅に変え、パイプラインはメモリをアイドル時間に変え、コンテキスト並列はシーケンス軸のメモリをリング通信に変えます。

ですから設定を選ぶ前にやるべきことは、フレームワークのドキュメントを読むことではなく、クラスタの帯域幅地図を描くことです。ノード内GPU間の実測帯域幅、ノード間の実測帯域幅、そしてその二つの比率。この三つの数字さえわかれば、上の表はほぼ自動的に決まります。次の記事では、この決定を実際にコードへ落とし込んでくれる学習フレームワークスタックを系統別に整理します。