借助动态上下文并行和 NVIDIA Megatron Core 加速可变长度训练

本文介绍了应用于 NVIDIA Megatron Core 中的一种新型调度方法 — — 动态上下文并行 (Dynamic-CP) ，由可灵团队和 NVIDIA 技术团队共同开发，其主要用于 LLM 后训练以及 DiT 预训练。Dynamic-CP 通过在微批次（micro-batch）粒度上动态选择上下文并行（CP）规模，高效应对可变长度序列训练问题，并在LLM训练上实现了高达 1.48 倍的加速。

在大模型训练中，真实数据集中序列长度的分布不均成为了一个经常被忽视的性能瓶颈。无论是在 LLM 训练还是在大规模视频生成任务中，序列长度均呈现出明显的长尾分布特征：少量超长样本带来了不成比例的计算开销和显存消耗。

在 LLM 训练中，这种现象表现为不同 batch 之间文本序列长度差异显著；在视频生成任务中也是如此——从几百 token 的短片段到数万 token 的高清长视频，序列长度跨度巨大。上述特性直接导致计算效率和显存占用在数据并行（DP）rank间、不同模态以及微批次之间严重失衡，从而显著制约了调度效率和整体资源利用率。

为应对可变长度的输入，现有训练系统通常采用样本级打包（sample-level packing）策略，即将多个较短序列拼接为一个微批次，并限制其总 token 长度不超过目标序列长度。如图 1 所示，几个不同长度的序列被打包成了三个具有相同长度的微批次。

图 1. 打包序列与未打包序列的对比

对比未打包序列和已打包序列

尽管这三个打包后的样本具有相同的总长度，但由于点积注意力的计算复杂度与序列长度是二次关系，它们的计算负载并不相同，如图 2 所示。打包后的样本之间计算负载的这种差异被称为数据并行计算不均衡。该不均衡会导致 GPU 空闲，即部分 DP rank 需要等待计算负载更高的 DP rank 完成梯度同步。此外，这种计算不均衡还会进一步加剧流水线并行（Pipeline Parallel，PP）气泡，从而降低整体训练效率。

图 2. 打包序列之间的 attention 计算不均衡

在打包序列中，每个样本的注意力计算工作负载各不相同，导致 GPU 空闲。

在图 3 中，NVIDIA Nsight Systems 的 profiling 显示了 VLM 训练中的负载不均衡。其中不同的图像/ 视频样本具有可变的序列长度，且已对序列进行了打包以减少填充。图3展示了不同 DP 通信组之间的同步开销。

图 3. 不同 DP 通信组之间同步的开销

该图展示了 Nsight System profiling timeline，其中由于序列长度可变带来的不均衡，产生了额外的同步开销。

此外，在使用上下文并行时，为了避免GPU内存不足，CP 的切分大小往往由最长的序列决定。结果是，那些本不需要上下文并行的较短序列也会被切分，即使这些序列在单个 GPU上就能存下，他们仍会因为存在更长的序列而被分片，从而带来不必要的 CP 通信开销。

通常情况下，CP 的通信开销可以被计算掩盖。但是当 CP 较大时 (尤其是跨 InfiniBand (IB) 域的通信时) ，若打包后的序列较短、计算负载较小时，通信开销就会”暴露“出来，这会使 CP 的计算低效。

以下示例展示了由于总序列长度较大，需要使用 TP2CP8 的配置。许多打包好的序列由较小的子序列组成，计算量不足以掩盖 CP 通信。

图 4. 由于计算量不足，无法在序列打包的情况下隐藏 CP 通信

该图展示了 Nsight System 捕获的 timeline：对所有序列都使用了相同的且较大的 CP 大小。可以看到对于某些样本，NCCL 相关算子比 cuDNN 相关的计算算子用时更长，这时 CP 通信无法被计算隐藏。

这些结果表明，需要一种对上下文并行（CP）进行动态调整的方法。由此，本文提出了动态上下文并行（Dynamic-CP），这种方法不是静态地将 CP 大小固定为一个最长的序列所需的值，而是在 micro batch 粒度动态地调整 CP 大小和打包策略。相关工作（如ByteScale 和 WLB-LLM）也在解决类似问题。

切换 CP 大小需要重新划分序列切片，并重组 attention 计算所需的CP通信组。与其他的动态并行方案 (例如根据序列长度调整 TP 或 PP 大小) 相比，Dynamic-CP 增加的额外开销最少，因为调整 TP/ PP 大小往往需要在不同的 GPU 上重新分发权重或重构流水线计算图，这些操作代价很大。

为了决定如何打包序列，需要设计一个求解器。该求解器旨在给定一组可变长度序列，确定如何打包它们并选择 CP 大小，以在不超过 GPU 显存限制的情况下最大限度地提高计算效率。求解器的功能是获取可变长度序列，并计算最佳打包策略和 CP 大小。这种方法可更大限度地提高计算效率，同时不超过GPU 显存限制。通过对计算和通信成本进行建模，求解器可以避免对短序列过度切分以减少不必要的 CP 通信，从而缓解数据并行不均衡和 CP 低效问题。

以下示例展示了 Dynamic-CP 的优势。若不使用 Dynamic-CP，负载不均会导致不同 micro batch 之间出现流水线空泡（pipeline bubble），进一步引发了不同 DP rank之间的不均衡。开启 Dynamic-CP 后，micro batch 和 DP rank之间的空泡显著减少。

图 5. Dynamic-CP 方法的优势

该图展示了 Dynamic-CP 方法的优势。原本在数据并行 (DP) rank 和 micro batch 之间会出现空泡，并在整个 pipeline 中进一步放大。而 Dynamic-CP 方法可以有效地减小这些空泡。

用于支持 Dynamic-CP 的 Megatron Core 框架修改

在这部分中，我们会介绍将 Dynamic-CP 集成到 Megatron Core 的整体流程。

图 6. 将 Dynamic-CP 集成到 Megatron Core 中

该图展示了为将 Dynamic-CP 集成到 Megatron Core 中所做的修改。

为每个 rank 构建多个 CP group

使用标准的上下文并行时，每个 rank 属于具有固定 cp_size 的单个组 (cp_group) ，且 cp_size 在初始化期间静态确定。但是，Dynamic-CP 在不同 iteration 和 micro batch 中具有不同的 cp_size。

为此，单个 rank 必须参与多个不同规模的 CP 组。在初始化过程中构建多个 CP 组，cp_size 介于 1 到 max_cp_size 之间，且限制为2的幂次方。这种设计允许在运行时根据打包和调度结果选择合适的 CP 组，不会产生动态创建通信组的开销。

动态重新调度和打包数据

预训练的数据通常使用 Batch x Sequence x Head x Dim (BSHD) 布局，而 Dynamic-CP 则在 THD 布局上运行。在这种格式中，变长序列会在长度约束下被打包在一起，将原始 BS 维度折叠成为 T 维度（total token）。

因此，micro batch 的数量不再是静态的。在 BSHD 中，可以通过下式算出 micro batch 的数量： num_microbatches= global_batch_size/dp_size/micro_batch_size。

而使用 THD 时，每个打包序列中包含的原始序列数量并不固定，这会导致 num_microbatches 在随着训练步数迭代发生变化。

Megatron-Core 会根据是否启用流水线并行（PP/VPP）提供多种训练调度器。为了尽量减少对现有调度逻辑的侵入式修改，我们引入了一个围绕原始 data_iterator 的轻量级 data_iterator_wrapper。该 wrapper 执行三个步骤：

重新调度与打包Global batch中的序列，使 DP rank 间的工作负载尽量均衡。
根据打包结果选择合适的 cp_size，以最大限度地减少 CP 通信的低效性。
为当前迭代返回有效的 num_microbatches。

通过这种方式，只需插入一个 wrapper，就能为所有调度器添加 Dynamic-CP 支持，同时保持原有调度代码基本不变。

在流水线各阶段间广播，并扩展 PackedSeqParams

由于 num_microbatches 会发生变化，并且只有 TP rank 0以及第一个和最后一个 PP stage 会有data_iterator ，因此框架需要向所有相关的 PP rank 广播 num_micro_batches、max_seqlen 和 cu_seqlens。确保在动态 micro batch 调度下流水线的一致执行。

使用 Dynamic-CP 时， cp_size 会随迭代变化，因此依赖全局静态 CP 设置是不合适的。为解决此问题，PackedSeqParams 被进一步扩展，增加了 cp_size 和 cp_group 两个变量。

所有依赖上下文并行的组件——例如位置编码和 Transformer Engine 的 attention——现在都会从 PackedSeqParams 中获取 CP 配置，替代原先的全局 CP 变量。这保证了所有与 CP 相关的操作都与动态选择的 CP 布局保持一致。

Loss 计算和 FLOPS 统计

在变长序列和 THD 布局下，不同序列包含的有效 token 数量不同。因此，loss 需要按 token 进行归一化计算：loss = loss_over_valid_tokens/total_number_valid_tokens。它可以避免填充引入的偏差。

此前版本的 Megatron-Core 在计算 FLOPs 时并未考虑 THD 布局，而是假设 max_seqlen 就是有效序列长度，导致在变长场景下对 FLOPs 产生系统性的高估。

数据调度建模

Transformer 工作负载与序列长度呈现二次关系 ( $\mathrm{O}(S^2)$ )。与此同时，激活显存随线性增长 ( $\mathrm{O}(S^2)$ )，这意味着即使序列长度只有很小的波动，也可能在不同 DP rank 和不同 micro-batch 之间造成显著的计算与显存不均衡。为了平衡一个大样本的工作负载，我们可能会把多个小样本打包在一起，但这会带来严重的显存压力。由于无法同时让 FLOPs 与显存做到完全均衡，这就决定了调度与 packing 策略必须在两者之间权衡。

我们的目标是实现理想的均衡分布：在 DP rank 与 micro-batch 之间尽可能均匀地切分工作负载和显存占用。在每个 DP rank 的 micro-batch 数量固定时，可以为每个 micro-batch 设定目标的工作量与显存配额。随后，一个三阶段调度器会在“工作负载”和“显存”目标之间交替优化，并在需要时为更重的样本增大 CP 大小，以改善计算与显存的平衡。

工作负载模型、求解器和模拟器的协作

完整的调度流程由三个组件组成：

工作负载模型：根据每个样本的序列长度来估算其执行时间，并对 Transformer 各类算子上单个样本工作负载进行建模。这定义了最基本的负载单元，其准确性会直接影响最终的性能收益。
求解器：以工作负载模型的输出作为输入，使用启发式算法为各样本确定近似最优的 packing 策略。随后将打包后的样本组成 micro-batch，并为其分配 CP 大小。每个 DP rank 的 micro-batch 数量会影响最终效果，包括 pipeline bubble、流水线并行不均衡导致的bubble 以及数据并行不均衡导致的bubble。通过遍历不同的“每 DP rank 的 micro-batch 数量”，可以找到最优结果。
模拟器：在分布式流水线并行调度下评估这些 micro-batch，选择执行时间最短（即工作负载最均衡）的方案，同时满足峰值显存约束。

建模过程和双目标平衡

理想的均衡分布，是在不同 DP rank 以及不同 micro-batch 之间，将计算工作量与显存占用尽可能均匀地切分。在各个 DP rank 采用相同 micro-batch 数量的前提下，可以确定每个 micro-batch 的目标工作量配额与显存配额。同时，pipeline bubble 的大小也会有所不同，需要将其在各个 micro-batch 之间尽量均匀分摊，才能实现端到端的整体均衡。

如下所示，均衡 DP rank中的端到端训练时间完全相同，需满足：

是第 i 个 DP rank 的 micro batch 数量
是第 i 个 DP rank 中一个 virtual pipeline stage 的每个 micro batch 的工作负载
虚拟流水线阶段的数量
是流水线阶段编号。

各 rank 的工作负载配额满足：

同时，全局批量样本在一个 virtual pipeline stage 的总工作负载可表示为：

结合 2 和 3，就可以确定不同 DP 的每个 micro batch 的工作负载。

随着序列长度增长，计算负载按 $\mathrm{O}(S^2)$ 增长，而内存消耗按 $\mathcal{O}(S)$ 级增长，因此很难同时实现工作负载和内存的平衡。因此，求解器会在不同阶段交替优化“工作量优先”和“显存优先”两个目标，逐步逼近一个兼顾二者的平衡解。

如果单个工作负载超过 micro batch 工作负载配额，系统会为其分配更大的 CP 大小。完成此步骤后，工作负载的不均衡会减小，而显存成为主要约束。然后，目标会转移到显存，选择计算量最少的样本来填充存储桶。其余样本按降序排序，并使用相同的启发式算法分配给每个 micro batch。

零开销运行

在运行时，调度工作流不能在训练循环中引入可感知的额外开销。实际系统需要解决两个问题：I/O 压力与求解器的运行开销。

I/O 压力

首先，构建调度方案需要对 global batch 额外进行一次 get_item 遍历，以收集序列长度与形状信息。为缓解 I/O 压力，系统采用两种互补技术：将探测用的 get_item 分摊到集群各节点/各 rank 上执行，并在额外的一次通信中只汇聚轻量级的形状与序列长度元数据。

求解器运行

为避免阻塞主要训练过程，求解器在 data_sampler 中异步运行，以便与训练迭代重叠执行。为控制开销，我们将穷举搜索替换为小型网格搜索，所有 DP rank 被约束使用相同的 micro-batch 数量，并将该数量从 PP*1 扫描到 PP 的一个小倍数范围内。

在 global batch size 固定的情况下，这个一维网格能够刻画“每个 micro-batch 的工作量”与“pipeline bubble”之间的权衡。图 7 显示：随着 micro-batch 数量增加，micro batch 的工作负载的变化速度会越来越慢，所以针对较大的 micro batch 数量，求解出来的效果和 micro batch 较小时基本一致，系统选择曲线上的“拐点（knee）”，并将搜索限制在其邻域，以保证求解器开销在可接受范围内。

图 7. Microbatch FLOPS 预算曲线

此图显示了 micro batch 数量与每个 micro batch 的工作负载之间的关系。在本例中，PP 大小为 4。

基准测试结果

引入所有改进后，可变长度序列分布引起的不平衡空泡可以显著减少。

在表 1 中，我们将仅做序列打包（THD）而不引入额外优化的版本作为基准，与Dynamic CP 进行对比：模型为llama-13B、global batch size 2048、PP=8、CP=8、并启用 full recompute。共进行 10 次迭代，丢弃第一次作为 warm-up，对剩余 9 次的迭代时间取平均。Dynamic CP 在 GitHub 与 CommonCrawl 数据集上分别获得 1.48 倍与 1.25 倍的加速。

在有数千张 GPU 的工业级集群中，Dynamic CP 带来了超过35%的端到端性能提升。

模型大小	数据集类型	方法	TFLOPS/GPU
Llama 13B	GitHub	Only Packing	195.88
Llama 13B	GitHub	Dynamic CP	289.32
Llama 13B	Commoncrawl	Only Packing	139.17
Llama 13B	Commoncrawl	Dynamic CP	174.39

表 1. 是否开启动态 CP 在不同数据集上的比较

了解详情

本文展示了：在 Megatron-Core 上，Dynamic CP 相比固定 CP 方法，能够提升变长序列训练的吞吐。结合序列 packing、4D 并行以及针对 GPU 优化的 kernel，Dynamic CP 能在不同模型规模下保证高训练效率。

开始使用：

访问 Megatron Core GitHub，使用 Megatron-Core 的优化来训练变长序列模型。
在GitHub查看调度程序和例子。

感谢每一位同事为此项目做出的贡献。

借助动态上下文并行和 NVIDIA Megatron Core 加速可变长度训练