大規模言語モデル（LLM）は数十億から数千億のパラメータを持ち、単一のGPUメモリには収まりません。このような巨大なモデルを学習・推論するために、モデル並列化（Model Parallelism）という技術が使われます。
モデル並列化には主にTensor並列（Tensor Parallelism）とPipeline並列（Pipeline Parallelism）の2種類があります。本記事では、これらの手法の仕組み、トレードオフ、実装の考え方を解説します。
本記事の内容
- データ並列化の限界
- Tensor並列（モデルの幅方向の分割）
- Pipeline並列（モデルの深さ方向の分割）
- 各手法の比較とトレードオフ
- 実用的な並列化戦略
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
モデルサイズとメモリ
LLMのパラメータ数
代表的なモデルのパラメータ数:
| モデル | パラメータ数 | 必要メモリ（FP16） |
|---|---|---|
| GPT-2 | 1.5B | 3 GB |
| Llama 2 7B | 7B | 14 GB |
| Llama 2 70B | 70B | 140 GB |
| GPT-4 (推定) | 1.7T | 3.4 TB |
float16で1パラメータ = 2バイトなので、70Bモデルでも140GBのメモリが必要です。現在の最大級GPU（H100）でも80GBなので、単一GPUには収まりません。
学習時のメモリ
学習時には、パラメータに加えて以下も必要です:
- 勾配: パラメータと同サイズ
- オプティマイザ状態: Adamの場合、パラメータの2倍（モメンタムと二乗平均）
- アクティベーション: バッチサイズと系列長に依存
Adam + FP16の場合、パラメータサイズの約16倍のメモリが必要になることがあります。
データ並列化とその限界
データ並列化の仕組み
最も基本的な並列化はデータ並列化（Data Parallelism）です。
GPU 0: 全パラメータのコピー, バッチの1/N
GPU 1: 全パラメータのコピー, バッチの1/N
...
GPU N-1: 全パラメータのコピー, バッチの1/N
各GPUで独立に順伝播・逆伝播を計算し、勾配を集約（All-Reduce）してパラメータを更新します。
限界
データ並列化は、モデル全体が各GPUに収まることが前提です。70Bのモデルは単一GPUに収まらないため、データ並列化だけでは対応できません。
Tensor並列（幅方向の分割）
基本アイデア
Tensor並列は、モデルの各層を複数のGPUに分割します。例えば、行列積 $\bm{Y} = \bm{X}\bm{W}$ を複数GPUで分散計算します。
行列積の分割
行列 $\bm{W} \in \mathbb{R}^{m \times n}$ を列方向に分割:
$$ \bm{W} = [\bm{W}_1, \bm{W}_2, \ldots, \bm{W}_N] $$
各GPU $i$ は $\bm{W}_i \in \mathbb{R}^{m \times (n/N)}$ を保持し、部分積を計算:
$$ \bm{Y}_i = \bm{X} \bm{W}_i $$
結果を結合: $$ \bm{Y} = [\bm{Y}_1, \bm{Y}_2, \ldots, \bm{Y}_N] $$
MLPレイヤーの分割
Transformerの Feed-Forward Network を例に考えます。
$$ \bm{Y} = \text{GeLU}(\bm{X}\bm{W}_1)\bm{W}_2 $$
ここで $\bm{W}_1 \in \mathbb{R}^{d \times 4d}$, $\bm{W}_2 \in \mathbb{R}^{4d \times d}$ です。
分割戦略:
- $\bm{W}_1$ を列方向に分割（各GPUが $d \times 2d$ を保持）
- GeLUを各GPUで独立に適用
- $\bm{W}_2$ を行方向に分割（各GPUが $2d \times d$ を保持）
- 結果を All-Reduce で集約
GPU 0: X @ W1_0 -> GeLU -> @ W2_0 -> 部分和
GPU 1: X @ W1_1 -> GeLU -> @ W2_1 -> 部分和
↓
All-Reduce
↓
Y
この戦略では、各GPUが独立に計算でき、最後にAll-Reduceで結果を集約します。
Self-Attentionの分割
Multi-Head Attentionは自然にヘッド方向に分割できます。
$$ \text{MultiHead}(\bm{Q}, \bm{K}, \bm{V}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\bm{W}^O $$
$h$ 個のヘッドを $N$ 台のGPUに分散:
GPU 0: head_1, head_2, ..., head_{h/N}
GPU 1: head_{h/N+1}, ..., head_{2h/N}
...
各GPUで部分的なヘッドを計算し、$\bm{W}^O$ も分割して、最後にAll-Reduceします。
通信パターン
Tensor並列では、以下の通信が発生します:
- Forward時: 各レイヤーの終わりでAll-Reduce（または All-Gather）
- Backward時: 勾配のAll-Reduce
通信は頻繁に発生するため、高速な相互接続（NVLink、NVSwitchなど）が必要です。
Pipeline並列（深さ方向の分割）
基本アイデア
Pipeline並列は、モデルを層（レイヤー）ごとに分割し、各GPUが異なる層を担当します。
GPU 0: Layer 0-7 (Embedding + 最初の8層)
GPU 1: Layer 8-15
GPU 2: Layer 16-23
GPU 3: Layer 24-31 (最後の8層 + Output)
データは「パイプライン」のように各GPUを順番に流れていきます。
ナイーブな実装の問題
素朴な実装では、1つのGPUが計算している間、他のGPUはアイドル状態になります。
時間 →
GPU 0: [計算] [待機] [待機] [待機]
GPU 1: [待機] [計算] [待機] [待機]
GPU 2: [待機] [待機] [計算] [待機]
GPU 3: [待機] [待機] [待機] [計算]
GPU使用率が非常に低く、効率的ではありません。
マイクロバッチ分割
効率を上げるため、バッチを複数のマイクロバッチに分割します。
時間 →
GPU 0: [μ1] [μ2] [μ3] [μ4] [待機] ...
GPU 1: [待機] [μ1] [μ2] [μ3] [μ4] ...
GPU 2: [待機] [待機] [μ1] [μ2] [μ3] ...
GPU 3: [待機] [待機] [待機] [μ1] [μ2] ...
パイプラインが「満たされる」と、すべてのGPUが並列に動作します。
バブル（Pipeline Bubble）
パイプラインの開始時と終了時には、一部のGPUがアイドル状態になります。これをバブルと呼びます。
バブルの割合: $$ \text{Bubble ratio} = \frac{(N_{\text{stages}} – 1)}{N_{\text{microbatches}}} $$
マイクロバッチ数を増やすとバブルは減りますが、メモリ使用量（アクティベーション保存）が増えます。
1F1B スケジュール
GPipe では、すべての Forward を完了してから Backward を開始します。より効率的な 1F1B（One Forward One Backward）スケジュールでは、Forward と Backward を交互に実行し、アクティベーションのメモリ使用量を抑えます。
時間 →
GPU 0: [F1] [F2] [F3] [F4] [B4] [B3] [B2] [B1]
GPU 1: [ ] [F1] [F2] [F3] [B3] [B2] [B1] [B4] ...
通信パターン
Pipeline並列では、隣接するGPU間でのみ通信が発生します:
- Forward時: GPU $i$ から GPU $i+1$ へアクティベーションを送信
- Backward時: GPU $i+1$ から GPU $i$ へ勾配を送信
通信は隣接間のみなので、NVLinkがなくてもある程度動作します。
各手法の比較
| 特性 | データ並列 | Tensor並列 | Pipeline並列 |
|---|---|---|---|
| 分割方向 | データ（バッチ） | モデル（幅） | モデル（深さ） |
| 通信パターン | All-Reduce | All-Reduce | Point-to-Point |
| 通信頻度 | 各ステップ1回 | 各層で発生 | 各マイクロバッチ |
| 必要な相互接続 | 中程度 | 高速（NVLink推奨） | 低速でもOK |
| 効率低下要因 | 勾配通信 | 頻繁なAll-Reduce | バブル |
| スケーラビリティ | バッチサイズに制限 | 通信がボトルネック | バブル増加 |
組み合わせ戦略：3D並列
実際の大規模学習では、3種類の並列化を組み合わせて使用します。
例：8ノード × 8GPU（64GPU）
- Tensor並列: ノード内の8GPUで分割（NVLinkで高速通信）
- Pipeline並列: 2ノードで1パイプライン（4ステージ）
- データ並列: 4パイプラインで同じモデルを複製
ノード0-1: Pipeline Stage 0-3, Data Replica 0
ノード2-3: Pipeline Stage 0-3, Data Replica 1
ノード4-5: Pipeline Stage 0-3, Data Replica 2
ノード6-7: Pipeline Stage 0-3, Data Replica 3
各ステージ内ではTensor並列（8-way）を使用。
メモリ効率化技術
アクティベーションチェックポイント
すべてのアクティベーションを保存する代わりに、一部のみ保存し、必要に応じて再計算します。
$$ \text{Memory} = O(\sqrt{L}) \quad (\text{通常は } O(L)) $$
計算量は増えますが、メモリ使用量が大幅に減少します。
ZeRO（Zero Redundancy Optimizer）
データ並列化において、オプティマイザ状態・勾配・パラメータの冗長コピーを排除します。
- Stage 1: オプティマイザ状態を分割
- Stage 2: 勾配も分割
- Stage 3: パラメータも分割
Stage 3では、各GPUがパラメータの一部のみを保持し、必要に応じて通信で取得します。
実装の考え方（PyTorch）
Tensor並列の概念的実装
import torch
import torch.nn as nn
import torch.distributed as dist
class ColumnParallelLinear(nn.Module):
"""列方向に分割されたLinear層"""
def __init__(self, in_features, out_features, world_size, rank):
super().__init__()
# 各GPUは out_features / world_size の列を担当
self.out_features_per_rank = out_features // world_size
self.weight = nn.Parameter(
torch.randn(in_features, self.out_features_per_rank)
)
self.rank = rank
def forward(self, x):
# 各GPUで部分的な計算
output = x @ self.weight
return output # (batch, out_features_per_rank)
class RowParallelLinear(nn.Module):
"""行方向に分割されたLinear層"""
def __init__(self, in_features, out_features, world_size, rank):
super().__init__()
# 各GPUは in_features / world_size の行を担当
self.in_features_per_rank = in_features // world_size
self.weight = nn.Parameter(
torch.randn(self.in_features_per_rank, out_features)
)
self.rank = rank
def forward(self, x):
# 入力は分割されている想定
output = x @ self.weight
# All-Reduce で結果を集約
dist.all_reduce(output, op=dist.ReduceOp.SUM)
return output
class ParallelMLP(nn.Module):
"""Tensor並列化されたMLP"""
def __init__(self, d_model, d_ff, world_size, rank):
super().__init__()
self.fc1 = ColumnParallelLinear(d_model, d_ff, world_size, rank)
self.fc2 = RowParallelLinear(d_ff, d_model, world_size, rank)
self.gelu = nn.GELU()
def forward(self, x):
# x: (batch, seq, d_model) - 全GPUで同じ
h = self.fc1(x) # (batch, seq, d_ff // world_size)
h = self.gelu(h) # 各GPUで独立に適用
out = self.fc2(h) # All-Reduce で集約
return out # (batch, seq, d_model) - 全GPUで同じ
Pipeline並列の概念的実装
class PipelineStage(nn.Module):
"""パイプラインの1ステージ"""
def __init__(self, layers):
super().__init__()
self.layers = nn.ModuleList(layers)
def forward(self, x):
for layer in self.layers:
x = layer(x)
return x
def pipeline_forward(stages, x, device_ids):
"""
パイプライン実行（概念的な実装）
Args:
stages: 各ステージのモデル
x: 入力テンソル
device_ids: 各ステージのデバイスID
"""
activations = [x]
# Forward pass
for i, (stage, device) in enumerate(zip(stages, device_ids)):
x = x.to(device)
x = stage(x)
activations.append(x)
return x, activations
実践的なツール
Megatron-LM
NVIDIAが開発したLLM学習フレームワーク。Tensor並列とPipeline並列をサポート。
DeepSpeed
Microsoftが開発した分散学習ライブラリ。ZeRO、Pipeline並列、混合精度学習をサポート。
FSDP（Fully Sharded Data Parallelism）
PyTorch標準のZeRO-3相当の実装。パラメータ、勾配、オプティマイザ状態をシャーディング。
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model)
まとめ
本記事では、大規模モデルのための並列化手法について解説しました。
- データ並列: バッチを分割し、同じモデルを複数GPUで複製。モデルがGPUに収まることが前提
- Tensor並列: モデルの各層を複数GPUに分割。頻繁な通信が発生するため、高速な相互接続が必要
- Pipeline並列: モデルを層ごとに分割。通信は隣接GPU間のみだが、バブルが効率を低下させる
- 3D並列: 3種類を組み合わせて大規模モデルを効率的に学習
- ZeRO: オプティマイザ状態の冗長コピーを排除してメモリ効率を向上
大規模LLMの学習・推論には、これらの技術を適切に組み合わせることが不可欠です。ハードウェア構成（GPU数、相互接続速度）に応じて最適な並列化戦略を選択しましょう。
次のステップとして、以下の記事も参考にしてください。