%%{init: {'theme': 'base', 'themeVariables': {'fontSize': '14px'}}}%%
graph TB
A[分布式训练范式] --> B[数据并行<br/>Data Parallelism]
A --> C[模型并行<br/>Model Parallelism]
A --> D[流水线并行<br/>Pipeline Parallelism]
A --> E[张量并行<br/>Tensor Parallelism]
A --> F[专家并行<br/>Expert Parallelism]
A --> G[混合并行<br/>Hybrid Strategy]
B --> B1["每张卡持有完整模型<br/>处理不同数据批次"]
C --> C1["模型按层切分<br/>各卡负责不同层"]
D --> D1["模型按阶段切分<br/>微批次流水执行"]
E --> E1["层内权重切分<br/>多卡协同计算"]
F --> F1["MoE 路由<br/>不同专家分布在不同卡"]
分布式训练范式
分布式训练范式
单卡训练的时代已经结束了。当你读完这一章，你会理解为什么。
为什么我们需要分布式训练
2020 年，GPT-3 用了 1750 亿参数，训练成本约 460 万美元。2024 年，前沿模型的参数量已经突破万亿，训练需要数千张 GPU 协同工作。一张 H100 的显存是 80GB——连一个万亿参数模型的权重都放不下。
分布式训练不是”锦上添花”，而是”唯一可行路径”。
但分布式训练的复杂性在于：你可以选择很多种方式把工作拆开。每种方式都有自己的代价。理解这些权衡，是设计 AI 基础设施的基本功。
4.1 数据并行（Data Parallelism）
数据并行是最直观的分布式策略：每张 GPU 持有完整的模型副本，各自处理不同的数据批次，然后通过 AllReduce 同步梯度。
工作原理
%%{init: {'theme': 'base'}}%%
sequenceDiagram
participant W0 as GPU 0
participant W1 as GPU 1
participant W2 as GPU 2
participant W3 as GPU 3
Note over W0,W3: 前向传播（各自独立）
W0->>W0: Forward(batch_0)
W1->>W1: Forward(batch_1)
W2->>W2: Forward(batch_2)
W3->>W3: Forward(batch_3)
Note over W0,W3: 反向传播（各自独立）
W0->>W0: Backward(batch_0)
W1->>W1: Backward(batch_1)
W2->>W2: Backward(batch_2)
W3->>W3: Backward(batch_3)
Note over W0,W3: AllReduce 梯度同步
W0->>W1: Gradient Sync
W1->>W0: Gradient Sync
W1->>W2: Gradient Sync
W2->>W1: Gradient Sync
W2->>W3: Gradient Sync
W3->>W2: Gradient Sync
Note over W0,W3: 优化器更新（各自独立）
W0->>W0: Optimizer Step
W1->>W1: Optimizer Step
W2->>W2: Optimizer Step
W3->>W3: Optimizer Step
PyTorch DDP 最小示例
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import os
def setup():
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
def cleanup():
dist.destroy_process_group()
def train():
setup()
model = MyModel().cuda()
# DDP 包装：自动处理梯度同步
model = DDP(model, device_ids=[local_rank])
# DistributedSampler 确保每个进程看到不同的数据
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for epoch in range(num_epochs):
# 关键：每个 epoch 必须设置 sampler
sampler.set_epoch(epoch)
for batch in loader:
loss = model(batch)
loss.backward() # DDP 自动在 backward 时同步梯度
optimizer.step()
optimizer.zero_grad()
cleanup()
if __name__ == "__main__":
train()启动 8 卡训练：
torchrun --nproc_per_node=8 train.py数据并行的代价
数据并行的核心瓶颈是 通信开销。每一步训练都需要一次 AllReduce 操作来同步梯度。当 GPU 数量增加到数百张时，通信可能占据 30%-50% 的训练时间。
Tip实践建议
何时使用数据并行：模型能放进单卡显存时（如 7B 模型在 80GB H100 上），数据并行是首选方案——简单、可靠、调试友好。
Warning注意
数据并行有一个硬限制：模型的参数量 + 优化器状态 + 梯度 + 激活值必须放进单卡显存。一个 70B 模型用 AdamW 优化器需要约 1.1TB 显存——远远超过单卡容量。这时候你需要模型并行。
4.2 模型并行（Model Parallelism）
模型并行将模型的不同部分放到不同的 GPU 上。最简单的形式是按层切分：把一个 100 层的 Transformer 切成 4 段，每段 25 层放在一张卡上。
简单模型并行
class ModelParallelTransformer(nn.Module):
def __init__(self):
super().__init__()
# 第一段层放在 GPU 0
self.layers_part1 = nn.ModuleList([
TransformerLayer() for _ in range(25)
]).to(f"cuda:0")
# 第二段层放在 GPU 1
self.layers_part2 = nn.ModuleList([
TransformerLayer() for _ in range(25)
]).to(f"cuda:1")
def forward(self, x):
x = x.to("cuda:0")
for layer in self.layers_part1:
x = layer(x)
x = x.to("cuda:1") # 跨卡传输
for layer in self.layers_part2:
x = layer(x)
return x致命缺陷：串行执行
简单模型并行的问题是同一时刻只有一张 GPU 在工作。当 GPU 0 在计算前 25 层时，GPU 1 完全空闲。这导致 GPU 利用率极低。
GPU 0: ████████░░░░░░░░░░░░ (计算) (等待...)
GPU 1: ........████████░░░░ (等待...) (计算)
Warning为什么不直接用简单模型并行
GPU 利用率 = 1/N（N 是 GPU 数量）。4 张卡只有 25% 的利用率，这在价值数千万的集群上是不可接受的浪费。解决方案是下一节要讲的流水线并行。
4.3 流水线并行（Pipeline Parallelism）
流水线并行是模型并行的进化版：它将输入数据切成微批次（micro-batch），让多张 GPU 像工厂流水线一样同时工作。
GPipe 方式：同步流水线
%%{init: {'theme': 'base', 'themeVariables': {'fontSize': '12px'}}}%%
gantt
title GPipe 流水线（4 个微批次，2 个阶段）
dateFormat X
axisFormat %s
section GPU 0 (Stage 0)
MB0 前向 :a1, 0, 1
MB1 前向 :a2, 1, 2
MB2 前向 :a3, 2, 3
MB3 前向 :a4, 3, 4
气泡 :b1, 4, 8
MB3 反向 :b2, 8, 9
MB2 反向 :b3, 9, 10
MB1 反向 :b4, 10, 11
MB0 反向 :b5, 11, 12
section GPU 1 (Stage 1)
气泡 :c1, 0, 1
MB0 前向 :c2, 1, 2
MB1 前向 :c3, 2, 3
MB2 前向 :c4, 3, 4
MB3 前向 :c5, 4, 5
MB3 反向 :c6, 5, 6
MB2 反向 :c7, 6, 7
MB1 反向 :c8, 7, 8
MB0 反向 :c9, 8, 9
上图中灰色的”气泡”就是流水线空闲时间。气泡比例的计算公式：
\[ \text{Bubble fraction} = \frac{p - 1}{m + p - 1} \]
其中 \(p\) 是流水线阶段数，\(m\) 是微批次数。当 \(m \gg p\) 时，气泡比例趋近于 0。
1F1B 方式：交错流水线
1F1B（One Forward, One Backward）策略通过更激进的调度来减少气泡：一旦第一个微批次的前向传播到达最后一个阶段，立即开始反向传播，同时后续微批次继续前向。
# 1F1B 调度伪代码（简化版）
def schedule_1f1b(num_microbatches, num_stages, stage_id):
warmup = num_stages - stage_id - 1
warmup = min(warmup, num_microbatches)
# Warmup 阶段：执行若干前向
for i in range(warmup):
forward_step(microbatch=i)
# Steady state：一前一后交替
remaining = num_microbatches - warmup
for i in range(remaining):
forward_step(microbatch=warmup + i)
backward_step(microbatch=i)
# Cooldown 阶段：完成剩余反向
for i in range(warmup):
backward_step(microbatch=num_microbatches - warmup + i)
Tip实践建议
微批次数量的经验法则：\(m \geq 4 \times p\)（微批次数至少是流水线阶段数的 4 倍）。这能将气泡比例控制在 25% 以下。但微批次越多，每个微批次越小，通信开销也会增加。
使用 Megatron-LM 的流水线并行
# 8 卡，2 路张量并行，4 路流水线并行
python pretrain_gpt.py \
--tensor-model-parallel-size 2 \
--pipeline-model-parallel-size 4 \
--num-layers 32 \
--hidden-size 4096 \
--num-attention-heads 32 \
--micro-batch-size 2 \
--global-batch-size 64 \
--bf164.4 张量并行与序列并行（Tensor/Sequence Parallelism）
张量并行是分布式训练中最”精细”的切分方式：它不是按层切，而是在每一层的内部把权重矩阵切开。
张量并行的工作原理
对于一个线性层 \(Y = XW\)，张量并行将权重矩阵 \(W\) 沉着列或行切分：
%%{init: {'theme': 'base'}}%%
graph LR
subgraph "列并行（Column Parallel）"
X1[X] --> W1["W[:, 0:n/2]"] --> Y1["Y[:, 0:n/2]"]
X1 --> W2["W[:, n/2:n]"] --> Y2["Y[:, n/2:n]"]
end
subgraph "行并行（Row Parallel）"
XA["X[:, 0:n/2]"] --> WA["W[0:n/2, :]"] --> YA["Partial Y"]
XB["X[:, n/2:n]"] --> WB["W[n/2:n, :]"] --> YB["Partial Y"]
YA --> SUM(("+"))
YB --> SUM
SUM --> YOUT[Y]
end
关键洞察：Megatron-LM 发现 Transformer 的 MLP 和 Attention 层天然适合张量并行，只需要 2 次通信（一次 All-Reduce）就能完成一层的前向和反向传播。
# Megatron 风格的列并行线性层（简化版）
class ColumnParallelLinear(nn.Module):
def __init__(self, in_features, out_features, tp_size):
super().__init__()
assert out_features % tp_size == 0
self.out_features_per_partition = out_features // tp_size
self.weight = nn.Parameter(torch.empty(
self.out_features_per_partition, in_features
))
def forward(self, x):
# 每张卡只计算输出的一个分片
return torch.nn.functional.linear(x, self.weight)
class RowParallelLinear(nn.Module):
def __init__(self, in_features, out_features, tp_size):
super().__init__()
assert in_features % tp_size == 0
self.in_features_per_partition = in_features // tp_size
self.weight = nn.Parameter(torch.empty(
out_features, self.in_features_per_partition
))
def forward(self, x):
# 每张卡计算部分结果，然后 AllReduce 求和
partial = torch.nn.functional.linear(x, self.weight)
return dist.all_reduce(partial, op=dist.ReduceOp.SUM)序列并行（Sequence Parallelism）
张量并行有一个隐藏的开销：在 LayerNorm 和 Dropout 等操作中，激活值在所有 TP 组成员之间是完整复制的，这浪费了大量显存。
序列并行（Korthikanti et al., 2022）将序列维度也切分开来。在非 Attention/MLP 的区域（如 LayerNorm），每个 TP rank 只处理序列的一段，从而减少激活值内存。
标准张量并行： [full_seq] → LayerNorm → [full_seq] → Attention(TP) → ...
序列并行： [seq/TP] → LayerNorm → All-Gather → Attention(TP) → Reduce-Scatter → ...
Tip何时使用张量并行
张量并行最适合节点内（intra-node）使用，因为它需要频繁的 AllReduce 通信。NVLink/NVSwitch 互联的节点内通信带宽可达 900GB/s，而跨节点 IB 通常只有 50-100GB/s。
4.5 专家并行与 MoE 训练（Expert Parallelism）
混合专家模型（Mixture of Experts, MoE）是一种”稀疏激活”架构：每个 token 只激活模型中的一小部分参数（专家），从而在不增加计算量的前提下大幅扩展模型容量。
MoE 的基本结构
%%{init: {'theme': 'base'}}%%
graph LR
T["输入 Token"] --> Router["路由器 / Gate Network"]
Router -->|Top-K 选择| E1["专家 1<br/>GPU 0"]
Router -->|Top-K 选择| E2["专家 2<br/>GPU 1"]
Router -->|Top-K 选择| E3["专家 3<br/>GPU 2"]
Router -->|Top-K 选择| E4["专家 4<br/>GPU 3"]
E1 --> W1["加权输出"]
E2 --> W1
E3 --> W1
E4 --> W1
W1 --> OUT["最终输出"]
专家并行
在专家并行中，不同的专家被放到不同的 GPU 上。当输入 token 到来时，路由器决定将 token 发送到哪些 GPU。
# 简化的 MoE 前向传播
class MoELayer(nn.Module):
def __init__(self, d_model, d_ff, num_experts, top_k=2):
super().__init__()
self.router = nn.Linear(d_model, num_experts)
self.experts = nn.ModuleList([
FeedForward(d_model, d_ff)
for _ in range(num_experts)
])
self.top_k = top_k
def forward(self, x):
# x: [batch, seq_len, d_model]
router_logits = self.router(x) # [batch, seq_len, num_experts]
# Top-K 选择
weights, indices = torch.topk(
router_logits.softmax(dim=-1), self.top_k, dim=-1
)
weights = weights / weights.sum(dim=-1, keepdim=True)
# 分发 token 到对应专家（实际使用 all-to-all 通信）
output = torch.zeros_like(x)
for k in range(self.top_k):
for expert_idx in range(len(self.experts)):
mask = (indices[..., k] == expert_idx)
if mask.any():
expert_input = x[mask]
expert_output = self.experts[expert_idx](expert_input)
output[mask] += weights[..., k:k+1][mask] * expert_output
return outputMoE 训练的三个核心挑战
- 负载均衡（Load Balancing）：如果所有 token 都涌向同几个专家，其他 GPU 就会空闲。解决方案是加入辅助损失（auxiliary loss），惩罚专家利用率的不均衡。
\[ \mathcal{L}_{aux} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i \]
其中 \(f_i\) 是分配给专家 \(i\) 的 token 比例，\(P_i\) 是路由器对专家 \(i\) 的平均概率。
Token 丢弃（Token Dropping）：当某个专家的容量溢出时，超出部分被丢弃。这会损失信息但保证训练不卡住。
All-to-All 通信：专家并行需要两次 all-to-all 通信（分发 token → 收集结果），这对网络互联是很大的压力。
WarningMoE 不一定更快
MoE 的吞吐量取决于稀疏度。如果 top_k=2 且有 8 个专家，每个 token 只使用 25% 的参数，但路由和 all-to-all 通信的额外开销可能吃掉节省的计算时间。在集群互联不够强时，稠密模型反而可能更快。
4.6 混合并行策略的选择与组合
实际生产中，没有任何单一并行策略能覆盖所有场景。前沿大模型训练通常同时使用 3D 甚至 4D 并行。
3D 并行：数据 × 张量 × 流水线
%%{init: {'theme': 'base'}}%%
graph TB
subgraph "集群：64 张 GPU"
subgraph "DP Group 0 (16 卡)"
subgraph "PP Stage 0 (4 卡 × TP=4)"
TP00["GPU 0<br/>TP Rank 0"]
TP01["GPU 1<br/>TP Rank 1"]
TP02["GPU 2<br/>TP Rank 2"]
TP03["GPU 3<br/>TP Rank 3"]
end
subgraph "PP Stage 1"
TP10["GPU 4-7<br/>TP=4"]
end
subgraph "PP Stage 2"
TP20["GPU 8-11<br/>TP=4"]
end
subgraph "PP Stage 3"
TP30["GPU 12-15<br/>TP=4"]
end
end
subgraph "DP Group 1 (16 卡)"
DP1["同样的结构<br/>GPU 16-31"]
end
subgraph "DP Group 2 (16 卡)"
DP2["同样的结构<br/>GPU 32-47"]
end
subgraph "DP Group 3 (16 卡)"
DP3["同样的结构<br/>GPU 48-63"]
end
end
配置选择决策树
| 模型规模 | GPU 数量 | 推荐 3D 配置 | 典型场景 |
|---|---|---|---|
| 7B | 8 | DP=8, TP=1, PP=1 | 单节点 fine-tune |
| 7B | 64 | DP=16, TP=4, PP=1 | 多节点训练 |
| 70B | 64 | DP=4, TP=8, PP=2 | 中等规模预训练 |
| 175B | 512 | DP=8, TP=8, PP=8 | GPT-3 规模 |
| 1T+ | 2048+ | DP=16, TP=8, PP=16+ | 万亿参数训练 |
经验法则
Tip3D 并行配置原则
- 张量并行（TP）：限制在节点内（通常 ≤ 8），匹配 NVLink 拓扑
- 流水线并行（PP）：跨节点使用，微批次数 ≥ 4 × PP stages
- 数据并行（DP）：用剩余的 GPU，尽量最大化以提升吞吐量
- TP × PP 应该刚好让单个模型副本的显存需求适配可用 GPU 组
使用 DeepSpeed 的混合并行
# deepspeed_config.json
{
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 8,
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "none"
},
"offload_param": {
"device": "none"
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": 5e8,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e5
},
"bf16": {
"enabled": true
},
"pipeline": {
"stages": 4,
"activation_checkpoint_interval": 1
},
"tensor_parallel": {
"tp_size": 4,
"tp_group_size": 4
}
}# 启动脚本
# 32 卡：TP=4, PP=4, DP=2
deepspeed --num_gpus 8 --num_nodes 4 \
train.py \
--deepspeed_config deepspeed_config.json \
--model_name "my-175b-model"小结
| 并行策略 | 切分维度 | 通信模式 | 内存节省 | 适用场景 |
|---|---|---|---|---|
| 数据并行 | 数据 | AllReduce | 无（需 ZeRO 配合） | 模型能放进单卡 |
| 模型并行 | 层 | 点对点 | 高（利用率低） | 已被流水线并行取代 |
| 流水线并行 | 层（微批次） | 点对点 + 同步 | 高 | 超大模型跨节点 |
| 张量并行 | 层内权重 | AllReduce | 中高 | 节点内、大层 |
| 序列并行 | 序列维度 | AllGather/RS | 中 | 配合 TP 减少激活值 |
| 专家并行 | 专家（路由） | All-to-All | 视稀疏度而定 | MoE 模型 |
选择并行策略的本质是在计算、通信、内存三角中找到平衡点。没有”最优”策略——只有”最适合你的模型、硬件和预算”的策略。
延伸阅读
- Megatron-LM 论文系列：Shoeybi et al. (2019), Narayanan et al. (2021), Korthikanti et al. (2022) — 张量并行、3D 并行、序列并行的奠基工作
- DeepSpeed：Rasley et al. (2020) — ZeRO 优化器与多种并行策略的集成框架
- GSPMD：Xu et al. (2021) — Google 的自动并行系统，理解单设备程序到并行程序的转换
- PyTorch FSDP：Zhao et al. (2023) — PyTorch 原生的 ZeRO-3 实现
- Switch Transformer：Fedus et al. (2021) — MoE 训练的工程实践
- DeepSeek-V3 技术报告 (2024) — 万亿参数 MoE 的工业级混合并行实践