Hierarchical Balance Packing: Towards Efficient Supervised Fine-tuning for Long-Context LLM

• 作者/机构: Yongqiang Yao (Shanghai Jiao Tong University), Jinru Tan (Central South University), Kaihuan Liang (SenseTime Research), Feizhao Zhang (SenseTime Research), Jiahao Hu (SenseTime Research), Shuo Wu (SenseTime Research), Yazhe Niu (The Chinese University of Hong Kong), Ruihao Gong (Beihang University, SenseTime Research), Dahua Lin (SenseTime Research), Ningyi Xu (Shanghai Jiao Tong University)

A1 主要贡献

核心问题: 训练长上下文大语言模型（Long-Context LLMs）具有挑战性，因为使用长短文混合数据进行训练常会导致工作负载不平衡。现有工作主要通过数据打包（data packing）来缓解此问题，但未能考虑注意力计算不平衡和通信开销浪费的问题。

研究目标与创新点: 本文提出了一种名为分层平衡打包（Hierarchical Balance Packing, HBP）的方法，旨在通过新颖的批次构建方法和训练策略来解决这些低效问题。HBP的核心贡献如下：
1. 多级数据打包组: HBP构建了多级数据打包组，每个组都针对一个特定的打包长度进行优化。它将训练样本分配到最优的分组中，并为每个分组配置最有效的设置，包括序列并行度（sequential parallelism degree）和梯度检查点（gradient checkpointing）配置。
2. 动态训练流程: 为了有效利用多级数据组，设计了一个专为HBP定制的动态训练流程，包括课程学习（curriculum learning）、自适应序列并行（adaptive sequential parallelism）和稳定的损失函数（stable loss）。
3. 解决现有打包方法的局限性:
* 不平衡的注意力计算: 传统的打包方法（朴素打包）直接混合长短样本，导致注意力计算复杂度的方差很高，从而引起工作负载不平衡。HBP通过分层分组来降低注意力复杂度的方差。
* 浪费的通信开销: 长上下文数据需要序列并行（SP）和集体通信，而短数据不需要。混合打包会迫使短上下文数据也参与序列并行通信，造成时间浪费。HBP将不同长度的数据分开，避免了不必要的通信。

下图展示了朴素打包与分层平衡打包之间的差异。其中，短、中、长代表不同长度的样本，SP Comm 指的是启用序列并行（SP）训练所引入的额外通信开销。ABR（注意力平衡比）衡量不平衡的注意力计算，CR（通信比）衡量额外的通信开销。

主要成果: 实验证明，HBP在多个数据集和开源模型上显著减少了训练时间，同时保持了强大的性能。例如，在最大的DeepSeek-V2 (236B) MoE模型上，HBP将训练速度提升了2.4倍，并取得了有竞争力的性能。

A3 背景知识与问题分析

本节首先在表1中定义了符号，并在3.1节介绍了性能指标。随后在3.2节对常用的打包方法和3.3节的训练策略进行了初步分析。

表1: 符号与定义

表2: 使用打包在不同序列长度下的ABR和CR结果。指标越低表示训练效率越高。

3.1 衡量指标

Dist Balance Ratio (DBR)。该指标由【8, Omnibal: Towards fast instruct-tuning for vision-language models via omniverse computation balance, 2024, arXiv preprint arXiv:2407.20761】提出，用于量化设备间基于长度的计算平衡性。较低的DBR值表示不同设备间的工作负载分布更均衡。

Padding Ratio (PR)。该指标由【8, Omnibal: Towards fast instruct-tuning for vision-language models via omniverse computation balance, 2024, arXiv preprint arXiv:2407.20761】提出，用于衡量设备内基于输入长度的填充所导致的计算浪费比例。较低的PR值表示填充的token数量较少。

Attention Balance Ratio (ABR)。我们提出了这个指标来量化不同设备间注意力计算的不平衡性。以前的指标仅根据完整注意力的输入长度来估算注意力计算成本。然而，在使用打包算法时，注意力计算成为总成本的重要因素。以一个4K序列的打包为例，{1K, 1K, 1K, 1K}和{2K, 2K}的总长度相同，但它们的实际注意力计算量却显著不同，复杂度分别为4K²和8K²。注意力平衡比（ABR）及其示例如公式2所示。较低的ABR表示不同DP组之间的注意力计算更均衡，从而减少了空闲时间。

$$\mathrm{ABR}=\frac{\sum_{i}^{N}\left(A_{\max }-A_{i}\right)}{A_{\max } \times N}, \quad \mathrm{ABR}(\text {example})=\frac{8 K^{2}-4 K^{2}}{8 K^{2} \times 2}=0.25$$

Communication Ratio (CR)。我们提出了这个指标来衡量序列并行（SP）通信的开销。公式3中的 $T_{comm_i}$ 表示当前迭代中需要SP通信的总token数，而 $T_i$ 表示当前迭代中的总token数。较低的CR值表示需要SP通信的token较少，从而减少了通信开销。

Average Tokens (Ave-T)。该指标定义为每次迭代处理的平均token数，用以衡量模型的工作负载。较大的Ave-T值表示更高的训练效率。

$$ \mathrm{CR}=\frac{\sum_{i}^{N} T_{\mathrm{comm}_{i}}}{\sum_{i}^{N} T_{i}}, \quad \text { Ave-T }=\frac{\sum_{j}^{\mathrm{Iter}_{\max }}\left(\sum_{i}^{N} \mathrm{~T}_{i}\right)}{\mathrm{Iter}_{\max } \times N} $$

3.2 打包分析

打包的重要性。在表3中，我们对三种批处理方案进行了比较：随机批处理、ISF打包批处理【8, Omnibal: Towards fast instruct-tuning for vision-language models via omniverse computation balance, 2024, arXiv preprint arXiv:2407.20761】（不同打包方法的比较见附录A）、以及排序批处理【6, LongAlign: A recipe for long context alignment of large language models, 2024, EMNLP】（确保每个批次内的序列长度相近），实验使用了Tulu3【14, Tülu 3: Pushing frontiers in open language model post-training, 2024, https://arxiv.org/abs/2411.15124】数据集，序列长度分别为4K、32K和128K。结果显示，排序和打包都能降低DBR和PR，从而加速训练。对于像128K这样的长序列，打包比排序能获得更高的Ave-T，提高了GPU利用率，并且对混合长度的数据集更有利 。

打包的局限性。尽管打包可以部分解决混合训练带来的效率问题，但仍存在一些局限性。注意力复杂度不平衡问题：打包不同序列长度的样本会导致不同的注意力计算复杂度。直接混合这些样本会造成工作负载不均衡和资源利用效率低下。如表2所示，随着序列长度的增加，ABR显著上升，表明设备空闲时间增多。更多细节见附录D。SP通信开销：长序列需要进行通信以完成注意力计算，而短序列则不需要。直接混合它们会导致额外的通信开销。如表2所示，当序列长度增加时，CR达到1，这表明所有短上下文数据都参与了不必要的通信。

表3: 不同序列长度下批处理策略的比较。在32K和128K设置下，排序批处理的本地批大小B是动态调整的。

表4: 不同序列长度下的SP、最小GC层数和内存成本结果。Iter Time表示十次迭代的平均时间。

3.3 训练策略分析

SP与GC的权衡。在给定的GPU资源和训练数据下，只要满足显存（VRAM）要求，我们可以从多种训练策略中进行选择。主要考虑的因素是序列并行（SP）的程度和梯度检查点（GC）的配置，即启用GC的层数。如果SP程度较小，显存需求会很高，这迫使我们增加GC层数，导致过多的额外计算。反之，如果SP程度较大，虽然显存需求降低了，但会引入额外的通信开销。因此，在SP程度和GC配置之间存在一个权衡。此外，最优策略因序列长度的不同而异。表4显示，32K、64K和128K的最优策略是不同的。

A2 方法细节

本节首先描述如何确定最优的打包长度组（第4.1节），然后解释每个样本如何被分配到一个组（第4.2节）。最后，我们介绍一个用于HBP的动态训练流程（第4.3节）。整体框架如图2所示。

4.1 分层组自动选择

基于性能剖析的自动选择算法。为了确定最优的打包长度组，我们设计了一个基于性能剖析（profile-based）的自动选择算法，如算法1所述。该算法分两个阶段运行：（1）为预定义的一组可能的序列长度（例如，8K, 16K, 32K, 64K, 128K）找到最佳训练策略，这一步基于朴素打包。（2）通过优化通信开销来推导出最终的打包组。

算法1 分层组自动选择

1: 输入: 长度L, 性能剖析时间P, 策略S
2: 阶段1: 寻找最佳训练策略
3: 初始化 P ← [], S ← []
4: for each l ∈ L do
5:   s = (sp, ckpt) ← FindBestSpCkpt(l)
6:   P.add(ProfileTime(s)), S.add(s)
7: end for
8: j ← argmin(P)
9: s_best ← S[j], l_best ← L[j]
10: s_max ← S[−1], l_max ← L[−1]
11: 阶段2: 优化打包组以减少通信
12: l1 ← ⌊l_best/l_best.sp⌋, l2 ← ⌊l_max/l_max.sp⌋
13: if l2 > l_best then
14:   Lp ← [l1, l_best, l2, l_max]
15: else
16:   Lp ← [l1, l_best, l_max]
17: end if
18: return Lp

算法2 平衡打包

1: 输入: 数据集 D = {x1, x2, . . . , xn}, 分层打包组 Lp = {l1, l2, . . . , ln}
2: G = {G1, G2, . . . , Gn}: 已打包的数据组
3: B = {B1, B2, . . . , Bn}: 最终的批处理数据
4: GroupData(D, Lp): 根据打包组Lp将D划分为子集[D1, D2, . . . , Dn]。
5: Packing Gi =(Di, li): 用长度li打包Di
6: GreedyFill Gi = (Gi, li, [Di−1, . . . , D1]): 从较小的组中填充数据以减少PR。
7: Balance Batching(Gi): 根据3.1节中的注意力复杂度A对已打包数据进行排序，构建平衡的批次。
8: 
9: 初始化 B ← [ ]
10: [D1, . . . , Dn] ← GroupData(D, Lp)
11: for i = n to 1 do
12:   Gi ← Packing(Di, li)
13:   Gi ← GreedyFill(Gi, li, [Di−1, . . . , D1])
14:   Bi ← Balance Batching(Gi), B.add(Bi)
15: end for
16: return Shuffle(B)

阶段1：寻找最佳训练策略。对于每个输入长度$l$，我们使用FindBestSpCkpt函数计算最优的SP程度和GC配置$s = (sp, ckpt)$，该函数能在SP程度和GC配置之间实现最佳权衡。关于FindBestSpCkpt的更多细节在附录I.1中展示。具体来说，给定一个打包长度$l$，我们遍历所有可能的sp程度和ckpt GC配置，并对它们的迭代时间进行性能剖析。然后通过贪心法找到最佳组合$(sp, ckpt)$。当所有输入长度都处理完毕后，算法通过找到使性能剖析时间P最小化的索引$j$来选择最佳策略。最优策略和输入长度为：

$$j = \operatorname{argmin}(P), \quad s_{\text{best}} = S[j], \quad l_{\text{best}} = L[j]$$

此外，最大输入长度$l_{max}$及其对应的策略$s_{max}$也会被保留下来。

阶段2：优化打包组以减少通信。在计算出最优的$l_{best}$和相应的$s_{best}$之后，我们可以进一步优化它以减少通信开销，并获得最终的打包组长度。以图3中SP=2的32K打包组为例，每个SP进程处理一个16K token的分割。对于长度超过16K的样本，无论如何都需要额外的通信。对于长度小于16K的样本，我们可以将它们分组到16K的打包组中，这相当于使用SP=1来训练16K的打包组，从而减少了通信。

$$l_{1}=\left\lfloor\frac{l_{\text {best }}}{s_{\text {best.sp }}}\right\rfloor \quad \text { and } \quad l_{2}=\left\lfloor\frac{l_{\max }}{s_{\max . \mathrm{sp}}}\right\rfloor$$

对于$l_{best}$，我们根据其最优的SP程度$s_{best}.sp$推导出最小的打包组$l_1$，以确保在序列并行期间没有通信开销。同样，我们也会考虑$l_{max}$的最小打包组$l_2$。如果$l_2$小于$l_{best}$，它将被合并到$l_{best}$的范围内。最终，得到分层打包组$L_p = {l_1, l_{best}, l_2, l_{max}}$。更多实现细节见附录I。

4.2 平衡打包

样本分组与优化。在获得最优的分层打包组后，我们将数据集样本分配到不同的组中，同时确保第3.1节中概述的指标（DBR, PR, ABR, 和 CR）得到良好优化，而这正是传统打包方法难以解决的问题。我们首先根据$L_p$将整个数据集划分为子数据集$[D_1, D_2, . . . , D_n]$。对于每个$D_i$，执行以下步骤：

1. 打包(Packing)：我们将$D_i$中的数据打包成长度为$l_i$。这确保了较低的PR和DBR。值得注意的是，任何打包方法都是可行的。
2. 贪婪填充(GreedyFill)：我们使用剩余未打包的数据进行贪婪填充，因为较大的组仅使用其自身的样本很难填满打包数据。一个简单的示例说明见附录E。
3. 平衡批处理(Balance Batching)：我们使用启发式解决方案——排序，根据注意力复杂度来构建平衡的批次$B_i$，这有助于在不同打包组之间保持注意力计算的平衡。对于大规模数据集，启发式解决方案是一种高效的近似方法。更多证据见附录D。

平衡打包流程。平衡打包的流程如算法2所示。由于我们的方法天生包含多个层次，即分层打包组，它能自动分离长短上下文数据，从而避免了浪费的通信开销和不平衡的注意力计算，显著降低了CR和ABR。得益于贪婪填充（GreedyFill），我们还在多个层级上实现了较低的PR和DBR。更多实现细节见附录J。

4.3 动态训练流程

多级输入的动态训练设计。由于HBP涉及多级输入，设计一个动态的训练流程至关重要，该流程能够实现不同打包组的热切换。我们提出了自适应SP（预先初始化多个SP，实现零开销）来确保多级打包的高效训练。同时，提出了课程学习和稳定的损失归一化器来提升混合训练的性能。

自适应序列并行。每个打包组都被分配一个最优的训练配置$s = (sp, ckpt)$。我们采用一种交替方案，为每个打包组选择最佳的SP程度和GC设置，如图2所示。

课程学习策略。在长上下文任务上进行训练具有挑战性，因为在没有指令遵循能力的情况下开始训练，可能会导致训练损失出现剧烈波动，如图4所示。得益于我们固有的分层结构，采用课程学习策略变得非常直接：在训练的早期阶段，从通用的短上下文数据开始。随着训练的进行，我们过渡到一种混合方法，交替训练短上下文和长上下文数据，如图2所示。

稳定的损失归一化器。由数据打包引入的训练稳定性是一个重要问题，因为它会影响数据分布。先前的工作【6, LongAlign: A recipe for long context alignment of large language models, 2024, EMNLP】分析了损失计算，确定了两种主要的损失归一化器：$L_{token}$（Token-Mean）和$L_{sample}$（Sample-Mean）：

$$ \mathcal{L}_{\text{token}} = \frac{\sum_{i}^{B_{l}} \text{loss}_{i}}{\sum_{i}^{B_{l}} T_{i}}, \quad \mathcal{L}_{\text{sample}} = \frac{\sum_{i}^{B_{l}} \frac{\text{loss}_{i}}{T_{i}}}{B_{l}}, \quad T_{\text{ave}} = \frac{\sum_{i}^{B_{g}} T_{i}}{B_{g}}, \quad \mathcal{L}_{\text{stable}} = \frac{\sum_{i}^{B_{l}} \text{loss}_{i}}{B_{l} * T_{\text{ave}}} $$

其中，$B_l$是DP组内的本地批大小，$T_i$是批次$i$中的损失token数。他们认为不稳定的原因在于归一化的不一致性。【27, T\" ulu 3: Pushing frontiers in open language model post-training, 2024, arXiv preprint arXiv:2411.15124】也提出了不进行归一化的求和损失（sum loss），以减轻归一化差异的影响。然而，求和损失引入了一个权衡：随着序列长度增加，梯度值会不成比例地急剧上升（达到1e+5），如图5左侧部分所示。为了解决上述问题，我们凭经验观察到，所有DP组的全局批大小$B_g$的平均token数$T_{ave}$可以作为一个稳定的损失归一化器。稳定的损失归一化器源于理论考量，其目标是确保每个token对总损失的贡献相等。这种形式减轻了由异构序列长度、变化的数据并行组大小和梯度累积策略引入的偏差。我们通过一个简单的设置来说明不同归一化策略的效果：一个大小为2的数据并行（DP）组，其中每个rank处理一个包含2个样本的本地批次。

Token-Mean。在token-mean方法中，每个DP rank的损失由其本地token数进行归一化：

$$\mathcal{L}_{\mathrm{dp}_{1}}=\frac{\operatorname{loss}_{1}+\operatorname{loss}_{2}}{T_{1}+T_{2}}, \quad \mathcal{L}_{\mathrm{dp}_{2}}=\frac{\operatorname{loss}_{3}+\operatorname{loss}_{4}}{T_{3}+T_{4}}$$

然后，全局损失在所有rank之间取平均：

$$\mathcal{L}_{\text {final }}=\frac{\mathcal{L}_{\mathrm{dp}_{1}}+\mathcal{L}_{\mathrm{dp}_{2}}}{2}=\frac{\frac{\text { loss }_{1}+\text { loss }_{2}}{T_{1}+T_{2}}+\frac{\text { loss }_{3}+\text { loss }_{4}}{T_{3}+T_{4}}}{2}$$

这种形式是按rank进行损失归一化，没有考虑全局token分布，当不同rank的序列长度不同时，可能会引入偏差。

Sample-Mean。在sample-mean方法中，损失首先按样本进行归一化，然后取平均：

$$ \mathcal{L}_{\mathrm{dp}_{1}}=\frac{\frac{\operatorname{loss}_{1}}{T_{1}}+\frac{\operatorname{loss}_{2}}{T_{2}}}{2}, \quad \mathcal{L}_{\mathrm{dp}_{2}}=\frac{\frac{\operatorname{loss}_{3}}{T_{3}}+\frac{\operatorname{loss}_{4}}{T_{4}}}{2} $$

全局损失变为：

$$\mathcal{L}_{\text{final}} = \frac{\mathcal{L}_{\text{dp}_1} + \mathcal{L}_{\text{dp}_2}}{2} = \frac{\frac{\text{loss}_1}{T_1} + \frac{\text{loss}_2}{T_2} + \frac{\text{loss}_3}{T_3} + \frac{\text{loss}_4}{T_4}}{4}$$

这种策略平衡了单个样本，但未能反映总token数，给予了较短序列过多的权重。

稳定的损失归一化器。为了消除这种偏差，我们通过全局平均token长度进行归一化：

$$\mathcal{L}_{\text{ave}} = \frac{T_{1} + T_{2} + T_{3} + T_{4}}{4}, \quad \mathcal{L}_{\text{dp}_{1}} = \frac{\text{loss}_{1} + \text{loss}_{2}}{2 \cdot T_{\text{ave}}}, \quad \mathcal{L}_{\text{dp}_{2}} = \frac{\text{loss}_{3} + \text{loss}_{4}}{2 \cdot T_{\text{ave}}}$$

每个DP的损失都由$T_{ave}$进行归一化，最终的损失为：

$$\mathcal{L}_{\text{final}} = \frac{\mathcal{L}_{\text{dp}_1} + \mathcal{L}_{\text{dp}_2}}{2} = \frac{\text{loss}_1 + \text{loss}_2 + \text{loss}_3 + \text{loss}_4}{4 \cdot \mathcal{L}_{\text{ave}}} = \frac{\text{loss}_1 + \text{loss}_2 + \text{loss}_3 + \text{loss}_4}{T_1 + T_2 + T_3 + T_4}$$

这种归一化确保了每个token的贡献是相等的，无论它属于哪个样本或rank，从而产生一个无偏的全局损失。

A4 实验环境

• 数据集:

  • Tulu3 (32K) 【14, Tülu 3: Pushing frontiers in open language model post-training, 2024, https://arxiv.org/abs/2411.15124 】: 用于通用任务，增强模型处理0.1K到128K token输入的能力。
  • LongCite (128K) 【15, Longcite: Enabling llms to generate fine-grained citations in long-context qa, 2024, arXiv preprint arXiv:2409.02897】: 用于长上下文任务。
  • 消融实验中，从Longsite数据集中均匀采样了约2k个样本。

• 模型架构:

  • LLaMA 3.1 【1, The llama 3 herd of models, 2024, arXiv preprint arXiv:2407.21783】
  • Qwen-2.5 【2, Qwen2. 5 technical report, 2024, arXiv preprint arXiv:2412.15115】
  • DeepSeek-V2 (236B)

• 硬件配置:

  • 多数模型使用32块H100 80GB GPU进行训练。
  • DeepSeek-V2 (236B) 使用256块H100 80G GPU进行训练。

• 软件配置:

  • 框架: DeepSpeed 【28, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20】, Megatron-LM 【29, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, arXiv preprint arXiv:1909.08053】 (用于DeepSeek-V2)。
  • 序列并行: 采用DeepSpeed-Ulysses 【10, Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models, 2023, arXiv preprint arXiv:2309.14509】的序列并行方法，ring-attention 【11, Blockwise parallel transformer for large context models, 2023, Advances in neural information processing systems】方法也适用。
  • 优化器: AdamW，学习率为1e-5，权重衰减为0.01。
  • 损失归一化: 基线模型默认使用Token-Mean，HBP使用Ave-Token。

• 评估: 使用OpenCompass 【30, Opencompass: A universal evaluation platform for foundation models, 2023】进行全面评估，覆盖通用任务（MMLU, MMLU PRO, CMMLU, BBH, Math等）和长上下文任务（Ruler, NeedleBench, LongBench, Longcite）。

A4 实验结果

5.2 主要结果

与基线方法的比较。在表5中，我们将HBP方法与LongAlign 【6, LongAlign: A recipe for long context alignment of large language models, 2024, EMNLP】和打包方法ISF 【8, Omnibal: Towards fast instruct-tuning for vision-language models via omniverse computation balance, 2024, arXiv preprint arXiv:2407.20761】进行了比较。我们注意到，尽管LongAlign在一定程度上提升了通用任务的平均性能，但它牺牲了在长任务（如Ruler-128K）上的表现。相比之下，我们的方法在通用短任务和长任务上都保持了强大的性能。这种改进在从8B到236B参数的各种模型尺寸上都是一致的。特别值得一提的是，对于最大的MoE模型DeepSeek-V2（236B），我们的方法实现了惊人的2.4倍训练加速，将训练时间从57.1 GPU天减少到23.8 GPU天。完整结果见附录C。

表5: 不同模型的结果。朴素打包基线ISF使用Token-Mean损失归一化器。LongAlign使用其提出的损失重加权方法。AVE代表通用任务的平均性能。由于资源限制，Deepseek-V2(236B)在32K设置下训练。

5.3 消融实验结果

混合训练的重要性。表9显示，短上下文和长上下文数据对于维持模型的通用能力和长文本处理能力都至关重要。仅使用短上下文数据（第1行）会损害长文本性能，而仅使用长上下文数据（第2行）会削弱通用能力。第3行使用部分短数据也证实了这些影响。

表9: 混合训练的重要性。

HBP组件的贡献。在表6中，我们展示了通过分层打包可以显著降低注意力平衡比（ABR）和通信比（CR）。具体来说，ABR从0.506降至0.288，CR从1.0降至0.173。通过将具有相似注意力计算复杂度的批处理数据进行平衡，我们获得了更加平衡的小批量数据，ABR从0.288降至0.002。总体上，我们实现了1.4倍的加速。在更大的LLaMA-3.1-70B模型上也观察到了类似的结果。

表6: HBP组件的结果。此实验使用Token-Mean损失归一化器。Hierarchical表示启用分层打包。Balance指平衡批处理。

课程学习的效果。表7展示了课程学习对HBP的影响。从短任务开始，然后逐渐混合长上下文任务，有助于训练和收敛。我们将类似的课程策略应用于朴素的ISF基线，也显示了性能提升，这证实了课程学习对于长上下文SFT的普遍有效性。值得注意的是，HBP天然地分离了短上下文和长上下文，使得课程学习更容易应用。

表7: 课程学习（CL）的结果。

课程学习的迭代次数。表8展示了不同模型下详细的消融研究。结果显示，即使是简单的课程学习设置，比如在早期阶段增加100步的短样本训练，就已经能带来明显的提升。为确保训练稳定，我们最终在实验中采用了500步纯短样本训练的设置。

表8: 不同CL迭代次数的结果。

稳定的损失归一化器。我们比较了几种损失归一化方法，如表10所示，Ave-Token损失归一化器在通用任务（AVE 57.6）和长上下文任务（LongBench 43.1, Ruler-128K 70.8, LongSite 71.2）上均取得了最高性能。

表10: 不同损失归一化器的结果。

分层组自动选择的重要性。表11比较了手动（第2行）和自动（第3行）选择打包组的效果。基于附录F的自动选择方法实现了更大的加速，将训练时间从4.25 GPU天减少到3.73 GPU天。

表11: 优化打包组以减少通信的消融结果。Groups表示训练中使用的打包组长度，SP表示序列并行度。

与FlexSP的比较。如表12所示，HBP比FlexSP实现了更低的ABR、DBR和PR，同时保持了相似的CR，这有助于更快的训练（从4.35 GPU天降至3.73 GPU天），即使只考虑模型计算时间。此外，FlexSP每个批次会引入5-15秒的分组开销，这进一步减慢了训练速度（见附录B）。HBP提前进行全局数据级操作，数据开销可忽略不计。

表12: 与FlexSP的消融结果。所有指标均基于官方代码计算。对于FlexSP，我们只考虑模型计算时间，忽略了数据分组时间的影响。

可忽略的开销。HBP的开销是可忽略的：自动选择约需3分钟，打包仅需5秒，两者总共占训练时间的不到2%（见附录G）。

数据集泛化性。为了验证我们方法的泛化性，我们在不同的数据集（OpenHermes【16, Openhermes 2.5: An open dataset of synthetic data for generalist llm assistants, 2023, https://huggingface.co/datasets/teknium/OpenHermes-2.5 】, LongWriter【23, Longwriter: Unleashing 10,000+ word generation from long context llms, 2024, arXiv preprint arXiv:2408.07055】）上进行了实验。我们的方法在这些数据集上也取得了有效的结果。结果见附录H。

A5 结论

本文提出了一种名为分层平衡打包（Hierarchical Balance Packing, HBP）的新策略，通过多级数据打包和动态训练流程来解决长上下文LLM训练中的工作负载不平衡问题。由于计算资源的限制，我们尚未在更长的上下文（如256K或512K）上进行实验。此外，我们也没有在其他后训练任务（如RLHF或DPO）上验证HBP。我们将这些探索留给未来的工作。

A6 附录

A 打包策略结果

不同打包策略的比较。表13展示了不同打包策略的复杂性及其对应的最终训练时间。它还表明，训练指标数据平衡比（DBR）、填充率（PR）和注意力平衡比（ABR）与训练时间密切相关，强调了这些指标的有效性。基于打包策略的计算复杂性、训练时间以及准确性表现，我们最终选择ISF作为我们的朴素打包基线。

表13: 在128K混合数据训练设置下不同打包策略的结果。N: 样本数; M: 每包样本数; S: 最大包长度; C: 迭代次数。

B FlexSP数据分组结果

FlexSP数据分组开销分析。实验在序列长度为128K、使用32个GPU的条件下进行。FlexSP的结果是通过直接测试其官方发布的代​​码获得的。FlexSP的数据分组方法每批次会产生5-15秒的不可忽略的开销，这部分开销并不总能被训练时间所掩盖，如表14的第3和第4行所示。

表14: FlexSP数据分组的结果。Data Time指与模型训练重叠后获取一个批次数据的时间。Model Time指模型训练所花费的时间。Total Time表示包括数据和模型处理在内的总时间。所有时间均为100次迭代的平均值。

C 完整结果

通用任务与长上下文任务的完整结果。表15和表16分别展示了我们在通用任务和长上下文任务上的完整结果。这些结果共同验证了我们HBP方法的有效性。其他密集模型都使用GQA架构，而Deepseek-V2采用MLA结构，由于K和V头的数量更多，这涉及到更多的通信，类似于MHA。HBP有效地减少了通信开销，使得加速效果更为显著。

D 注意力不平衡问题的更多细节

注意力复杂度的不平衡性。根据Megatron-LM论文【29, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, arXiv preprint arXiv:1909.08053】的推导，给定一个包含n个子序列的打包序列，每个transformer层的计算成本可以近似为：$24Bsh^2 + 4B \sum s_i^2$，其中$s_i$是第i个子序列的长度，B是批大小，h是隐藏层大小。因此，注意力复杂度的方差$\sum s_i^2$直接影响计算成本。

平衡批处理的有效性。排序确实是一个启发式解决方案，并且无法从数学上严格证明它能保证完美的平衡。然而，根据在不同数据集上的结果，它实现了接近完美的平衡（低ABR值接近0）。启发式解决方案是解决大规模数据集问题的常用且高效的近似方法。我们的方法确保在超过99.6%的训练迭代中，在不同数据集上都能保持注意力平衡，如表17所示，这证明了我们方法的强大泛化能力。

平衡证据。图6的左侧显示了基线打包在DP组之间存在较高的注意力成本方差，而右侧显示了我们基于排序的方法，该方法显著降低了这种方差并提高了效率。

表15: 通用任务的完整结果。朴素打包基线ISF使用Token-Mean损失归一化器。LongAlign使用其提出的损失重加权方法。AVE代表通用任务的平均性能。由于资源限制，Deepseek-V2(236B)在32K训练设置下进行训练。

表16: 长任务的完整结果。朴素打包基线ISF使用Token-Mean损失归一化器。LongAlign使用其提出的损失重加权方法。AVE代表通用任务的平均性能。由于资源限制，Deepseek-V2(236B)在32K训练设置下进行训练。

表17: 使用HBP在不同数据集上的ABR比较。

E GreedyFill的更多细节

GreedyFill的目标与工作方式。GreedyFill旨在减少包内的填充token并提高Ave-T。它在本地打包组内操作，通过从其他较短的打包组中选择少量序列来有效地填充剩余空间，如图7所示。

F 分层组自动选择的更多细节

自动选择的必要性。表18展示了一个32K token长度训练的例子，其中包含了各种SP程度和GC配置。第二行突出了满足内存限制的最小配置s = (2, 28)。虽然此配置符合内存要求，但未能达到最佳性能。相比之下，第四行展示了一个更平衡且有效的配置s = (8, 8)，实现了最快的速度。这些实验表明，对于给定的打包长度，不同的SP程度和GC配置之间存在显著的性能差异。

最终选择的组结果。表19展示了不同序列长度L下，最优SP程度和GC配置的训练速度。表的第二行显示，我们的最优配置是groups = 16K，对应的训练策略是s = (1, 28)。这表明，对于不同的打包组，最优训练策略是不同的。以上证据强调了搜索最佳分组及其对应训练策略的重要性，即自动分组选择（Auto-Group Selection）。

表18: 不同SP程度和GC配置的结果。Iter Time是十次迭代的平均时间。

表19: 使用混合数据在不同打包组下的结果。Iter Time是十次迭代的平均时间。

G 开销分析

平衡打包和自动选择的开销。相对于训练时间，平衡打包（Balance Packing）和自动选择（Auto-Selection）的开销。这两个步骤加起来的开销始终可以忽略不计：自动选择最多需要15分钟，打包仅需5秒，而训练时间从168分钟到1400分钟不等。总体而言，合并后的开销低于训练时间的2%，证实了预处理成本与模型训练相比是微不足道的。

表20: 平衡打包和自动选择相对于训练时间的开销。

G.1 性能剖析的开销

内存性能剖析。内存性能剖析的成本取决于设备处理的序列长度集（Seq）。

时间性能剖析。搜索策略（S）的总数由打包长度集（L）和序列并行（SP）设置决定。profile_iter可以根据实际迭代时间进行调整，通常取值范围为3到10。

总成本（示例）。例如，当profile_iter = 5，L = {16K, 32K, 128K}，SP = {2, 4, 8}，Seq = {4K, 8K, 16K}时，对于LLaMA3.1-8B模型：

$$\begin{aligned} \begin{aligned} \text{Cost}_t &= 3 \times 3 \times 5 = 45 \times \text{iteration\_time}, \\ \text{Cost}_m &= 3 \times 5 = 15 \times \text{iteration\_time}, \\ \text{Cost}_{\text{total}} &= \text{Cost}_t + \text{Cost}_m = 60 \times \text{iteration\_time}. \end{aligned} \end{aligned}$$

考虑到总共有3000多次训练迭代，60次迭代（2%）的开销与整体训练时间相比是微不足道的。

G.2 平衡打包的开销

平衡打包的实际开销。对于我们包含100万样本的训练数据集，平衡打包引入的开销始终在5秒以内。

复杂度。平衡打包的计算复杂度可以表示为：

$$\text{Cost} = L \cdot (C \cdot O(N + M) + M \cdot \log M),$$

其中N是样本数量，M是打包组的数量，C是ISF迭代次数，L表示选定的打包集。由于$C \ll N$, $M \ll N$ 且 $L \ll N$，整体复杂度相对于N降至亚线性级别。因此，扩展到更大数据集时仅会产生极小的额外开销。

G.3 所提方法的可扩展性

大规模数据集。我们的性能剖析仅在少量迭代（< 100）中进行。随着数据集规模的增加，总开销比例（< 1%）会变得更低。同时，平衡打包的复杂度为O(N)，因此随着数据集的扩展，引入的额外成本可以忽略不计。

H 数据集泛化性

H.1 OpenHermes

在OpenHermes上的实验结果。我们还提供了在OpenHermes上的一些实验结果。例如，表21显示，在不同打包策略下，结果与Tulu3一致。表22显示，我们的HBP在128K训练设置下同样有效。

表21: 在OpenHermes数据集上4K训练设置下不同打包策略的结果。

表22: 在OpenHermes + Longcite数据集上HBP组件的结果。此实验使用Token-Mean损失归一化器。Balance指启用注意力平衡排序，Hierarchical指激活分层打包。

H.2 LongWriter

在Tulu3和LongWriter上的混合训练结果。图8和表23展示了我们的HBP在Tulu3和LongWriter上进行混合训练的结果。x轴表示用户指令要求的长度，y轴表示根据Long-Writer论文模型输出的长度。更强的线性相关性表明模型具有更好的指令遵循能力。图的左侧显示了未使用Long-Writer数据集的原始基线结果。中间显示了朴素打包的结果，右侧显示了HBP的结果。实验表明，我们的方法同样有效，在两个模型上都实现了一致的加速。

表23: Tulu3和LongWriter混合训练的模型评估结果。

I 分层组自动选择的实现细节

I.1 FindBestSpCkpt

算法3 FindBestSpCkpt函数

1: 初始化 P ← [], O ← []
2: for each sp ∈ SP do
3:   ckpt ← GreedyProfileCkpt(l)
4:   P.add(ProfileTime(l, sp, ckpt)), O.add(sp, ckpt)
5: end for
6: j ← argmin(P)
7: return O[j]

算法4 GreedyProfileCkpt

1: 输入: l, sp, c_min, c_max
2: s1 ← (sp, l, c_min)
3: s2 ← (sp, l, c_max)
4: mr1 ← ProfileMemory(s1)
5: mr2 ← ProfileMemory(s2)
6: m_ave ← (mr2 - mr1) / (c_max - c_min)
7: c_o ← c_max - mr2 / m_ave
8: return c_o

FindBestSpCkpt函数。算法3通过评估所有可能的sp策略来确定最优的梯度检查点策略。
- 初始化：首先用空列表P和O分别存储性能剖析时间和配置。
- 遍历策略：对于每个策略sp ∈ SP：
- 使用GreedyProfileCkpt函数计算最佳的梯度检查点配置（ckpt）。
- 使用ProfileTime对给定配置的模型执行时间进行性能剖析，并将其追加到P中。

• 寻找最优策略：使用argmin(P)找到P中最小性能剖析时间的索引j。
• 返回最佳配置：返回SP程度和对应的梯度检查点配置O[j]。

I.2 GreedyProfileCkpt

GreedyProfileCkpt算法。算法4估算了给定策略l和sp所需的梯度检查点层数。这里我们提供更详细的解释。我们使用"pynvml"库来监控GPU内存使用情况以进行内存分析。

内存性能剖析:
- 步骤1：给定输入长度L，为所有层（例如7B模型为32层）启用梯度检查点，并记录剩余的GPU内存为m1。
- 步骤2：给定相同的输入长度L，仅为一部分层（例如20层）启用梯度检查点，并记录剩余的GPU内存为m2。
- 步骤3：根据以上信息，我们可以估算出每层使用梯度检查点所节省的平均内存：ave_m = (m1 - m2) / (32 - 20)。
通过仅运行10次迭代，我们就可以获得使用梯度检查点时每层节省内存的可靠估计。

内存与时间消耗报告:
对于给定的输入长度L，我们设置不同的序列并行（SP）配置。根据剩余的GPU内存和先前估算的ave_m，我们确定应用梯度检查点的层数。然后，我们运行10次迭代来记录相应的内存使用情况和迭代时间。

完整工作流程:
- 初始化：使用最小和最大梯度检查点层数（基于经验观察）获取配置：

$$s_{1}=(sp, l, c_{min}), \quad s_{2}=(sp, l, c_{max})$$

• 内存性能剖析：对s1（mr1）和s2（mr2）的剩余内存进行性能剖析。
• 内存斜率计算：计算平均内存斜率（mave）为：

$$m_{ave}=\frac{m_{2}^{r}-m_{1}^{r}}{c_{max}-c_{min}}.$$

• 检查点层数估算：估算所需的检查点数（co）为：

$$c_o=c_{max}-\frac{m_2^r}{m_{ave}}.$$

J 平衡打包的实现细节

J.1 GroupData

GroupData函数。给定一个数据集D和预定义的分层长度Lp，我们评估每个数据集x的长度，以确定其所属的区间$(l_{i-1}, l_i]$，并将其分配给相应的组Di。详细过程在算法5中概述。

J.2 GreedyFill

GreedyFill函数。对于给定的打包组Gi及其对应的长度li，我们遍历较小的数据集分区$(D_{i-1}, D_{i-2}, . . . , D_1)$，并贪婪地填充当前Gi内的组g。详细过程如算法6所示。

J.3 注意力平衡排序函数

注意力平衡排序函数。首先，我们计算给定打包组G内所有数据的注意力复杂度。然后，我们根据它们的注意力复杂度对Gi中的元素进行排序，并根据全局token数的要求构建小批量数据，如算法7所示。

算法5 GroupData函数

1: 输入: 数据集 D = {x1, x2, . . . , xN }, 分层打包长度 L = {l1, l2, . . . , ln}
2: 初始化: (D1, D2, . . . , Dn) ← ([], [], . . . , [])
3: for x ∈ D do
4:   if len(x) ∈ (li−1, li] then
5:     Di.add(x)
6:   end if
7: end for
8: return (D1, D2, . . . , Dn)

算法6 GreedyFill函数

1: for g ∈ Gi do
2:   for j = i - 1 → 1 do
3:     for x ∈ Dj do
4:       if Σs∈g len(s) + len(x) ≤ li then
5:         g.add(x)
6:         从Dj中移除x
7:       end if
8:     end for
9:   end for
10: end for
11: return Gi

算法7 注意力平衡排序函数

1: 初始化 A ← []
2: for g ∈ G do
3:   a = Σx∈g(len(x))^2
4:   A.add(a)
5: end for
6: 根据A对G进行排序
7: return G