生产级分布式训练实战 - 第4节
一、模型并行的动机：显存不够，怎么办？
做分布式训练，第一个绕不开的问题就是：模型太大，单卡装不下。
我刚开始接触大模型时，也天真地以为“多买几张卡就行了”。直到有一次，我把一个 10B 参数的模型往 A100 上塞，结果 OOM 了。嗯，那时候我才真正意识到——显存不是万能的，但没有显存是万万不能的。
你想想看，一个 175B 的 GPT-3，光参数就要 700GB 的显存（FP32）。就算你用混合精度，也得 350GB。而一张 A100 才 80GB。怎么办？
答案就是：把模型拆开，分到多张卡上。这就是模型并行（Model Parallelism）的核心思想。
核心动机：单卡显存不足，无法容纳整个模型。通过将模型切分到多张 GPU 上，每张卡只负责一部分计算。
二、朴素模型并行：Layer-wise 切分
最直观的做法是什么？按层切。
比如一个 12 层的 Transformer，我把前 4 层放到 GPU0，中间 4 层放到 GPU1，最后 4 层放到 GPU2。数据从 GPU0 开始，一层层往后传。这就是朴素模型并行，也叫 Layer-wise 切分。
# 伪代码示例：朴素模型并行
class ModelParallelTransformer(nn.Module):
def __init__(self, num_layers=12, num_gpus=3):
super().__init__()
# 每张卡负责 4 层
self.layers_0 = nn.Sequential(*[TransformerLayer() for _ in range(4)])
self.layers_1 = nn.Sequential(*[TransformerLayer() for _ in range(4)])
self.layers_2 = nn.Sequential(*[TransformerLayer() for _ in range(4)])
# 分配到不同 GPU
self.layers_0 = self.layers_0.to('cuda:0')
self.layers_1 = self.layers_1.to('cuda:1')
self.layers_2 = self.layers_2.to('cuda:2')
def forward(self, x):
# 数据串行经过每张卡
x = x.to('cuda:0')
x = self.layers_0(x)
x = x.to('cuda:1')
x = self.layers_1(x)
x = x.to('cuda:2')
x = self.layers_2(x)
return x这个方法简单吧？但问题也很明显——同一时刻只有一张卡在工作。其他卡都在闲着等数据。GPU 利用率极低。
我曾经踩过的坑：第一次用朴素模型并行跑 40 层 BERT，发现 8 张卡的总吞吐量还不如单卡。原因就是通信开销太大，而且 GPU 利用率不到 20%。说白了，这就是“伪并行”。
三、流水线并行：让所有卡都忙起来
既然朴素模型并行有“空闲等待”的问题，那怎么解决？把数据切成微批次（micro-batch），让流水线转起来。
这就是流水线并行（Pipeline Parallelism）的核心思想。我习惯把它比作工厂流水线——每个工位（GPU）处理完一个零件（micro-batch），就传给下一个工位，同时开始处理下一个零件。
3.1 GPipe：Google 的方案
GPipe 的思路很直接：把一个大 batch 切成多个 micro-batch，依次送入流水线。
举个例子：假设 batch size = 8，我切成 4 个 micro-batch，每个大小为 2。GPU0 处理完 micro-batch 0 后，立刻传给 GPU1，同时开始处理 micro-batch 1。这样所有 GPU 都能同时工作。
# GPipe 流水线调度示意
# 假设 4 个 micro-batch，3 个 GPU 阶段
# 时间步 1: GPU0 处理 mb0
# 时间步 2: GPU0 处理 mb1, GPU1 处理 mb0
# 时间步 3: GPU0 处理 mb2, GPU1 处理 mb1, GPU2 处理 mb0
# 时间步 4: GPU0 处理 mb3, GPU1 处理 mb2, GPU2 处理 mb1
# ... 以此类推GPipe 的优点是实现简单，但有个问题：流水线气泡（pipeline bubble）。在流水线启动和排空阶段，部分 GPU 还是空闲的。
我的经验：micro-batch 数量越多，气泡占比越小。但 micro-batch 太多会增加通信次数。我一般建议 micro-batch 数量是 GPU 数量的 2-4 倍，这样能在利用率和通信开销之间取得平衡。
3.2 PipeDream：更激进的方案
PipeDream 和 GPipe 最大的区别是什么？它不等待所有 micro-batch 完成前向，就开始反向传播。
说白了，GPipe 是“同步流水线”——所有 micro-batch 的前向跑完，再统一跑反向。而 PipeDream 是“异步流水线”——每个 micro-batch 的前向一结束，立刻开始反向。
这样做的好处是：减少了内存占用。因为不需要保存所有 micro-batch 的中间激活值。
但代价也很明显：权重更新不一致。因为不同 micro-batch 可能基于不同版本的权重计算梯度。PipeDream 用了 1F1B（One-Forward-One-Backward）调度策略来解决这个问题。
# PipeDream 1F1B 调度示意
# 每个 GPU 交替执行前向和反向
# GPU0: F(mb0) -> F(mb1) -> B(mb0) -> F(mb2) -> B(mb1) -> ...
# GPU1: 等待 -> F(mb0) -> B(mb0) -> F(mb1) -> B(mb1) -> ...
# 注意：反向传播的梯度会立即用于更新权重我曾经踩过的坑：用 PipeDream 时，发现模型收敛变慢了。原因是异步更新导致梯度噪声增大。后来我加了梯度累积和 warmup 策略，才稳定下来。所以，PipeDream 虽然吞吐高，但对调参要求也更高。
四、GPipe vs PipeDream：怎么选？
| 对比维度 | GPipe | PipeDream |
|---|---|---|
| 调度策略 | 同步流水线 | 异步 1F1B |
| 内存占用 | 高（需保存所有 micro-batch 激活值） | 低（及时释放） |
| 吞吐量 | 中等（气泡较大） | 高（气泡小） |
| 收敛稳定性 | 好（同步更新） | 一般（异步更新） |
| 实现复杂度 | 简单 | 复杂 |
| 适用场景 | 对收敛要求高的场景 | 对吞吐要求高的场景 |
我个人习惯：刚开始做流水线并行时，先用 GPipe。等把整个训练流程跑通、调稳了，再考虑切换到 PipeDream 来压榨性能。毕竟，稳定收敛比那点吞吐提升更重要。
五、核心知识体系：一张图看懂
下面这张图总结了本章的核心逻辑。我把它画成了流程图，方便你理解模型并行的演进路径。
从这张图你可以看到：模型并行的演进路径很清晰。从朴素模型并行到 GPipe，再到 PipeDream，每一步都在解决前一步的问题。但每个方案也都有自己的 trade-off。
一句话总结：模型并行解决的是“放不下”的问题。朴素模型并行是基础但效率低，GPipe 用同步流水线提升利用率，PipeDream 用异步调度进一步压榨性能。选哪个，取决于你对收敛稳定性和吞吐量的权衡。
嗯，这一章的内容就到这里。下一章我们会聊张量并行——另一种更细粒度的模型切分方式。到时候你会发现，流水线并行和张量并行经常搭配使用，效果更好。