流水线并行与气泡分析:把模型按深度切到 rank

流水线并行与气泡分析:把模型按深度切到 rank

本节摘要:张量并行把矩阵乘切到 rank,流水线并行把模型切到 rank,每 rank 一段(stage),微批(microbatch)流过流水线。开头与结尾的空闲时间是气泡,压小它是全部手艺。本节把一个串行模型切成 N 段、在 N rank 上模拟 forward 流水线,用 GPipe 调度(先全填 forward 再反向)排程 M 个微批并算气泡占比 (N-1)/(M+N-1),对比 Megatron-LM/PipeDream 的交错 1F1B 调度。读完本节,你能说清为什么 M=8/N=4 气泡 27% 而 M=64/N=4 仅 4.5%、为什么均摊每段**计算量**比均摊参数量更重要(embedding 参数多计算少)、以及为什么 send-recv 要偶 rank 先发后收、奇 rank 先收后发以防死锁。

对应原课程:Phase 19 · Lesson 79 · pipeline-parallel(原英文 phases/19-capstone-projects/79-pipeline-parallel/docs/en.md)。本节属第 20 章「毕业项目」的分布式训练赛道。

学习目标

阅读完本节,你应当能够:

1. 把串行模型切成 N 段,在 N rank 上模拟 forward 流水线。
2. 用 GPipe 调度(先全填 forward 再反向)排程 M 个微批,并算气泡占比。
3. 对比 Megatron-LM 与 PipeDream 用的交错 1F1B 调度。
4. 辩护段分配:每段计算量相等比每段参数量相等更重要。

一、问题与直觉

一个 700 亿参数模型 fp16 光参数就要 140GB,没有消费级 GPU 装得下。ZeRO-3 把参数分片到 rank,但每 forward 仍要每 rank 对每层 allgather 全层,付每层 log(N) 跳。流水线并行走另一条路:把模型切成 N 段,每段放一个 rank。第 0 段 forward 完把激活张量交给 rank 1;rank 1 跑第 2 层交 rank 2;依此类推,反向逆序流。每 rank 只持一段,显存线性降;计算是串行的,这就是气泡问题。

气泡是流水线开头(等首批微批到末段)与结尾(等末批反流回来)的空闲。M 个微批、N 段时,每段气泡占比 (N-1)/(M+N-1)。M=8、N=4 是 27%;M=64、N=4 是 4.5%。气泡随每步微批数增多而缩,意味着每微批 batch 要小,这是驱动微批设计的约束。

二、从零实现

GPipe 调度

先把全部 M 个微批的 forward 填满流水线再启动任何反向,然后反向逆序排空。每个微批的激活须保留到它的反向,故显存随 M 线性增长。forward 占 M+N-1 周期,反向再 M+N-1 周期;每段有用功 2M 周期,每段气泡 2(N-1) 周期。每 forward 与 backward 各一单位时间时,气泡占比 (N-1)/(M+N-1)。取 M 远大于 N 即可藏气泡。

1F1B 调度

交错:某微批 forward 一到末段就启动它的反向,让它回流。调度每段交替一个 forward 一个 backward。气泡仍是 N-1,但激活显存以流水线深度为界而非微批数。生产流水线用 1F1B(Megatron、PipeDream)。本节先实现 GPipe(更简单),1F1B 作练习。

为什么每段计算量相等要紧

若段 0 耗 50ms、段 1 耗 100ms,每个周期被段 1 门控,其他段每周期闲 50ms 等段 1 释放。每段参数量相等是错轴:transformer 的计算由每层注意力加 MLP 主导,embedding 层参数多但计算少。段分配应均摊每段 FLOPs,而非每段权重。

微批 vs batch

流水线跑 M 个大小 B 的微批,有效 batch 是 MB。流水线步末梯度是 MB 个样本合并的梯度。气泡占比取决于 M,优化器看到 M*B。调 M 是在气泡(高 M 低)与每微批显存(GPipe 高 M 高激活)间换。

实现骨架

class Pipeline:
    def __init__(self, stages, num_microbatches):
        self.stages, self.M = stages, num_microbatches

    def gpipe_schedule(self):
        # forward 填满:M + N - 1 周期
        for m in range(self.M):
            for s in range(len(self.stages)):
                self.stages[s].forward(m)         # 微批 m 在段 s
        # backward 排空:M + N - 1 周期,逆序
        for m in reversed(range(self.M)):
            for s in reversed(range(len(self.stages))):
                self.stages[s].backward(m)

def bubble_fraction(num_stages, num_microbatches):
    return (num_stages - 1) / (num_microbatches + num_stages - 1)
​

code/main.py 实现:PipelineStage(持一段参数、暴露 forward(activation))、Pipeline(GPipe 调度,模拟每段墙钟)、bubble_fraction(闭式 (N-1)/(M+N-1))、4 段 demo 打印每微批轨迹与实测气泡。运行 python3 code/main.py 输出段×微批甘特图与气泡百分比对比闭式预测。

三、框架对比

业界对比:Megatron-LM 是大规模流水线参考,用 1F1B,支持张量+流水线+数据并行组合;DeepSpeed Pipeline 与 ZeRO 集成,ZeRO-1 + 流水线是最大开源模型常见组合;PyTorch Pipe 是 PyTorch 原生流水线 wrapper(基于 torch.distributed.pipeline.sync.Pipe)。它们的共识与本节一致:气泡占比 (N-1)/(M+N-1)、生产用 1F1B、段按 FLOPs 均摊、M 远大于 N 藏气泡。

四、生产中的硬化模式

三种模式把流水线并行硬化到可上线。

激活检查点与流水线配对。 GPipe 上 M 个微批在飞,激活显存是单微批的 M 倍。激活检查点在反向时重算 forward,用计算换显存;这个组合让长序列的流水线可行。

段均衡是测的,不是假设的。 生产团队跑一次剖析,测目标硬件上每层实际计算(FLOPs 与墙钟),再按测量分区。Megatron-LM 的 --num-layers-per-stage 接受列表,允许段每层成本不同时不均层数。

send-recv 调度必须防死锁。 每段先发后收的流水线在线缆上死锁。标准修法是交错:偶 rank 段先发后收,奇 rank 段先收后发。本节显式排程 rank 让模式可见。

五、可复用产物

• PipelineStage + Pipeline + bubble_fraction:本节是 81 节端到端 demo 的流水线那一半(为运行时把流水线保持模拟),也是 bubble_fraction 闭式公式的可复用工具。
• 段分配原则(按 FLOPs 不按权重):任何流水线划分的起点。
• GPipe/1F1B 调度模板:GPipe 先实现(简单),1F1B 作生产升级(激活显存有界)。

六、练习

1. 实现 1F1B:验证气泡占比与 GPipe 相同但激活显存有界。
2. 剖析真实每段时间:在更深模型上测,按墙钟再均衡段。
3. 加梯度累积:跨流水线微批累梯度,验证等于等价全批 forward 的梯度。
4. 配激活检查点:测显存降 vs 计算成本。
5. 流水线 + DDP:每个流水线 rank 在数据并行组上复制,推理二维调度。

本节要点回顾

1. 流水线并行按深度切:每 rank 一段,激活逐段流;显存线性降,计算串行是气泡问题。
2. 气泡占比 (N-1)/(M+N-1):M=8/N=4 是 27%,M=64/N=4 是 4.5%;**M 远大于 N 藏气泡**。
3. GPipe 调度:先全填 forward 再反向排空,激活随 M 线性;每段有用功 2M、气泡 2(N-1)。
4. 1F1B 调度:forward 一到末段就反向,激活以流水线深度为界;生产用(Megatron/PipeDream)。
5. 段按 FLOPs 均摊不按权重:embedding 参数多计算少;段 0 快段 1 慢则每周期门控在慢段。
6. 微批 vs batch:有效 batch = MB,优化器看 MB,气泡看 M;调 M 是气泡 vs 激活显存。
7. 三大硬化:激活检查点配对、段均衡实测、send-recv 偶发奇收防死锁。

下一节,我们将进入「分片检查点与恢复」——把 ZeRO 与流水线的分片状态存成带清单的检查点,在崩溃后按同 world size 恢复,让万步训练不怕中断。