第 70 题：如何处理PP中的负载不均衡？recompute和no-recompute层的分配？
题目
如何处理PP中的负载不均衡？recompute和no-recompute层的分配？
完整讲解
一、PP 负载不均衡的来源
不同 stage 的层数、参数量、计算量不同（如 attention 与 FFN 比例、embed 与 head 数），导致各 stage 的 F/B 时间不一致，快的 stage 等慢的，bubble 或空闲增加。此外，recompute（activation checkpointing）只对部分层做，有 recompute 的 stage backward 时要多一次重算，时间变长，若分配不当会加重不均衡。
二、Recompute 与 no-recompute 的分配
- Recompute 层：计算相对便宜、激活大的层（如部分 attention、大 hidden 的 FFN）适合做 checkpoint，用算力换显存；若把这些层集中到少数 stage，这些 stage 的 backward 会明显变长。
- No-recompute 层：小激活或计算贵的层不 checkpoint；若集中到某 stage，该 stage 显存压力大但算得快。策略：尽量让各 stage 的「计算+通信+重算」时间接近，例如把部分 recompute 层与 no-recompute 层交错到不同 stage，或按实测时间微调 stage 切分点。
三、其他手段
- Stage 划分：按层数或按「预估 F+B 时间」切分，使各 stage 耗时接近；工具可 profile 各层耗时再划分。
- Interleaved：同一设备多子 stage 交错，从统计上平滑单设备内负载。
- Pipeline 调度：1F1B 等已能减轻 bubble；再结合均衡的 stage 与合理的 recompute 分配，可进一步减少等待。
面试要点
- 负载不均衡来自各 stage 层数/计算/重算不同；需让各 stage 总耗时接近。
- Recompute 放「大激活、小计算」层；避免把大量 recompute 集中到同一 stage。
- 手段：按耗时划分 stage、interleaved、合理分配 recompute/no-recompute。
记忆要点
- 不均衡 = 各 stage F+B+recompute 时间不一致；按耗时划分 stage。
- Recompute 分散到各 stage，避免单 stage 重算过多。
- Interleaved + 1F1B + 均衡划分 一起用。