流水线并行：按层切分模型的训练策略
流水线并行教程：按层切分模型的训练策略
什么是流水线并行？
流水线并行（Pipeline Parallelism）是一种将深度学习模型按层（或层组）切分到多个计算设备上的模型并行策略。它可以让你训练那些单个设备（如GPU）显存放不下的巨型模型。核心思想很像工厂流水线：每个设备负责模型的一部分连续层，数据像产品一样依次流过各个流水线阶段，每个阶段完成自己的计算后将中间结果传送给下一阶段。
为什么需要流水线并行？
随着大语言模型和视觉Transformer的发展，模型参数动辄上百亿甚至千亿。单个GPU显存有限，无法容纳完整模型、优化器状态和中间激活值。常见的并行策略有：
- 数据并行：每个设备拥有完整模型副本，处理不同数据。显存压力未减小。
- 张量并行：将单层内的矩阵运算切分，通信极密集，扩展性受限于服务器内高速互联。
- 流水线并行：按层切分，每个设备只存储部分层，通信量相对较低，天然适合跨节点扩展。
流水线并行是训练超大模型不可或缺的技术之一，通常与数据并行、张量并行混合使用（3D并行）。
核心原理：按层切分的流水线
朴素流水线及其“气泡”问题
假设一个4层网络，分配到4个GPU上：
- GPU0 存储并执行 Layer 1
- GPU1 存储并执行 Layer 2
- GPU2 存储并执行 Layer 3
- GPU3 存储并执行 Layer 4
正向传播时，一个batch的数据先由GPU0计算，结果传给GPU1，再传给GPU2，最后GPU3计算损失。反向传播时，梯度从GPU3依次传回GPU0。这种串行执行模式导致大部分时间设备处于空闲等待，形成巨大的“流水线气泡”（pipeline bubble），设备利用率极低。
时间 →
GPU0: [F1] [B1] ...
GPU1: [F2] [B2] ...
GPU2: [F3] [B3] ...
GPU3: [F4] [B4] ...
大量空白区即为气泡。我们需要用更智能的调度来填充这些气泡。
微批次（Micro-batch）与 GPipe
GPipe 是谷歌提出的经典流水线并行方案。它将一个训练批次（mini-batch）进一步切分成多个更小的微批次（micro-batch）。假设将mini-batch切为 M 个微批次。
在前向传播中，所有微批次依次注入流水线，各设备一旦完成当前微批次的计算即可立即处理下一个微批次（同阶段），无需等待整个批次走完。反向传播则等所有微批次前向完成后统一进行，梯度最后累积并更新参数。
M=4，4个设备：
GPU0: [F1][F2][F3][F4] [B1][B2][B3][B4]
GPU1: [F1][F2][F3][F4] [B1][B2][B3][B4]
GPU2: [F1][F2][F3][F4] [B1][B2][B3][B4]
GPU3: [F1][F2][F3][F4][B1][B2][B3][B4]
气泡仍然存在（在阶段开始和结束处），但当微批次数量 M 远大于设备数时，气泡占比显著减小。GPipe的优势是实现简单，每次只更新一次参数，优化器状态一致。但缺点是所有激活值需要保留到反向传播结束，造成高峰显存占用与微批次数量成正比。
1F1B 调度与 PipeDream
为了解决GPipe的显存问题，PipeDream 提出了 1F1B（one forward, one backward）调度策略。核心是在一个微批次完成前向传播后，尽早安排其对应的反向传播，而不是等所有微批次前向都完成。
1F1B调度下，设备在稳定状态时交替执行一个前向和一个反向，可以及时释放部分中间激活，降低显存峰值。
M=4，1F1B示意：
GPU0: [F1][F2][F3][F4][B1][B2][B3][B4]
GPU1: [F1][F2][F3][B1][F4][B2][B3][B4]
GPU2: [F1][F2][B1][F3][B2][F4][B3][B4]
GPU3: [F1][B1][F2][B2][F3][B3][F4][B4]
反向传播更早开始，因此显存中无需长期保留全部微批次的激活。PipeDream还允许使用异步通信，进一步掩盖通信延迟。
更高效的现代变种：零气泡与交错的1F1B
- 交错的1F1B（Interleaved 1F1B）：每个设备不仅负责一个连续层段，而是负责多个小的层块（交错分配）。例如4个设备，模型分为8个流水段，每个设备持有两个不连续的块。这样可以进一步减少气泡，但实现更复杂，需要额外的通信。
- 零气泡（Zero Bubble）：通过将反向计算拆分为权重梯度计算和输入梯度计算，利用自动微分中的依赖关系，在流水线空档插入权重梯度计算，将气泡压缩到几乎为零。这是目前前沿研究，被DeepSpeed等框架采用。
如何实现流水线并行？
模型切分
首先需要将模型按顺序划分成多个阶段（stages）。常见做法是：
- 分析模型结构，按层（如Transformer的Encoder层）均匀分割，使每个设备上的计算量大致均衡。
- 处理跨层依赖，如残差连接、层归一化等，必须确保逻辑完整。通常一个阶段包含完整的注意力层和前馈网络层。
通信原语
流水线并行中的通信主要是点对点传递（send/recv）中间激活和梯度。在数据流过相邻阶段时，当前设备将输出张量发送给下一设备，同时从上一设备接收输入张量。使用高效的点对点通信API（如torch.distributed.send和recv），并可与计算重叠来隐藏延迟。
PyTorch 原生流水线API示例（概念性）
PyTorch 2.0+ 提供了 torch.distributed.pipelining 模块简化流水线并行：
from torch.distributed.pipelining import pipe_split, pipeline
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = ...
self.layer2 = ...
self.layer3 = ...
self.layer4 = ...
def forward(self, x):
x = self.layer1(x)
x = pipe_split() # 标记切分点
x = self.layer2(x)
x = pipe_split()
x = self.layer3(x)
x = pipe_split()
x = self.layer4(x)
return x
# 使用pipeline封装，自动处理调度和通信
pipe = pipeline(MyModel(), chunks=8) # chunks即微批次数量
框架会自动进行模型分片和调度，你可以指定调度策略（如GPipe或1F1B）。
与数据并行、张量并行结合
实际训练千亿模型时，单一并行策略不够。常见的3D并行：
- 流水线并行：跨节点，将模型按层分布在不同节点。
- 张量并行：节点内，将层内矩阵乘法切分到多个GPU，利用NVLink高速互联。
- 数据并行：复制整个流水线或流水线组，处理更多数据增加吞吐。
三维组合可以极大扩展训练规模，例如训练GPT-3或 Llama 等模型的标准做法。
挑战与最佳实践
负载均衡
不同层的计算量可能差异大（如注意力与FFN），若按相同层数切分会导致某些GPU等待。应按照计算时间或参数量精细划分阶段，必要时一个阶段包含不同数量的层。部分框架支持自动平衡策略。
显存与通信权衡
- 微批次数量越多，气泡越小，但GPipe会榨取更多显存。1F1B可有效缓解。
- 现代方案可使用激活检查点（Activation Checkpointing）配合流水线，在后向传播时重新计算激活，进一步降低显存。
- 使用高效通信库（如NCCL）和尽可能的重叠通信。
批量归一化（BatchNorm）的问题
BatchNorm依赖跨设备的统计信息，在流水线并行中难以处理。对于Transformer架构多使用LayerNorm，天然适合流水线。如图像模型用到BatchNorm，需要特殊处理，如使用同步BatchNorm或改用GroupNorm。
收敛性和优化器设计
不同调度可能改变梯度计算的数值顺序，但最终梯度累积后是等价的。确保在参数更新前正确累积所有微批次的梯度（GPipe自然保证；1F1B需要显式累积）。使用与数据并行兼容的优化器（如Distributed Optimizer），注意参数更新时间点。
总结
流水线并行通过按层切分模型，将大模型分布到多设备上，突破了单设备显存限制。GPipe和PipeDream（1F1B）是两种基础调度策略，分别侧重简单性与显存效率。现代训练框架集成流水线、张量、数据三种并行，开发者可根据硬件拓扑和模型规模灵活组合。掌握流水线并行，是迈入超大模型训练领域的核心技能。
希望本教程让你对流水线并行有了清晰认识，为将来的模型扩展开启一扇大门。