AI 原生开放教材

外观

AI 工程从零到一（中文）源码研读·中文原文 ↗原件 ↓仓库 ↗

Pipeline Parallel 与 Bubble 分析 ​

Tensor parallelism 把矩阵乘法切到各个 rank 上。Pipeline parallelism 把模型切到各个 rank 上，一个 rank 一个 stage。Microbatch 在 pipeline 里流动。开头和结尾那段空闲就是 bubble；把它压到最小，就是这门手艺的全部。

类型： Build 语言： Python 前置要求： 阶段19 Track C 第42-49课 预计时间： ~90 分钟

学习目标 ​

  • 把一个顺序模型切成 N 个 stage，模拟一条横跨 N 个 rank 的 forward pipeline。
  • 用 GPipe 调度（只填 forward，再做 backward）把 M 个 microbatch 调度过 pipeline，并算出 bubble 比例。
  • 把 bubble 和 Megatron-LM、PipeDream 用的交错式 1F1B 调度做对比。
  • 论证 stage 分配：每个 stage 计算量相等，比每个 stage 参数量相等更重要。

问题背景 ​

一个 fp16 的 700 亿参数模型，光参数就要 140 GB。没有消费级 GPU 装得下。ZeRO-3 把参数 shard 到各个 rank 上，但每次 forward 仍要每个 rank 把整层 allgather 出来，每层付 log(N) 跳。Pipeline parallel 走另一条路：把模型切成 N 个 stage，每个 rank 放一个 stage。第 1 层的 forward 在 rank 0 上算完，把激活张量交给 rank 1；rank 1 跑第 2 层再交给 rank 2；以此类推。Backward 反向流动。内存线性下降，因为每个 rank 只持有一个 stage；计算是顺序的，这就是 bubble 问题。

Bubble 是 pipeline 开头的空闲（等第一个 microbatch 到达最后一个 stage）和结尾的空闲（等最后一个 microbatch 反向排空）。有 M 个 microbatch 和 N 个 stage 时，每 stage 的 bubble 比例是 (N-1)/(M+N-1)。在 M=8、N=4 时是 27%。在 M=64、N=4 时是 4.5%。每步 microbatch 越多 bubble 越小，这意味着每 microbatch 的 batch size 要小，而这个约束正是驱动 microbatch 设计的因素。

核心概念 ​

mermaid
flowchart LR
  R0[rank 0：stage 0 / layer 0] --> R1[rank 1：stage 1 / layer 1]
  R1 --> R2[rank 2：stage 2 / layer 2]
  R2 --> R3[rank 3：stage 3 / loss]
  R3 -.backward.-> R2
  R2 -.backward.-> R1
  R1 -.backward.-> R0

GPipe 调度 ​

在开始任何 backward 之前，先用全部 M 个 microbatch 把 pipeline 正向填满；然后反向排空 backward。每个 microbatch 的激活值必须留到它的 backward 才能释放，所以内存随 M 线性增长。Forward 占 M+N-1 个周期，backward 再占 M+N-1 个周期。每 stage 的有用工作是 2M 周期；每 stage 的 bubble 是 2(N-1) 周期。当每次 forward 和 backward 各占一个时间单位时，bubble 比例是 (N-1)/(M+N-1)。选 M 远大于 N 就能把 bubble 藏起来。

1F1B 调度 ​

交错进行：一个 microbatch 的 forward 一到达最后一个 stage，就启动它的 backward 并让它反向流回。这个调度在每个 stage 上交替做一次 forward 和一次 backward。Bubble 仍是 N-1，但激活内存被 pipeline 深度限住，而不是 microbatch 数量。生产 pipeline 用 1F1B（Megatron、PipeDream）。本课先实现 GPipe 因为它更简单，1F1B 留作练习。

为什么每 stage 计算量相等很重要 ​

如果 stage 0 要 50 ms、stage 1 要 100 ms，每个周期都被 stage 1 卡着。其他 stage 每周期空等 50 ms 等 stage 1 放行。参数量相等是错误的衡量轴：transformer 的计算由每层的 attention 加 MLP 主导，而 embedding 层参数多但计算少。Stage 分配应该让每个 stage 的 FLOPs 相等，而不是权重相等。

Microbatch 与 batch 的关系 ​

一条 pipeline 跑 M 个大小为 B 的 microbatch。有效 batch size 是 MB。Pipeline 一步结束时的梯度，就是在合起来的 MB 个样本上的梯度。Bubble 比例取决于 M；optimizer 看到的是 M*B。调 M 就是在 bubble（M 高则更低）和每 microbatch 内存（在 GPipe 下 M 高则激活内存更高）之间权衡。

figure
cd-pipeline-bubble

动手构建 ​

code/main.py 实现了：

  • PipelineStage：一个小 nn.Module，持有一个 stage 的参数并暴露 forward(activation)。
  • Pipeline(stages, num_microbatches)：用每个 stage 的模拟墙钟时间，在模拟 stage 上编排 GPipe 调度。
  • bubble_fraction(num_stages, num_microbatches)：闭式解 (N-1)/(M+N-1)。
  • 一个 4-stage demo，打印逐 microbatch 的 trace 和实测的 bubble 比例。

运行：

bash
python3 code/main.py

输出：一张 stage 对 microbatch 的甘特图，以及 bubble 百分比与闭式预测的对比。

真实世界中的生产模式 ​

有三个模式能把 pipeline parallel 打磨到可以上线。

activation checkpointing 要和 pipeline 配套。 在 GPipe 上有 M 个 microbatch 在飞时，激活内存是单个 microbatch 的 M 倍。Activation checkpointing 在 backward 时重算 forward，拿计算换内存；正是这个组合让 pipeline 对长序列变得可行。

stage 平衡是测出来的，不是假设出来的。 生产团队会跑一遍 profiling，在目标硬件上测出实际的逐层计算量（FLOPs 和墙钟），再按这个测量值来切分。Megatron-LM 的 --num-layers-per-stage flag 接受一个列表，允许在各 stage 单层成本不同时使用不均的层数。

send-recv 调度必须避免死锁。 一条每个 stage 都先 send 再 recv 的 pipeline 会在线路上死锁。标准修法是交错：偶数 rank 的 stage 先 send 再 recv，奇数 rank 的 stage 先 recv 再 send。本课显式地排 rank 的顺序，让这个模式看得见。

实际使用 ​

生产模式：

  • Megatron-LM。 大规模 pipeline parallel 的标杆。用 1F1B，支持 tensor + pipeline + data parallel 三者组合。
  • DeepSpeed Pipeline。 和 ZeRO 集成；ZeRO-1 + pipeline 是最大那批开源模型的常见组合。
  • PyTorch Pipe。 PyTorch 原生的 pipeline wrapper，构建于 torch.distributed.pipeline.sync.Pipe。

拿去用 ​

第 80 课把每个 stage 的参数 shard 存进 sharded checkpoint。第 81 课在端到端 demo 上把 DDP + ZeRO + pipeline 组装起来（在精神上如此；为了运行时间，demo 把 pipeline 保持为模拟）。

练习 ​

  1. 实现 1F1B，验证 bubble 比例与 GPipe 一致，但激活内存被限住。
  2. 在一个更深的模型上 profile 真实的每 stage 时间，按实测墙钟重新平衡 stage。
  3. 加上跨 pipeline microbatch 的梯度累积，检查梯度等于等价的整 batch forward 的梯度。
  4. 把 pipeline 和 activation checkpointing 配在一起，测量内存下降与计算代价的对比。
  5. 把 pipeline 和 DDP 组合（每个 pipeline rank 在一个数据并行组里被复制），推理这套 2D 调度。

关键术语 ​

术语大家怎么说实际含义
Pipeline"沿深度的模型并行"一个 rank 一个 stage，激活值在 stage 间流动
Bubble"pipeline 空闲时间"开头 + 结尾各 (N-1) 步，部分 stage 没活干
Microbatch"batch 的一片"一个 forward/backward 单元；M 越大 bubble 越小
GPipe"先填后排"全部 M 个 forward 在任何 backward 之前；激活内存高
1F1B"交错调度"每个 stage 一次 forward 一次 backward；激活内存有界

延伸阅读 ​