训练一个千亿参数的大模型，需要数千GPU协同工作数周。如何高效利用这些GPU，是决定训练成本和速度的关键。分布式训练策略就是这个问题的核心答案。本文将系统解析各种并行策略的原理与工程实践。
一、为什么单GPU不够
1.1 显存瓶颈
一个70B参数的模型：
- 模型参数（FP16）：140GB
- 梯度（FP16）：140GB
- 优化器状态（Adam, FP32）：560GB
- 激活值：取决于batch和序列长度
总计约840GB+ ——远超单张A100的80GB显存。
1.2 计算瓶颈
训练70B模型一个epoch（1万亿token）：
- 单A100算力：312 TFLOPS (FP16)
- 训练所需FLOPS：~6×10²³
- 单卡理论时间：~22000天
需要2048张A100并行才能在11天内完成。
二、数据并行（DP）
2.1 基本原理
最直观的并行——每张GPU持有完整模型副本，处理不同数据：
GPU 0: 完整模型 + Batch[0:32]
GPU 1: 完整模型 + Batch[32:64]
GPU 2: 完整模型 + Batch[64:96]
GPU 3: 完整模型 + Batch[96:128]
前向传播 → 各GPU独立计算loss
反向传播 → AllReduce同步梯度
更新 → 各GPU同步更新参数
2.2 问题：大模型装不下
数据并行要求每张GPU能装下完整模型+优化器状态+梯度。对于70B模型，需要840GB——单卡80GB根本装不下。
2.3 ZeRO优化：显存突破
DeepSpeed的ZeRO（Zero Redundancy Optimizer）通过切分训练状态来突破显存限制：
ZeRO-1: 切分优化器状态
每GPU只存 1/N 的优化器状态
70B模型: 560GB/N → N=8时70GB/卡
ZeRO-2: 切分优化器状态 + 梯度
每GPU只存 1/N 的优化器状态和梯度
70B模型: (560+140)GB/N → N=8时87.5GB/卡
ZeRO-3: 切分优化器状态 + 梯度 + 参数
每GPU只存 1/N 的所有状态
70B模型: (560+140+140)GB/N → N=8时105GB/卡
参数在前向/反向时按需All-Gather
# ZeRO-3的参数按需获取
class ZeRO3Layer:
def forward(self, x):
# 按需All-Gather本层参数
full_weight = self.all_gather(self.shard_weight)
output = F.linear(x, full_weight)
# 立即释放完整参数
del full_weight
return output
def backward(self, grad_output):
# 按需All-Gather本层参数
full_weight = self.all_gather(self.shard_weight)
grad_input = grad_output @ full_weight.T
# Reduce-Scatter梯度
self.shard_grad = self.reduce_scatter(grad_output.T @ self.input)
return grad_input
三、张量并行（TP）
3.1 原理
将单个层的参数矩阵切分到多张GPU上：
class TensorParallelLinear:
"""张量并行的线性层"""
def __init__(self, in_features, out_features, n_gpus):
# 按列切分权重
self.shard_size = out_features // n_gpus
self.weight_shard = nn.Parameter(
torch.randn(self.shard_size, in_features) / math.sqrt(in_features)
)
def forward(self, x):
# 各GPU独立计算部分输出
local_output = F.linear(x, self.weight_shard) # (batch, shard_size)
# All-Gather合并完整输出
full_output = self.all_gather(local_output, dim=-1) # (batch, out_features)
return full_output
def backward(self, grad_output):
# 按列切分梯度
local_grad = grad_output[..., self.rank*self.shard_size:(self.rank+1)*self.shard_size]
grad_weight = local_grad.T @ self.input
grad_input = local_grad @ self.weight_shard
return grad_input, grad_weight
3.2 注意力层的张量并行
注意力层更适合按"头"切分：
8个注意力头，4张GPU：
GPU 0: 头 0,1
GPU 1: 头 2,3
GPU 2: 头 4,5
GPU 3: 头 6,7
每张GPU只计算自己的注意力头，最后All-Reduce合并结果。
3.3 TP的通信开销
- 前向：每层需要1次All-Reduce
- 反向：每层需要1次All-Reduce
- 通信量与隐藏维度成正比
TP适合在节点内使用（NVLink连接），跨节点通信开销会抵消并行收益。通常TP_size = 节点内GPU数（如8）。
四、流水线并行（PP）
4.1 原理
将模型按层切分到不同GPU：
GPU 0: Layer 1-12 GPU 1: Layer 13-24 GPU 2: Layer 25-36 GPU 3: Layer 37-48
数据流:
Batch 1 → GPU0 → GPU1 → GPU2 → GPU3 → 输出1
Batch 2 → GPU0 → GPU1 → GPU2 → GPU3 → 输出2
...
4.2 气泡问题
朴素流水线的问题：GPU 0处理Batch 1时，GPU 1-3都在等待——形成"气泡"：
时间: 0 1 2 3 4 5 6 7
GPU0: [B1][B2][..][..][..][..] ← 等待
GPU1: [..][B1][B2][..][..][..] ← 等待
GPU2: [..][..][B1][B2][..][..] ← 等待
GPU3: [..][..][..][B1][B2][..] ← 等待
气泡占比：1 - (micro_batches / (micro_batches + stages - 1))
4.3 1F1B调度
1F1B（One Forward, One Backward）调度减少气泡和显存占用：
阶段1: 先填满流水线（前向）
GPU0: [F1][F2][F3][F4]
GPU1: [F1][F2][F3][F4]
GPU2: [F1][F2][F3][F4]
GPU3: [F1][F2][F3][F4]
阶段2: 1前向1反向交替（稳态）
GPU0: [F5][B1][F6][B2][F7][B3]...
GPU1: [F5][B1][F6][B2][F7][B3]...
阶段3: 排空流水线（反向）
GPU0: [B4][B5][B6][B7]
...
4.4 Interleaved PP
将每个GPU负责的层分成多段，进一步减少气泡：
传统PP: GPU0=L1-12, GPU1=L13-24, GPU2=L25-36, GPU3=L37-48
交错PP: GPU0=L1-6,L25-30, GPU1=L7-12,L31-36, ...
好处: 流水线段数翻倍，气泡减小
代价: 通信次数翻倍
五、3D并行：DP+TP+PP
5.1 组合策略
生产级大模型训练通常同时使用三种并行：
总GPU: 1024张
DP = 16 (数据并行组数)
TP = 8 (张量并行大小，节点内)
PP = 8 (流水线并行段数)
16 × 8 × 8 = 1024 ✓
GPU排列:
节点0: TP组0 (GPU 0-7)
节点1: TP组1 (GPU 8-15)
...
节点0-7: PP组0 (8个节点组成一条流水线)
节点8-15: PP组1
...
16条流水线，每条8节点，每节点8卡
5.2 通信模式
TP通信: 节点内All-Reduce (NVLink, 高带宽)
PP通信: 节点间P2P (InfiniBand, 中带宽)
DP通信: 节点间All-Reduce (InfiniBand, 中带宽)
关键原则：通信最频繁的TP放在带宽最高的节点内，通信较少的DP和PP跨节点。
六、序列并行（SP）
6.1 原理
对于长序列训练，激活值的显存占用很大。序列并行将序列维度切分：
class SequenceParallelAttention:
def forward(self, hidden_states):
# 输入: (batch, seq_len, dim) → 切分为 (batch, seq_len/N, dim)
# 1. 注意力计算需要全局信息 → All-Gather
full_hidden = self.all_gather(hidden_states, dim=1)
# 2. 注意力计算
attn_output = self.attention(full_hidden)
# 3. Reduce-Scatter回切分状态
local_output = self.reduce_scatter(attn_output, dim=1)
return local_output
6.2 与TP的配合
SP通常与TP配合使用——TP切分参数维度，SP切分序列维度，两者正交。
七、实战配置建议
7.1 不同规模模型的推荐配置
| 模型规模 | GPU数 | TP | PP | DP | ZeRO |
|---|---|---|---|---|---|
| 7B | 8 | 1 | 1 | 8 | ZeRO-2 |
| 13B | 16 | 2 | 1 | 8 | ZeRO-2 |
| 70B | 128 | 8 | 4 | 4 | ZeRO-3 |
| 175B | 512 | 8 | 8 | 8 | ZeRO-3 |
| 530B | 1024 | 8 | 16 | 8 | ZeRO-3 |
7.2 超参数调优
# DeepSpeed配置示例
deepspeed:
train_micro_batch_size_per_gpu: 4
gradient_accumulation_steps: 8 # 有效batch = 4 × 8 × DP
zero_optimization:
stage: 3
overlap_comm: true # 通信与计算重叠
contiguous_gradients: true
sub_group_size: 1e9
reduce_bucket_size: 5e8 # 通信桶大小
allgather_bucket_size: 5e8
cpu_offload: false # CPU卸载（省显存但慢）
tensor_parallel:
size: 8
pipeline_parallel:
size: 4
num_stages: 4
pipeline_type: "1f1b_interleaved"
num_chunks: 2 # 交错PP的chunk数
7.3 通信优化
# 通信与计算重叠
class OverlappedCommunicator:
async def forward_with_overlap(self, x):
# 1. 启动当前层的All-Gather（非阻塞）
gather_future = self.all_gather_async(self.weight_shard)
# 2. 同时做上一层的计算
if self.prev_output is not None:
self.prev_output = self.compute(self.prev_output)
# 3. 等待当前层权重就绪
full_weight = await gather_future
# 4. 当前层计算
output = F.linear(x, full_weight)
return output
八、训练效率监控
关键指标:
MFU (Model FLOPs Utilization): 实际FLOPS / 峰值FLOPS
- 优秀: >50%
- 合格: 40-50%
- 需优化: <40%
通信占比: 通信时间 / 总训练时间
- 优秀: <20%
- 合格: 20-35%
- 需优化: >35%
气泡占比: 空闲时间 / 总时间
- 优秀: <10%
- 合格: 10-20%
- 需优化: >20%
九、前沿：异步训练与弹性训练
9.1 异步训练
传统DP是同步的——所有GPU等待梯度同步后才更新。异步训练允许GPU独立更新：
- 优点：消除等待，提高GPU利用率
- 缺点：梯度staleness（旧梯度），可能影响收敛
- 缓解：staleness threshold（限制最大延迟步数）
9.2 弹性训练
训练过程中GPU可能故障。弹性训练允许：
- 动态增减GPU，不中断训练
- 故障GPU自动隔离，剩余GPU继续
- 使用checkpoint恢复机制
结语
分布式训练是大模型工程的核心技术。从简单的数据并行到复杂的3D并行，每一步优化都在为"更大模型、更快训练、更低成本"这个目标服务。理解各种并行策略的原理和适用场景，是构建大模型训练能力的基石。在实践中，没有"最好"的并行策略，只有"最适合"特定模型规模和硬件配置的策略。
本文同步发布于 硅基AGI论坛