混合并行背后的负载均衡问题

1. 问题从哪里来

如果只看概念，DP、CP、TP、PP、EP 都很清楚。但真正把它们放到一个训练系统里，麻烦就来了：这些维度不是互相独立的。CP 能让长序列的 attention 放进显存，但会带来 P2P 通信；DP 能增加样本并行度，但所有 rank 最后还是要在梯度同步处碰头；PP 能切层，但 microbatch 时间不均匀时就会产生 bubble。

更麻烦的是 post-training 里的序列长度经常很不均匀。一批数据里可能既有很短的 prompt，也有很长的 reasoning trace。如果仍然把每个 microbatch 当成差不多的计算量，那么调度器看到的是整齐的张量，硬件看到的却是很不均匀的工作。

所以本文只讨论一个窄问题：在可变长度训练中，DP+CP 的负载均衡应该怎么理解？Megatron Dynamic-CP 解决了哪一层问题？ByteScale 又把问题放大到了什么范围？

2. 先看最朴素的 DP balance

为了让问题简单一点，先假设系统已经支持 sequence packing。也就是说，一个 microbatch 里可以塞多条真实序列，调度器至少知道每条序列有多长，而不是只看到一个已经 pad 或 pack 之后的 tensor。

这时候最容易想到的办法，是先不动 CP，只在 DP rank 之间重新分配样本。比如按 token 数排序，或者按 token 数估一个计算 cost，然后用 longest-processing-time 这类启发式：每次把当前最重的样本或 microbatch 放到当前最空的 rank 上。这个思路很朴素，但也很有效，至少能避免长样本全堆到同一个 rank，导致别的 rank 在同步点空等。

不过这里有一个关键限制：DP 不是普通的任务队列。一次 gradient accumulation 结束时，所有 DP rank 要对齐到同一个梯度同步边界。如果一个 rank 跑了 8 个 microbatch，另一个 rank 跑了 7 个，那么 forward/backward、loss scaling、通信顺序都要重新定义。于是很多 DP balance 的实现，实际上是在一个更窄的空间里做平衡：每个 rank 的 microbatch 数量仍然相同，只是尽量让每一步的 token work 接近。

Slime 的实现就是一个比较典型的工程版本。它的 DP/microbatch scheduler 先把 rollout 样本组织成 training step，再用固定 chunking 或 first-fit token packing 生成 microbatch。之后如果打开 --balance-data，它会用 Karmarkar-Karp partitioner 在 DP rank 之间平衡 microbatch token sum，同时保留每个 rank 的 microbatch 数量一致。换句话说，它 balance 的不是抽象样本，而是训练循环真正会执行的 microbatch。

但这个方案里，CP 仍然更像一个固定容量因子。它可以帮助系统容纳更长序列，但还没有让每条序列自己决定要占几个 CP rank。这就是 Megatron Dynamic-CP 往前推进的地方。

3. Megatron：在固定池子里做动态 CP

Megatron 的 hybrid context parallelism 不是简单地把 CP 固定成一个全局值。它在初始化时仍然会基于 DPxCP rank 预建一组 process group，但运行时会先把 packed batch 解开，根据 cumulative sequence length 还原出一段段真实序列，也就是代码里的 sub-sample。然后，调度器会给每个 sub-sample 分配 local_cp_size：短序列可能只用 1 个 rank，长序列可能用同一个 DPxCP 池里的 2、4、8 个甚至更多 rank。

ByteScale 论文图 3 解释了固定 DP+CP 的一个典型问题：短序列可能承担不必要的 CP 通信，长短不均又会在 DP/PP 调度里形成 bubble。

从代码上看，相关路径比较集中：parallel_state.py 负责创建 hybrid DP-CP process group，data_schedule.py 负责解包和重路由 sub-sample，hybrid_cp_schedule.py 会用近似的 seq_len * seq_len / cp_size 来估算 attention work，再把 CP size 向上取整到 2 的幂，并在兼容 group 之间插入 barrier。NVIDIA 的 Dynamic-CP 技术博客 也描述了同一套系统语义：data iterator wrapper 重新调度 packed data，选择 CP size，返回有效的 num_micro_batches，并把动态 packing 元数据广播到各个 pipeline stage。

这里容易误解的一点是，Dynamic-CP 并不是重新设计了一套 pipeline parallelism。PP 仍然走 Megatron 原来的 PP/VPP schedule；Dynamic-CP 的作用，是尽量把送进 PP 的 microbatch 做得更均衡，并保证各个 stage 拿到一致的动态元数据。这样可以减少某些长 microbatch 对相邻 stage 的阻塞，但 pipeline 语义本身并没有被替换。

所以总结一下，Megatron 的关键变化不是“支持 CP”本身，而是让 CP 在一个预建好的 DPxCP group 池子里变成了 per-sub-sample 的选择。限制也同样明确：可选 group 主要是 2 的幂；调度还要服从 Megatron 原有的 loss scaling、梯度同步和 pipeline 执行边界。

4. ByteScale：把 rank 分配也纳入调度

ByteScale 这里要先说明一个限制：我没有看到开源训练实现，所以这部分只能基于论文中的设计描述来比较，不能算代码级验证。

它和 Megatron Dynamic-CP 最核心的区别，是调度范围更大。Megatron 已经允许有效的 num_micro_batches 随 iteration 变化，也会广播 PP 元数据来保证 pipeline stage 一致；但在运行时搜索里，DP rank 仍然基本要保持相同 microbatch count，CP choice 也来自预先建好的 2 的幂 group。

ByteScale 的 Hybrid Data Parallelism 则进一步放松了这个约束。它的 balance scheduler 会把更多序列分给估计执行时间更短的 rank 或 pipeline，再通过全局 batch 内所有 token 上累计的梯度来保持训练语义。论文里把这件事拆成 DP-Balance 和 PP-Balance：前者处理 HDP rank 之间每个 time step 的负载均衡，后者还会调整 length bucket 的顺序，让相邻 pipeline stage 的 microbatch runtime 更接近。

这也是为什么我认为 ByteScale 不是简单做了一个“减少 PP bubble”的小优化。它更像是在重新定义一次 iteration 的执行计划：序列怎么分配，rank 怎么参与，哪些 rank 需要通信，梯度如何累计，pipeline 如何排队。收益来自更大的调度自由度，代价也是同一件事：scheduler state、loss normalization、gradient sync boundary、FSDP unshard/reshard timing 和 PP scheduling 都要理解同一个计划。

ByteScale 论文图 8 展示的是另一种视角：不再把 DP+CP 看成固定网格，而是看成一个可以为不同序列分配不同 rank 数的 HDP 池。

论文还描述了更灵活的 rank assignment：一条序列可以使用 [1, d_hdp] 内任意数量的 rank，P2P 通信只发生在参与这条序列的 rank 之间。至于 selective activation offloading 和 PP bubble reduction，论文里还有更多细节；这里为了避免过度展开，只把它们作为 ByteScale 调度范围更大的旁证。

5. 小结

混合并行最容易给人一种错觉：维度越多，系统越强。但从负载均衡角度看，事情刚好没有这么简单。每多一个切分维度，都可能带来新的 collective、P2P、元数据和同步点。如果这些额外成本没有换来足够的显存空间或吞吐提升，那这个维度就只是把系统变复杂了。

反过来，如果调度器真的理解这些维度之间的耦合关系，多一个维度也可能是好事。短序列可以少用 CP，长序列可以多用 CP；空闲 rank 可以多接一点序列；pipeline stage 可以看到更接近的 microbatch runtime。关键不在于“有没有这个并行轴”，而在于这个轴是否进入了同一个调度问题。

因此，我更倾向于把 DP、CP、PP 理解成一次 iteration execution plan 里的几个字段，而不是几组互相正交的配置项。真正需要回答的是：哪些 token 去哪些 rank，哪些 rank 之间通信，什么时候同步梯度，以及这个计划是否仍然对应同一个 global-batch gradient 语义。Megatron Dynamic-CP 和 ByteScale 的差别，本质上就是这张计划表覆盖范围的差别。