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

本文将介绍动态上下文并行（Dynamic Context Parallelism，Dynamic-CP），这是 NVIDIA Megatron Core 中用于大语言模型后训练或 DiT 预训练的一种调度方法。该方法通过动态选择每个微批量的上下文并行（CP）大小，高效处理可变长度序列，在真实数据集上实现了高达 1.48 倍的加速效果。

在大规模模型训练中，现实世界数据集中序列长度的可变性带来了一个常被忽视的瓶颈。无论是大语言模型训练还是大规模视频生成，其序列长度均呈现出显著的长尾分布特征。少数极长的样本在计算负载与内存消耗中占据了不成比例的较大比重。

在 LLM 训练中，这会导致文本序列的长度在不同批量之间差异显著。在视频生成中，高分辨率、多秒的视频可能涵盖数万个 tokens。这种情况使得数据并行的 rank、模式以及微批量中的样本级 FLOP 和内存使用量出现不均衡，从而阻碍了高效的调度与资源利用。

为了管理可变长度的输入，训练系统通常采用样本级打包策略，将多个较短的序列组合成一个微批量，使其总 token 数不超过目标序列长度。如图 1 所示，这些序列被打包为等长的格式。

Diagram that explains unpacked samples and packed samples. — *图 1. 与打包序列相比，未打包序列*

由于点积注意力具有二次复杂度，即使三个已封装样本的长度相同，其计算工作负载仍可能不同，如图 2 所示。在打包样本中，这种计算工作负载的差异被称为数据并行（DP）计算不平衡。该不平衡会导致 GPU 出现空闲，因为部分 DP 秩需等待计算负载较高的其他秩完成梯度同步。同时，它也会加剧管道并行（PP）中的气泡问题。

The attention compute workload per sample in packed sequences varies, leading to GPU idling. — *图 2. 注意封装序列之间的计算不平衡问题*

在图 3 中，NVIDIA Nsight Systems 配置文件展示了 VLM 训练过程中的负载不均衡现象。由于不同的图像/视频样本具有可变的序列长度并采用了序列打包，导致不同 DP 组之间出现同步开销，这一情况在性能捕获中得以体现。

This diagram shows the Nsight System profiling timeline capture, with sync overhead caused by an imbalance in variable sequence lengths. — *图 3. 同步不同 DP 组的开销*

此外，使用上下文并行时，CP 分片大小由批量中最长的序列决定，以避免 GPU 之间发生内存不足错误。因此，对于无需上下文并行的较短序列也会进行分片处理。即使这些序列能够容纳在单个 GPU 中，但由于同一批次中存在较长序列，仍会被划分，从而导致不必要的 CP 通信开销。

通常，计算会掩盖 CP 通信的开销。然而，当 CP 较大时（尤其是在跨 InfiniBand (IB) 域通信的情况下），若封装序列较短且计算负载较轻，通信开销便会显现，从而导致 CP 计算效率低下。

以下示例展示了因总序列长度较大而需采用 TP2CP8 的情况。许多打包后的序列由多个较小子序列组成，且计算量不足以掩盖 CP 通信开销。

Diagram shows the capture of the Nsight System profiling timeline with large CP for all sequences. For some samples, the NCCL related kernel is longer than the cuDNN related compute kernel. CP communications cannot be hidden by the computation process. — *图 4. 计算能力不足，无法在封装下隐藏 CP 通信*

这些观察表明，需要一种动态的上下文并行方法。这种方法并非静态地将 CP 大小固定为微批量中长度最大的序列，而是采用针对每个微批量的打包策略来调整 CP 大小。诸如 ByteScale 和 WLB-LLM 等相关工作也探讨了类似问题。

切换 CP 大小需要重新划分序列切片，并重新构建用于注意力运算的 CP 通信组。与其他动态并行方案（例如根据序列长度调整张量并行或流水线并行的规模）相比，Dynamic-CP 带来的额外开销极小，因为调整 TP 或 PP 规模通常涉及权重的重新分配或流水线执行图的重构，这些操作的代价较高。

该求解器旨在针对一组可变长度序列，确定如何进行序列打包并选择合适的 CP 大小，以在不超出 GPU 显存限制的前提下，尽可能提升计算效率。求解器能够接收可变长度序列，并计算出最优的填充方式与 CP 大小。该方法在严格满足 GPU 显存约束的同时，有效提升计算资源的利用率。通过对计算与通信开销进行建模，求解器可避免因过度分片导致的短序列问题以及冗余的 CP 通信，从而缓解数据并行中的负载失衡与 CP 通信效率低下的问题。

以下示例展示了使用 Dynamic-CP 的优势。在应用工作负载均衡之前，由于负载不均，不同微批量之间会出现工作流气泡，进而导致 DP 秩间的 DP 不平衡。经过负载均衡后，微批量与 DP 秩之间的气泡得以减少。

This diagram shows the advantages of the Dynamic-CP method. Previously, imbalanced bubbles occurred between data-parallel (DP) ranks and micro-batches, and were further amplified throughout the pipeline. The Dynamic-CP method effectively reduces these bubbles. — *图 5. Dynamic-CP 方法的优势*

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

本节将介绍把 Dynamic CP 集成到 Megatron Core 的具体流程。

A figure showing how Dynamic-CP integrates into Megatron-Core. — *图 6. 将 Dynamic-CP 集成至 Megatron Core 中*

为每个秩构建多个上下文并行组

在标准上下文并行中，每个秩都属于一个在初始化期间静态确定的固定组（cp_group），该组具有固定的 cp_size。然而，动态上下文并行在不同迭代和微批量中具有变化的 cp_size。

为此，单个秩需参与多个不同规模的 CP 组。初始化时构建多个 CP 组，其中 cp_size 的范围从 1 到 dp × cp，且受限于 2 的幂次。该设计使得运行时可根据打包和调度结果灵活选择合适的 CP 组，同时避免产生动态创建通信组的开销。

动态重新调度和打包数据

与通常采用 Batch × Sequence × Head × Dim (BSHD) 布局的预训练不同，Dynamic-CP 在 THD 布局上运行。在此格式中，可变长度序列在长度限制下被打包在一起，原始的 BS 维度被折叠为 token T 维度。

因此，微批量的数量不再固定。在 BSHD 中，微批量的数量由 num_micro_batches = global_batch_size/dp_size/micro_batch_size 给出。

使用 THD 打包时，每个打包序列中所包含的原始序列数量并不固定，这可能导致 num_micro_batches 在迭代过程中发生变化。

Megatron Core 提供多种训练调度程序，具体取决于是否启用工作流并行 (PP/VPP)。为尽可能减少对现有调度逻辑的侵入性更改，我们在原始 data_iterator 周围引入了轻量级的 data_iterator_wrapper。它执行三个步骤：

全局批量中的重新调度与打包序列在 DP rank 中构建均衡的工作负载。
依据打包结果选取合适的 cp_size，以充分减少 CP 通信的低效。
返回当前迭代的有效 num_micro_batches。

通过这种方法，只需插入一个封装器，即可将 Dynamic-CP 支持添加到所有调度程序中，从而基本保持原始调度代码不变。

跨工作流阶段广播并扩展 packedSeqParams

由于 num_micro_batches 各不相同，仅 TP rank 0 参与，并且第一个和最后一个 PP 阶段负责 Megatron Core 中的调度，因此框架会将 num_micro_batches、max_seqlen 和 cu_seqlens 广播至所有相关的 PP rank。这确保了在动态微批量调度下，整个工作流阶段能够一致执行。

借助 Dynamic-CP，有效的 cp_size 在迭代过程中可能发生变化，因此依赖全局静态的 CP 设置存在风险。为解决此问题，PackedSeqParams 增加了对 cp_size 和 cp_group 的支持。

现在，所有依赖上下文并行的组件（例如位置嵌入和 Transformer 引擎注意力）均可从 PackedSeqParams 中获取 CP 配置，以替代原有的全局 CP 变量。这确保了所有与 CP 相关的计算均与动态选择的 CP 布局保持一致。

损失计算和 FLOPS 计算

给定可变长度序列和 THD 布局，不同序列会产生不同数量的有效 tokens。因此，采用按 token 计算的损失函数：loss = loss_over_valid_tokens / total_number_valid_tokens，可有效避免因填充 tokens 而引入的偏差。

早期版本的 Megatron Core 未考虑 THD 布局，在计算 FLOPS 时假设 max_seqlen 为有效序列长度，导致在可变长度场景中出现系统性高估。

数据调度程序建模

Transformer 工作负载随序列长度 S $\mathrm{O}(S^2)$ 呈二次增长。与此同时，激活内存以线性方式增加 $\mathcal{O}(S)$ ，这意味着即使是微小的差异，也会导致 DP rank 与微批量之间的计算与内存出现严重失衡。为了平衡大样本的工作负载，我们可能将小样本打包处理，但这会带来显著的内存压力。由于难以同时均衡 FLOPS 与内存，这一挑战推动了对调度与打包策略的优化需求。

他们的目标是实现理想的均衡分布，使工作负载和内存均匀分布在 DP rank 与微批量之间。由于每个 DP rank 的微批量数量固定，因此为每个微批量设定了相应的工作负载和内存配额。随后，三级调度程序在工作负载与内存目标之间交替调度，根据需要增加较重样本的 CP 大小，以实现计算与内存的平衡。

成本模型、求解器和模拟器的协作

完整的调度程序工作流程由三个部分组成：

成本模型根据每个样本的序列长度估算其执行时间，并对 Transformer 操作中每个样本的工作负载进行建模。这定义了基本负载单元，其准确性将影响最终的性能提升。
求解器以成本模型的输出作为输入，采用启发式算法为每个样本确定接近最优的打包策略。随后，将打包后的样本分组为微批量，并分配上下文并行的 CP 大小。每个 DP 秩的微批量数量会影响结果、工作流气泡、工作流并行失衡气泡以及数据并行失衡气泡。通过迭代不同 DP 秩的微批量数量，可获得较优结果。
模拟器依据分布式工作流并行调度对这些微批量进行评估。它选择执行时间较短（即负载更均衡）的调度方案，同时满足峰值内存限制。

建模过程和双目标平衡

理想的均衡分布是在不同的 DP rank 和不同的微批量之间均匀分配工作负载与内存。假设各 DP rank 中的微批量数量相同，则需确定每个微批量的目标工作负载与内存配额。此外，工作流气泡的分布亦存在差异，应将其均匀分配至各个微批量，以实现端到端的平衡。

均衡 DP rank 中的端到端训练时间表明：

$W_1 \cdot (m_1 V + p - 1) = W_2 \cdot (m_2 V + p - 1)$

$m_i$ 是 DP 等级的微批次数量
$W_i$ 是第 $i$ -th DP 等级的每个微批次的配额
$V$ 是虚拟管道阶段的数量
$p$ 是管道阶段数

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

$W_2 = W_1 \cdot \frac{m_1 V + p - 1}{m_2 V + p - 1}$

同时，全局批量样本的总工作负载可表示为：

$\sum_{i=1}^{DP} W_i \cdot m_i = \mathrm{sum}_{\mathrm{batch}}(\mathrm{workload})$

结合 2 和 3，可以确定各个 DP 中每个微批量的工作负载。

由于计算工作负载随序列长度呈 $\mathrm{O}(S^2)$ 级增长，而内存消耗呈 $\mathcal{O}(S)$ 级增长，因此难以同时实现工作负载与内存的平衡。相反，求解器会在不同阶段交替优化面向工作负载和面向内存的目标，逐步逼近均衡的解决方案。

如果工作负载超过微批量工作负载配额，系统会为其分配更大的 CP 大小。完成此步骤后，工作负载失衡情况得到缓解，内存成为主要限制因素。随后，目标将转向内存，选择计算量相对较少的样本来填充存储桶。其余样本按降序排列，并采用相同的启发式算法分配至各个微批量。

零用度执行

在运行时，调度工作流不应给训练循环带来显著的开销。在实践中，系统需要克服 I/O 压力和求解器运行时问题。

I/ O 压力

首先，构建调度计划需要对全局批量进行额外的 get_item 传递，以收集序列长度和形状信息。通过在集群中分配探测 get_item，两种互补技术可缓解 I/O 压力，并借助额外的通信步骤，仅收集轻量级的形状与序列长度元数据。

求解器运行时

为避免阻塞主要训练过程，求解器在 data_sampler 中异步运行，从而与训练迭代实现重叠。为便于管理空间，我们用小型网格搜索替代了详尽搜索。所有 DP 秩均受限于使用相同数量的微批量，且该数量将从 PP = 1 扫描至 PP 的较小倍数。

在固定的全局批量大小下，此一维网格可在微批量工作负载与工作流气泡之间进行权衡。图 7 显示，随着微批量数量的增加，工作负载差异迅速减小。选取该曲线上“膝盖”位置的点，并将搜索范围限制在其邻近区域，以确保求解器开销处于可接受水平。

Diagram showing the relationship between the number of micro-batches and the workload per micro-batch. In this example, the PP size is 4. — *图 7. Microbatch FLOPS 预算曲线*

基准测试结果

引入所有增强功能后，由可变长度序列分布导致的不平衡气泡可大幅减少。

在表 1 中，基于以下设置对动态 CP 进行纯包装基准评估：llama-13B、全局批量大小为 2048、PP = 8、CP = 8，并采用完整重新计算。运行 10 次迭代，舍弃第一次迭代作为预热阶段，对剩余 9 次迭代的运行时间取平均值。实验结果表明，动态 CP 在 GitHub 和 CommonCrawl 数据集上分别实现了 1.48 倍和 1.25 倍的加速效果。

在拥有数千个 GPU 的工业环境中，Dynamic CP 方法可将端到端性能提升超过 35%。

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

表 1. 动态 CP 与纯打包方法在不同数据集上的比较

了解详情

本文展示了采用 Megatron 核心后端的 Dynamic CP 相较于固定 CP 方法，在可变长度序列的训练中提升了吞吐量。通过序列打包、4D 并行以及 GPU 优化内核，Dynamic CP 能够在不同模型规模下保持高效的训练性能。

开始使用：

Megatron Core GitHub，开始使用 Megatron Core 进行优化，支持可变长度序列训练模型。
调度程序也已在 GitHub 上发布。

感谢每一位同事为该项目付出的努力。

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