三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

1F1B(One-Forward-One-Backward)

1F1B(One-Forward-One-Backward)

1F1B(One-Forward-One-Backward)是流水线并行(Pipeline Parallelism)中最经典的微批次调度策略,主要用于解决传统流水线训练中的显存占用过高气泡率(Bubble Ratio)过大问题。


核心痛点:传统 GPipe 的缺陷

在普通的流水线并行(如 GPipe)中,系统会先一口气把所有微批次(Micro-batch)的前向传播(Forward)全部跑完,再统一跑反向传播(Backward)

  • 显存爆炸:设备必须在显存中同时保存所有 Micro-batch 的激活值(Activations),直到反向传播时才能释放。
  • 气泡严重:卡与卡之间存在大量的空闲等待时间。

1F1B 的工作机制

1F1B 的核心思想是:让前向传播和反向传播在时间线上升高交叉,每跑完一个 Micro-batch 的前向,就紧接着跑一个最老 Micro-batch 的反向。

整个调度过程分为三个阶段:

  • 1. 预热阶段(Warmup)

  • 处于流水线前面的卡先连续跑若干个(通常等于流水线深度p pp)Micro-batch 的前向计算,直到把流水线“填满”。

  • 2. 稳定阶段(1F1B Steady State)

  • 一前一后交替:每执行完1 个 Micro-batch 的 Forward,就立刻接1 个最老的 Micro-batch 的 Backward

  • 显存及时释放:由于 Backward 执行完毕后该 Micro-batch 的激活值就可以立即从显存中销毁,因此显存占用被锁死在一个固定上限(仅需保存约p pp个 Micro-batch 的激活值)。

  • 3. 消退阶段(Cooldown)

  • 所有 Forward 跑完后,把剩余 Micro-batch 的 Backward 执行完毕。


1F1B vs GPipe 性能对比

维度GPipe (Naive PP)1F1B Pipeline
激活值显存峰值O ( m ) O(m)O(m)m mm为总 Micro-batch 数量)O ( p ) O(p)O(p)p pp为流水线 Stage 数量/深度,p ≪ m p \ll mpm
气泡率(Bubble Ratio)p − 1 m \frac{p - 1}{m}mp1p − 1 m \frac{p - 1}{m}mp1(与 GPipe 相当,但显存大幅降低)
调度复杂度简单(顺序执行)中等(需精细管理内部队列和通信)

进阶演进:交错式 1F1B(Interleaved 1F1B)

Megatron-LM 等框架在 1F1B 基础上提出了Interleaved 1F1B

  • 让每张物理 GPU 负责模型中多个非连续的物理层(例如 16 层模型,卡 0 负责第 1 层和第 9 层)。
  • 通过在单卡内部交错调度不同 Virtual Stage 的前向和反向,进一步把流水线气泡率降低到了原本的1 / v 1/v1/vv vv为每张卡上的虚拟 Stage 数量)。
← 返回列表