BigMac 前向按微批切分、反向按流水排空

文章导读
一个微批在前向被按流水段切开、依次穿过所有段,反向却不跟着前向逐段回传,而是等若干微批在途后才从最后一段开始、按相反顺序放行。判断配置是否对得上,只需要盯住三个数:段数 S、微批数 M、反向开始前需要在途的微批数 W。这三个数一旦确定,时间线、队列容量和日志里的编号顺序都应该是可推导的,推不出来就是配置和调度没对齐。
📋 目录
  1. 一 画出一个微批从入口到出口经过的阶段序列
  2. 二 标出反向开始前需要等几个微批进入
  3. 三 用伪代码写出前向入队与反向出队的循环
  4. 四 在日志里按微批编号串起一次完整生命周期
  5. 五 对照时间线检查配置里的段数与微批是否匹配
A A

一个微批在前向被按流水段切开、依次穿过所有段,反向却不跟着前向逐段回传,而是等若干微批在途后才从最后一段开始、按相反顺序放行。判断配置是否对得上,只需要盯住三个数:段数 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说明
21第 0 个微批进入第 1 段即可回传
43稳态区间开始出现 F 与 B 交替
87预热与收尾都明显变长,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 后才从最后一段出队,并按微批编号的相反顺序推进。

BigMac 前向按微批切分、反向按流水排空
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 编号顺序被打乱:反向遍历方向写反了,应从最后一段开始。

对照时间线检查配置里的段数与微批是否匹配

匹配判据可以先用这几条卡一遍,任何一条不满足都先在配置里改,而不是去调计算逻辑:

  1. M 大于 S,并且通常建议 M 明显大于 S,否则稳态区间太短,前向与反向基本没有重叠。
  2. 排空深度 W 与段数对应:1F1B 式取 S-1,交错式按实际在途上限设置,两者不要混用。
  3. 全局 batch 等于单微批大小乘微批数(若有梯度累积再乘累积步数),三者对不上时最后一个微批会变小,时间线尾部会出现长短不一的格子。
  4. 配置里的段数与实际参与计算的设备数、拓扑分组一致,段数多一个就会有一段日志全空。

不匹配时的日志表现也比较固定:M 小于 S 时,反向日志几乎紧贴前向最后一段出现,看不到交替区,收尾段却很长;W 设得过小时反向提前进入,激活还没算完就去取梯度;W 设得过大时在途微批数持续高于设定,前向队列积压、显存占用缓慢上涨;段数与设备数不符时,某一段在整份日志里第一次出现就是反向,前向记录缺失。按这四个现象回头找对应的配置项,比逐行读调度代码更快。改动之后建议保持同一份打点不变,重新跑一小段并在日志里对齐 mb 编号顺序,确认时间线恢复成“斜线预热、交替稳态、单侧收尾”的形状,再放大规模。