自动并行策略搜索 PP 流水并行建模与负载均衡能力支持

本文档描述自动并行策略搜索中 PP 流水并行建模与负载均衡能力的设计需求，包括 pp_degree、micro batch number、pipeline bubble、stage partition、PPB 负载均衡算法、layer offset 和 layer recompute 等能力。不包含具体代码实现，侧重 PP 维度在端到端自动并行策略搜索方案中的建模语义、接口契约、测试设计与验收标准。

1. 背景

Pipeline Parallelism，PP，通过将模型层切分到多个 pipeline stage 上，降低单卡显存压力并支持更大模型训练。PP 的实际性能不仅取决于 pp_degree，还强依赖 stage partition、micro batch number、pipeline bubble、stage 间通信、layer offset 和 layer recompute 策略。

在自动并行策略搜索中，PP 不应只是一个简单的 pp_degree 变量，而应包含以下能力：

• 估算每个 pipeline stage 的显存；
• 估算每个 pipeline stage 的计算时间；
• 建模 micro batch number 对 pipeline bubble 和 activation 保留数量的影响；
• 根据 layer cost 自动生成更均衡的 stage partition；
• 支持 layer offset 调整 stage 边界；
• 支持 layer recompute 降低 activation 显存并影响计算时间；
• 将上述 PP 策略作为端到端自动并行策略搜索方案的一部分参与搜索。

其中，PPB 是 PP 维度下的流水并行负载均衡算法能力，用于根据层级代价信息生成更均衡的 stage partition。

2. 目标与非目标

2.1 目标

• 支持 pp_degree 和 micro_batch_num 进入自动并行策略搜索空间；
• 支持 stage_partition 的表达、校验和估算；
• 支持 PPB 流水并行负载均衡算法，根据 layer cost 生成更均衡的 stage partition；
• 支持 layer offset 作为 stage 边界调整维度进入搜索；
• 支持 layer recompute 作为显存优化维度进入搜索；
• 支持每个 pipeline stage 的显存估算；
• 支持每个 pipeline stage 的 compute time 和 communication time 估算；
• 支持 pipeline bubble 估算；
• 支持 1F1B 等典型调度下的 step time 估算；
• 支持 stage 显存超限过滤；
• 支持 PP schedule在不同负载情况下的可视化模拟；
• 补充单元测试和集成测试。

2.2 非目标

• 本 issue 不实现新的 pipeline runtime；
• 本 issue 不负责 DP/TP/EP/CP 的独立估算；
• 本 issue 不保证覆盖所有 pipeline schedule；
• 本 issue 不负责训练执行器的底层调度改造。

3. 建模语义

3.1 PP 内存估算

PP 内存估算应覆盖：

• 每个 pipeline stage 的参数显存；
• 每个 stage 的梯度显存；
• 每个 stage 的 optimizer state 显存；
• 每个 stage 的 activation 显存；
• micro batch number 对 activation 保留数量的影响；
• layer recompute 对 activation memory 的降低效果；
• layer recompute 带来的额外 compute cost；
• stage partition 和 layer offset 对每个 stage memory 的影响。

3.2 PP 性能估算

PP 性能估算应覆盖：

• 每个 stage 的 forward compute time；
• 每个 stage 的 backward compute time；
• stage 间 activation / gradient 通信开销；
• pipeline bubble；
• micro batch number 对流水效率的影响；
• 1F1B 等调度方式下的端到端 step time；
• layer recompute 带来的额外计算时间；
• stage 负载不均衡对 step time 的影响。

3.3 PPB 负载均衡算法

PP 建模 / PPB 负载均衡算法输出应包含以下信息：

输入	说明
model layer list	模型层级结构，至少包含 layer / block 的顺序信息。
layer compute cost	每层计算代价，可来自 symbolic model、profiling 结果或用户配置。
layer memory cost	每层显存代价，包括参数、激活、临时 buffer 等估算值。
layer communication cost	层间通信或 stage 边界通信代价，可选。
pp_degree	流水并行 stage 数量。
memory_limit	每个 stage 的显存上限，可选。
layer_recompute_config	重计算策略对激活显存与计算代价的影响。
layer_offset_config	stage 边界偏移搜索范围。
micro_batch_num	micro batch 数量，用于评估 pipeline bubble。

PPB 算法输出应包含以下信息：

输出	说明
stage_partition	每个 layer / block 到 pipeline stage 的映射关系。
stage_compute_cost	每个 stage 的累计计算代价。
stage_memory_cost	每个 stage 的累计显存代价。
stage_comm_cost	每个 stage 或 stage 边界通信代价，可选。
layer_offset	实际采用的 stage 边界偏移策略。
layer_recompute	实际采用的 layer recompute 策略。
imbalance_score	stage 间负载不均衡指标，例如 max / avg 或 max - min。
is_feasible	是否满足显存等硬约束。

3.4 layer offset 建模

layer offset 用于调整流水 stage 边界，使 stage partition 不局限于简单均分。应支持：

• 在给定 pp_degree 下调整各 stage 的 layer 数；
• 在不改变 layer 顺序的前提下移动 stage 边界；
• 与 stage memory、stage time、imbalance_score 联动评估；
• 作为候选策略的一部分输出；
• 对非法 offset 返回 infeasible 或可读错误。

3.5 layer recompute 建模

layer recompute 用于通过重计算降低 activation 显存。应支持：

• 指定哪些 layer / block 开启 recompute；
• 估算 recompute 对 activation memory 的降低效果；
• 估算 recompute 对 compute time 的额外开销；
• 与 micro batch number、stage memory 和 stage time 联动评估；
• 作为候选策略的一部分输出。

3.6 约束

应检查：

• pp_degree 不超过可切分 layer 数；
• stage_partition 覆盖所有 layer 且不重复；
• 每个 stage 至少包含一个有效 layer，除非配置允许空 stage；
• layer offset 不应导致非法 stage partition；
• layer recompute 不应引用不存在的 layer；
• global_batch_size、micro_batch_size、dp_degree、micro_batch_num 关系合法；
• 每个 stage 的显存估算不超过 memory limit。

4. 实现要点

• 应输出 stage 负载、pipeline bubble 和 schedule 相关的可视化数据，供上层报告或图表工具使用；
• PP 建模模块应支持 stage-level memory breakdown；
• PP 建模模块应支持 stage-level time breakdown；
• PPB 输出格式应与自动并行策略搜索器保持一致；
• layer offset 和 layer recompute 应作为 PP 策略配置的一部分参与搜索；
• 对 stage 显存超限策略，应返回 infeasible 状态和具体 stage 信息；
• 估算结果应可与 DP/TP/EP/CP 的估算结果组合。

5. 测试设计

用例 ID	描述	期望
AP-PP-01	pp_degree=2/4，默认均匀切分	能输出每个 stage 的 memory 和 time。
AP-PP-02	使用 PPB 生成 stage_partition	stage 负载比简单均分更均衡。
AP-PP-03	改变 micro_batch_num	pipeline bubble 和 estimated step time 趋势正确。
AP-PP-04	stage 显存超限	策略被标记为 infeasible，并说明超限 stage。
AP-PP-05	stage_partition 覆盖不完整	配置校验失败。
AP-PP-06	pp_degree 大于可切分 layer 数	配置校验失败或返回明确错误。
AP-PP-07	开启 layer offset 搜索	stage 边界可被调整，且输出 offset 策略。
AP-PP-08	非法 layer offset	策略校验失败并说明原因。
AP-PP-09	开启 layer recompute	activation memory 下降，compute time 上升趋势符合预期。
AP-PP-10	非法 layer recompute 配置	策略校验失败并说明原因。
AP-PP-11	输出 PP 可视化数据	能生成 stage 负载、pipeline bubble 或 schedule 相关的可视化数据。

6. 验收标准

• 能输出 stage 负载、pipeline bubble 或 schedule 相关的可视化数据；
• pp_degree 和 micro_batch_num 可进入自动并行策略搜索空间；
• PPB 能在给定 layer cost 和 pp_degree 条件下输出合法 stage_partition；
• layer offset 可进入 PP 策略搜索，并能调整 stage 边界；
• layer recompute 可进入 PP 策略搜索，并能影响 activation memory 和 compute time；
• 能输出每个 stage 的 memory cost 和 performance cost；
• 能估算 pipeline bubble 和 step time；
• stage 显存超限可被识别并过滤；
• AP-PP-01～AP-PP-10 测试通过；
• 文档中包含 PP 建模范围、stage_partition 输入输出格式、micro batch 影响、layer offset、layer recompute 和 PPB 负载均衡说明。

7. 参考

• GPipe / PipeDream：Pipeline Parallelism；
• Megatron-LM：1F1B pipeline schedule；
• HyperParallel：自动并行策略搜索相关模块；
• Pipeline stage partition 与流水并行负载均衡算法。