一个微批在前向被按流水段切开、依次穿过所有段,反向却不跟着前向逐段回传,而是等若干微批在途后才从最后一段开始、按相反顺序放行。判断配置是否对得上,只需要盯住三个数:段数 S、微批数 M、反向开始前需要在途的微批数 W。这三个数一旦确定,时间线、队列容量和日志里的编号顺序都应该是可推导的,推不出来就是配置和调度没对齐。
「前向按微批切分、反向按流水排空」的本质是把一个全局 batch 拆成 M 个微批,前向逐段推进,反向要等约 S-1 个微批就位后才从最后一段开始回传。适用场景是段数较多、希望前向与反向计算重叠的训练配置。操作上先按 S-1 设排空深度,再用微批编号日志核对实际顺序;边界是 M 太小时流水排不满,M 与 S 不匹配时等待开销会吃掉重叠收益,具体取值需要结合环境确认。
画出一个微批从入口到出口经过的阶段序列
先建立编号参照:前向记 F 加段号,反向记 B 加段号,微批本身用独立的 mb 编号标记,两套编号不要混写。这样时间线上的每个格子都能唯一对应“第几个微批、第几段、前向还是反向”。
| 阶段编号 | 动作 | 输入 | 输出 |
|---|---|---|---|
| F0 | 第 0 段前向 | 该微批的原始输入(token 或已嵌入张量) | 第 0 段激活,交给第 1 段 |
| F1 … F(S-2) | 中间段前向 | 上一段传来的激活 | 本段激活,交给下一段 |
| F(S-1) | 最后一段前向 | 第 S-2 段激活 | 该微批的 loss 或输出张量 |
| B(S-1) | 最后一段反向 | loss 对输出的梯度 | 本段参数梯度 + 传入梯度 |
| B(S-2) … B1 | 中间段反向 | 下一段回传的梯度 | 本段参数梯度 + 继续回传的梯度 |
| B0 | 第 0 段反向 | 第 1 段回传的梯度 | 本段参数梯度,该微批生命周期结束 |
这张表的价值在于:一个微批完整走完需要 2S 个阶段动作,前向 S 步、反向 S 步。当你在日志里看到某个 mb 编号只出现了 F 段而没有对应 B 段,多半是这一轮还没排到它,而不是丢批。
标出反向开始前需要等几个微批进入
在常见的 1F1B 式排空下,最后一个流水段必须先把第 0 个微批的前向算完,才能开始它的反向;而第 0 个微批要穿过 S 段才到得了最后一段,所以反向启动前大约需要 S-1 个微批已经进入流水。这就是排空深度 W 的来源。
| 段数 S | 反向开始前在途微批数 W | 说明 |
|---|---|---|
| 2 | 1 | 第 0 个微批进入第 1 段即可回传 |
| 4 | 3 | 稳态区间开始出现 F 与 B 交替 |
| 8 | 7 | 预热与收尾都明显变长,M 需要足够大 |
如果调度允许每个段同时挂多个在途微批(交错式或多微批在途),排空深度可以小于 S-1,这时必须以实际调度日志为准,不要照搬 S-1。下面的时间线按 1F1B 画,段 0 在最上方、时间向右:
时间 →
段0 F0 F1 F2 F3 B0 B1 B2 B3
段1 F0 F1 F2 F3 B0 B1 B2 B3
段2 F0 F1 F2 F3 B0 B1 B2 B3
段3 F0 F1 F2 F3 B0 B1 B2 B3
^
第 3 个微批进入最后一段,B0 才能开始
图里斜线部分就是预热,中间 F 与 B 交替的一段是稳态,最后只剩 B 的尾巴是收尾。判断配置是否合理,先看有没有稳态:稳态宽度约为 M - (S-1),稳态越窄,重叠带来的收益越容易被预热和收尾吃掉。
用伪代码写出前向入队与反向出队的循环
下面是不绑定具体框架的调度骨架,替换点用注释标出。核心只有两条:新微批只从第 0 段入队;反向只在在途数达到 W 后才从最后一段出队,并按微批编号的相反顺序推进。
S = cfg.pipeline_stages # 段数
M = cfg.micro_batches # 每轮微批数
W = cfg.drain_depth or (S - 1) # 排空深度
fwd_q = [deque() for _ in range(S)] # 每段待处理的前向微批编号
bwd_q = [deque() for _ in range(S)] # 每段待处理的反向微批编号
next_mb = 0
while done_mb < M:
# 1) 前向入队:新微批只进第 0 段,且该段队列不满
if next_mb < M and len(fwd_q[0]) == 0:
fwd_q[0].append(next_mb)
next_mb += 1
# 2) 每段消费队首做前向,产出投递给下一段
for s in range(S):
if fwd_q[s] and has_buffer(s + 1):
mb = fwd_q[s][0]
act = forward(s, mb) # 替换点:本段前向计算
if s + 1 < S:
push_activation(s + 1, mb, act)
else:
bwd_q[S - 1].append(mb) # 最后一段前向完成,进入反向候选
fwd_q[s].popleft()
# 3) 反向出队:在途数够 W 才启动,从最后一段往第 0 段推进
if in_flight() >= W:
for s in reversed(range(S)):
if bwd_q[s] and has_grad(s - 1):
mb = bwd_q[s].popleft()
g = backward(s, mb) # 替换点:本段反向计算
if s > 0:
push_gradient(s - 1, mb, g)
else:
done_mb += 1 # 该微批生命周期结束
两个容易写错的地方:一是把 next_mb 的递增放在投递之后,导致重复入队;二是反向循环里从第 0 段开始遍历,那样在途微批会被打乱顺序,表现为日志里 B 段的 mb 编号忽大忽小。修改后用工单里的 mb 编号日志确认顺序,再谈调参。
在日志里按微批编号串起一次完整生命周期
打点方式建议固定四个字段:mb、stage、phase(F 或 B)、event(enter 或 exit),再加一个单调递增的 seq 方便排序。每个段单独输出一行会更好比对,多段混在一个文件里时要保证 seq 全局递增。下面是字段形态示意,用于说明该长什么样,不是某次运行的实际输出。
seq=0001 mb=0 stage=0 phase=F event=enter
seq=0002 mb=0 stage=0 phase=F event=exit
seq=0003 mb=1 stage=0 phase=F event=enter
...
seq=00xx mb=0 stage=3 phase=F event=exit
seq=00xx mb=0 stage=3 phase=B event=enter
seq=00xx mb=0 stage=2 phase=B event=enter
按 mb 分组、按 seq 排序,你应当看到两段单调:前向阶段 stage 从 0 递增到 S-1,反向阶段 stage 从 S-1 递减到 0;同一 mb 的 B 段起点必须晚于它的 F 段终点。出现下面几种情况就说明调度和时间线不一致:
- 同一 mb 的 B enter 早于它的 F(S-1) exit:排空深度设得太浅,或者反向启动条件判断错。
- mb 编号跳号或重复:入队位置写错,通常是 next_mb 递增时机不对。
- 某段长期只有 enter 没有 exit:该段计算没走完或被下游阻塞,检查缓冲区容量。
- 反向阶段的 mb 编号顺序被打乱:反向遍历方向写反了,应从最后一段开始。
对照时间线检查配置里的段数与微批是否匹配
匹配判据可以先用这几条卡一遍,任何一条不满足都先在配置里改,而不是去调计算逻辑:
- M 大于 S,并且通常建议 M 明显大于 S,否则稳态区间太短,前向与反向基本没有重叠。
- 排空深度 W 与段数对应:1F1B 式取 S-1,交错式按实际在途上限设置,两者不要混用。
- 全局 batch 等于单微批大小乘微批数(若有梯度累积再乘累积步数),三者对不上时最后一个微批会变小,时间线尾部会出现长短不一的格子。
- 配置里的段数与实际参与计算的设备数、拓扑分组一致,段数多一个就会有一段日志全空。
不匹配时的日志表现也比较固定:M 小于 S 时,反向日志几乎紧贴前向最后一段出现,看不到交替区,收尾段却很长;W 设得过小时反向提前进入,激活还没算完就去取梯度;W 设得过大时在途微批数持续高于设定,前向队列积压、显存占用缓慢上涨;段数与设备数不符时,某一段在整份日志里第一次出现就是反向,前向记录缺失。按这四个现象回头找对应的配置项,比逐行读调度代码更快。改动之后建议保持同一份打点不变,重新跑一小段并在日志里对齐 mb 编号顺序,确认时间线恢复成“斜线预热、交替稳态、单侧收尾”的形状,再放大规模。