📚 课程目录第3章核心概念

3. 核心概念

在流水线并行中，模型被切分成多个 stage，分布到不同 GPU 上，micro-batch 像流水线上的工件一样依次流过各 stage。理想情况下，每个 stage 的计算时间完全相等，流水线可以一直保持满载。但真实训练中，这种理想状态几乎不存在，负载不均衡会直接转化为气泡，吞噬吞吐量。理解负载不均衡的来源，是掌握 DeepSpeed 流水线调度的第一步。

一、负载不均衡的来源

负载不均衡的本质，是各 stage 在每个 micro-batch 上消耗的时间不一致。造成这种不一致的原因主要有以下几类：

模型切分粒度不均：按层数平均切分时，不同层的参数量、计算量差异很大。例如 Transformer 中注意力层与前馈层的 FLOPs 不同，若切分边界落在计算量差异大的位置，stage 间就会失衡。
算子与内存差异：某些 stage 包含 embedding、LayerNorm 或大矩阵乘，计算密度高；另一些 stage 可能受显存带宽限制，实际执行时间更长。
通信开销不对称：相邻 stage 之间的 P2P 通信量取决于激活张量大小。若某条边传输的激活特别大，该 stage 的等待时间就会增加。
硬件异构：集群中 GPU 型号、NVLink 拓扑、PCIe 带宽不一致时，相同计算量在不同设备上的耗时不同。
动态输入长度：NLP 任务中序列长度可变，长序列 micro-batch 会让对应 stage 的计算时间骤增。

负载不均衡的度量可以用“最大 stage 时间 / 平均 stage 时间”表示。该比值越接近 1，流水线越均衡；比值越大，气泡越多。

二、静态调度与动态调度对比

静态调度在训练开始前就确定每个 micro-batch 在每个 stage 上的执行顺序，典型代表是 1F1B（One Forward One Backward）及其变体。它实现简单、通信模式固定，但无法感知运行时各 stage 的实际耗时变化。动态调度则在运行时根据各 stage 的队列长度、剩余计算量和通信状态，动态决定下一个执行 forward 还是 backward，从而填补空闲槽位。

对比维度	静态调度	动态调度
调度决策时机	训练前固定	运行时逐 micro-batch 决策
对负载变化的适应	弱，依赖人工切分	强，可自动填补气泡
实现复杂度	低	高，需维护状态与优先级
通信模式	规整、可预测	可能不规则，需处理竞争
典型代表	1F1B、GPipe	DeepSpeed 动态调度、PipeDream
适用场景	stage 耗时稳定、同构集群	异构、变长输入、负载波动大

需要强调的是，动态调度并非总是优于静态调度。当各 stage 耗时非常稳定时，静态调度的规整通信反而能减少调度开销和显存碎片。选择哪种方式，取决于负载波动的幅度与频率。

三、气泡与空闲时间

气泡（bubble）指流水线中某个 stage 因等待上游数据或下游梯度而处于空闲的时间片。气泡主要出现在两个阶段：流水线启动时的 warm-up 阶段，以及排空时的 cool-down 阶段。此外，运行中若某 stage 计算慢，其他 stage 也会被迫等待，形成运行中气泡。

以 GPipe 的朴素流水线为例，若 stage 数为 p，micro-batch 数为 m，则气泡比例约为 (p-1)/(m+p-1)。增大 m 可以稀释气泡，但会带来显存压力。1F1B 通过交错执行 forward 和 backward，把 warm-up 和 cool-down 压缩到 p-1 个 micro-batch 的规模，显著降低气泡比例。

# 气泡比例估算示例
p = 8      # stage 数
m = 32     # micro-batch 数
bubble_ratio = (p - 1) / (m + p - 1)
print(bubble_ratio)   # 约 0.179，即 17.9% 的时间是气泡

降低气泡的常用手段：增大 micro-batch 数、采用 1F1B 或交错式 1F1B、使用动态调度填补空闲槽位、对 stage 进行更细粒度的负载均衡切分。

四、吞吐量与延迟的权衡

吞吐量指单位时间内完成的样本数，延迟指单个样本从进入流水线到输出结果的时间。两者在流水线并行中往往相互制约。

增大 micro-batch 数：吞吐量上升，气泡比例下降，但单个样本的端到端延迟增加，因为要排队等待更多 micro-batch。
减少 stage 数：流水线深度变浅，延迟降低，但单卡显存压力增大，可能被迫减小 batch size，反而影响吞吐。
动态调度：通常能提升吞吐量，但调度决策本身引入额外开销，可能轻微增加延迟。
梯度累积：与流水线结合时，累积步数越多，吞吐越高，但延迟线性增长。

因此，调优目标不是单独最大化吞吐或最小化延迟，而是在给定显存和集群条件下找到帕累托最优点。训练任务通常更关注吞吐量，而在线推理服务更关注延迟。

五、案例与复盘：一次负载不均衡的排查

某团队在 8 卡上训练一个 24 层 Transformer，采用按层均分，stage 划分是 3-3-3-3-3-3-3-3。训练日志显示 GPU4 的利用率长期在 95% 以上，而 GPU7 只有 60% 左右，整体吞吐比预期低约 25%。

排查步骤：

用 torch.profiler 采集各 stage 单 micro-batch 耗时，发现 stage5 的 forward 时间是 stage0 的 1.8 倍。
检查模型结构，发现 stage5 包含两个大 FFN 层，参数量和 FLOPs 明显高于其他 stage。
重新按计算量而非层数切分，将 stage5 的一层移到 stage6，使各 stage 耗时方差下降约 70%。
同时把 micro-batch 数从 16 提高到 32，气泡比例从 30% 降到约 18%。

常见误区：只按层数平均切分就认为负载均衡。实际上应按 FLOPs、参数量、激活大小和通信量综合评估。另一个误区是盲目增大 micro-batch 数，导致显存溢出或延迟不可接受。

复盘结论：负载不均衡的根因往往在模型切分阶段就已埋下。静态调度下，切分质量决定上限；动态调度可以在一定程度上补偿，但不能替代合理的切分。理解气泡的来源和吞吐-延迟的权衡关系，才能在 DeepSpeed 配置中做出正确取舍。