跳到主要内容

数据截至 (上游 commit e79cb4c1bae1)

03 · 流水并行

这一章讲什么： PP 把模型按层切成 p 段分到 p 张卡后，真正的问题不再是「怎么切」而是「怎么排」——哪个 microbatch 在哪张卡、什么时候算前向、什么时候算反向。Megatron 给三个调度器：无流水、1F1B、interleaved 1F1B。读完你会理解「气泡」从哪来、interleaved 怎么把它再砍一刀，以及段间 P2P 通信的工程细节。

1. 它要解决的小问题​

把 64 层模型切 8 段，每张卡只放 8 层。朴素做法：把 batch 拆成 m 个 microbatch，先全部前向流完 8 段、再全部反向流回来。

问题有两个：

  • 气泡（bubble）：流水刚启动时，只有第 1 段有活干，后面 7 张卡干等——前 p-1 个时间步的算力被浪费，反向收尾时同理。气泡占比约 (p-1)/m。
  • 显存：先跑完所有前向意味着第 1 段要同时存住全部 m 个 microbatch 的激活，等反向来了才能释放。m 一大就爆。

两个问题的解药都是「排期」：让反向尽早开始（在途 microbatch 数就有上限，激活显存就有界），以及让每卡持有的段变碎（同样的 pipeline 深度下，每段算得更快，气泡窗口更短）。

2. 思路：1F1B 与 interleaved​

1F1B（one-forward-one-backward）​

把 m 个 microbatch 的执行分成三段：

阶段干什么在途激活数
warmup只做前向，把流水灌满递增到上限
稳态 1F1B每算 1 个前向，立刻算 1 个「最老」的反向恒定
cooldown只做反向，把流水排空递减到 0

warmup 的长度 = p - rank - 1：越靠前的段 warmup 越长（它要等后面所有段都准备好才能开始收反向梯度）。稳态时「前向一个、反向一个」让每卡同时在飞的激活数恒定——激活显存与 m 解耦，这是 1F1B 相比 GPipe 式「先全前向」的最大收益。

interleaved（虚拟流水）​

每张卡不再持有一段连续的层，而是持有 v 个不连续的「模型块」（model chunk）。比如 p=4、v=2、16 层：GPU0 拿第 [1,2] 层和第 [9,10] 层（parallel_state.py:600 的 docstring 里有这个完整例子）。

效果是流水深度从 p 变成 p×v，每「段」的计算量除以 v：

  • 气泡占比从 (p-1)/m 降到约 (p-1)/(v·m)；
  • 代价是 P2P 通信次数乘 v（每个 chunk 边界都要收发）。

图示：三种排期的形态​

怎么读：横轴时间、纵轴是 4 个流水段；F=前向，B=反向，数字是 microbatch 号，.= 空转（气泡）。

GPipe 式 (先全 F 后全 B):        1F1B:                      interleaved (v=2):
rank0 F0 F1 F2 F3 B0 B1 B2 B3    rank0 F0 F1 F2 B0 F3 B1    每卡两块, 段更碎
rank1 .  F0 F1 F2 .  B0 B1 B2    rank1 .  F0 F1 B0 F2 B1    气泡窗口按比例缩短
rank2 .  .  F0 F1 .  .  B0 B1    rank2 .  .  F0 B0 F1 B0
rank3 .  .  .  F0 .  .  .  B0    rank3 .  .  .  F0 B0 ...
      └ 启动/收尾大气泡 ┘             └ warmup 短, 稳态 1F1B ┘

3. 原理演示：稳态循环的骨架​

1F1B 稳态的逻辑，简化后不到 20 行。这段演示「前向一个 → 收反向梯度 → 反向最老的一个 → 把它的梯度发回去」的交替：

# 示意，非源码
for i in range(num_microbatches_remaining):
    output = forward_step(next(data_iter), recv_tensor)   # ① 前向一个新 microbatch
    grad_of_oldest = send_forward_recv_backward(output)   # ② 发走输出, 换回最老 mb 的输出梯度
    stash.append((recv_tensor, output))                   #    新的入队尾
    oldest_in, oldest_out = stash.pop(0)                  # ③ 队首 = 最老, 轮到它反向
    input_grad = backward_step(oldest_in, oldest_out, grad_of_oldest)
    recv_tensor = send_backward_recv_forward(input_grad)  # ④ 发回梯度, 顺便收下一个输入

重点看 ②④ 两步各把一次 send 和一次 recv 焊在一起：发前向输出的同时等反向梯度、发反向梯度的同时等下一个前向输入。这就是「稳态时通信刚好配对」的代码形态。

4. 真实实现​

4.1 三选一调度器​

get_forward_backward_func（megatron/core/pipeline_parallel/schedules.py:53）末尾就是全部选择逻辑（schedules.py:161-168）：

  • pp_size > 1 且有虚拟流水 → forward_backward_pipelining_with_interleaving（:1019）；
  • pp_size > 1 无虚拟流水 → forward_backward_pipelining_without_interleaving（:2147）；
  • 否则 → forward_backward_no_pipelining（:723）。

4.2 非交错 1F1B：三段循环​

forward_backward_pipelining_without_interleaving（schedules.py:2147）的结构和 §2 的表一一对应：

  • warmup：num_warmup_microbatches = total_stages - current_stage - 1（schedules.py:2277），然后 for i in range(num_warmup_microbatches) 纯前向 + send_forward（:2333 起）。warmup 里收到的 input_tensor/output_tensor 逐个 append 进列表——这就是反向时要排队的「在途激活」。
  • 稳态：for i in range(num_microbatches_remaining)（:2379），循环体就是 §3 那段示意代码的真实版：send_forward_recv_backward（:2416）换回梯度 → input_tensors.pop(0) 取最老 → backward_step → send_backward_recv_forward（:2446）。
  • cooldown：schedules.py:2450 起，只反向 + send_backward，把队列清空。

4.3 interleaved：warmup 公式里的 v​

交错版的 microbatch 计数多乘一个 num_model_chunks：total_num_microbatches = num_microbatches * num_model_chunks（get_pp_rank_microbatches，schedules.py:929、:951）。

warmup 公式在 schedules.py:964-965：

num_warmup_microbatches = (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
num_warmup_microbatches += (num_model_chunks - 1) * microbatch_group_size_per_vp_stage

与非交错的 p - rank - 1（:957）对比着看：(p-rank-1)*2 是环形往返的代价，(v-1)*group_size 是先把每个 chunk 的第一轮前向铺完的代价。交错版的三段循环本体分别在 schedules.py:1542（warmup）、:1731（稳态）、:1954（cooldown）。

4.4 段间 P2P：batch_isend_irecv​

相邻段的通信收在 P2PCommunicator（megatron/core/pipeline_parallel/p2p_communication.py:145）里。构造时算出自己在 PP 组里的前驱/后继全局 rank（:157-164），之后所有 send_forward/recv_backward/send_forward_recv_backward（:494、:462、:536）都汇到 _communicate（:280）。

_communicate 的核心只有一件事：先 torch.empty 开好接收 buffer（requires_grad=True，dtype 用 config.pipeline_dtype，:336-352），再调 _batched_p2p_ops（:22）把本拍要发/要收的全部 isend/irecv 塞进 torch.distributed.batch_isend_irecv 一次发起——成批发起是为了让 NCCL 把多个 P2P 操作融合调度，避免两两死锁。

4.5 一个容易被忽略的显存技巧​

输出张量发给下一段之后，本段只剩反向还要用它的 grad_fn，数据本体已经没用了。deallocate_output_tensor（schedules.py:171）把 .data 设成标量「伪释放」——docstring 原话：「only useful for its .grad_fn field」。在 warmup/稳态每次 send_forward 后调用（如 :2368、:2423）。流水越深、microbatch 越多，这个技巧省得越多。

5. 关键细节与坑​

  • 收发双方形状必须事先谈拢。 接收 buffer 是 torch.empty(shape) 开出来的，形状错了直接错位。固定序列长度时形状由 get_tensor_shapes（schedules.py:2115）统一算出；变长序列要走 _communicate_shapes（p2p_communication.py:191）先发形状再发数据——这是 variable_seq_lengths 配置的用途。
  • interleaved 要求 model 是 list。 forward_backward_pipelining_with_interleaving 开头直接断言 isinstance(model, list) 且每个 chunk 都要有独立 data_iterator（schedules.py:1074-1082）——从「一个模型」到「模型块列表」是调用方的责任。
  • overlap_p2p_comm 与 batch_p2p_comm 互斥。 交错版里两个都开会直接 ValueError（schedules.py:1086-1087）。
  • microbatch 数太少时 interleaved 退化。 warmup 数被 min 到 total（:976-977），m 太小会导致全部 microbatch 都在 warmup/cooldown，稳态为空——气泡优势归零。经验上 m 要显著大于 p。
  • 梯度同步被拆进调度缝隙。 稳态里 enable_grad_sync() 只在「最后一个 microbatch 的反向快完成」时打开（:2434）——之前所有 microbatch 的 DP 梯度都靠 buffer 分桶异步发（见 04 章），调度器只负责临门一脚。
  • 反向用的是 output_tensor_grad 配对，不是重算。 队列里存的 (input, output) 对和 recv 回来的梯度按 FIFO 严格配对（pop(0)），任何一跳的收发顺序错了就是静默错梯度——P2P 的成批配对（§4.4）正是为了堵这个洞。

6. 代码地图​

主题文件路径符号名
调度器三选一megatron/core/pipeline_parallel/schedules.pyget_forward_backward_func
非交错 1F1Bmegatron/core/pipeline_parallel/schedules.pyforward_backward_pipelining_without_interleaving
interleaved 1F1Bmegatron/core/pipeline_parallel/schedules.pyforward_backward_pipelining_with_interleaving
warmup/稳态 microbatch 计数megatron/core/pipeline_parallel/schedules.pyget_pp_rank_microbatches、get_schedule_table
无流水（PP=1）megatron/core/pipeline_parallel/schedules.pyforward_backward_no_pipelining
输出张量伪释放megatron/core/pipeline_parallel/schedules.pydeallocate_output_tensor
收发形状计算megatron/core/pipeline_parallel/schedules.pyget_tensor_shapes
P2P 通信器megatron/core/pipeline_parallel/p2p_communication.pyP2PCommunicator、_communicate、_communicate_shapes
成批 isend/irecvmegatron/core/pipeline_parallel/p2p_communication.py_batched_p2p_ops
前向/反向单步megatron/core/pipeline_parallel/schedules.pyforward_step、backward_step

下一章：04 · 分布式优化器——梯度收下之后，ZeRO 式分片怎么把优化器状态摊到 DP 组。