Sample Packing 综述：LLM 效果与效率的 Tradeoff

本文中我们通过几篇论文来具体介绍 Sample Packing 相关的方案和对应的各种问题，比如 GraphCore 的PackedBert、Meta 的 In-Context-Pretraining、智谱 AI 的 LongAlign、Amazon 的 Fewer Truncations 以及 IBM 的 Packing with FlashAttention。

一、背景

上一篇文章（​​​Sample Packing：长序列 LLM 训练的 Attention 问题及优化​​）中我们简单介绍了 Sample Packing 相关的问题和部分简单实验。本文中我们通过几篇论文来具体介绍 Sample Packing 相关的方案和对应的各种问题，比如 GraphCore 的PackedBert、Meta 的 In-Context-Pretraining、智谱 AI 的 LongAlign、Amazon 的 Fewer Truncations 以及 IBM 的 Packing with FlashAttention。

二、方法

Sample Packing 可以看成是一个经典的 Bin Packing 组合优化（Combinatorial Optimization）问题，核心思想是将一组不同体积的物品放入容量固定的箱子中，目标是最小化所需箱子数量。

Bin Packing 是一个典型的 NP-Hard 问题，因此往往需要关注其效率。对于 N 个 Sample 的数据集，通常需要先排序，对应 O(N*logN) 的时间复杂度；然后 Packing，需要 O(N*logN) 的时间复杂度。不过由于 Sample 的长度通常比较有限，可以采用计数排序的方式进行优化；此外，整个训练集通常有多个数据子集组成，为了保证每个数据子集具有不同的采样权重，往往一个 Sequence 中的 Sample 来自同一个数据子集，一个大的 Bin Packing 问题也变成了多个小的 Bin Packing 问题，可以进一步降低复杂度。本文中我们就不再具体介绍 Bin Packing 的优化问题。

如果从物品顺序角度考虑，Bin Packing 可以分为两类：

• Online Bin Packing：物品按顺序到达，必须立即决定放入哪个箱子，无法预知后续物品的大小。对于 LLM 训练而言，Online 方式依然可以实现 Batch 内（窗口内）的乱序，也可以通过梯度累加增加 Batch 的大小。
• Offline Bin Packing：预先知道所有物品大小，可以全局排序。对于 LLM 训练而言，相当于训练前就预先知道所有序列的长度，并对所有 Sample 打包。

当然，Bin Packing 方案也可能有不同的约束条件：

• 一次性 Packing：每个物品只能放进一个箱子，不允许拆分物品。对于 LLM 而言，可以理解为 Sample 不能截断。
• 多维 Packing：物品和箱子不仅有体积，还有其他属性，比如重量、形状，需要同时满足多维约束。对于 LLM，可以约束同一个 Batch 内不同 Sequence 的计算量尽量类似，以实现更好的负载均衡。

整体来说，对于 LLM 训练可以从数据分布和计算效率的角度考虑相关的方案，数据分布可能影响模型训练的效果，计算效率会影响训练的速度，往往需要综合考量。

影响数据分布的几个因素：

• 是否随机采样，比如排序机制可能引入分布的不一致。
• 是否交叉污染，比如 Packing 是否采用 Document Level 的 Mask（Block Diagonal Mask，也可以通过 Position ID 区分）。
• 是否有 Sample 的截断。

影响计算效率的几个因素：

• 是否 Padding，Padding 是否参与计算。
• Packing 是否采用 Document Level Mask。
• 是否存在负载不均衡问题（主要是指稀疏度不同，计算量不同）。

三、GraphCore PackedBert

3.1 概述

GraphCore 在 2021 年 07 月的 [2107.02027] Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance 中已经讨论了 Sample Packing 导致的交叉污染（Cross Contamination）问题。作者研究了新的 Packing 算法，并通过修改 Attention Mask 和 Position ID 来避免交叉污染，提升 Bert 模型的训练效率。

如下图 Table 1 所示，不同的 Packing 方式有不同的 EFF（有效率），其中 Baseline 的 None 表示有 50% 左右的 Padding Token，而 SPFHP 和 NNLSHP（本文提出） 能获得相对比较高的 EFF。

3.2 结果

如下图 Figure 4 所示可以看出，需要相应修改 Attention Mask 才能保证精度：

如下图 Figure 5 所示，随着 Accelerator 的增加，本文的方案能获得比较稳定的加速，非常接近理论速度（最上的蓝线）；而非 Padding 的方案可能随着 Accelerator 的增加出现明显的降速。

四、Meta In-Context-Pretraining

4.1 概述

Meta 作者在 [2310.10638] In-context Pretraining: Language Modeling Beyond Document Boundaries 也关注了 Sample Packing 的问题。作者指出，之前的预训练流程在训练时将随机的短文档拼接起来形成输入上下文，这些文档之间没有提供预测下一个文档的信号，导致计算效率不高。因此作者提出了新的方案：In-context Pretraining，通过改变文档的顺序，使得每个上下文包含相关的文档，从而明确鼓励模型跨文档边界进行阅读和推理。

4.2 方案

如下图所示，在预训练前会计算文档的相似性，在 Packing 时利用上这种相似性，保证 Sequence 中文档尽可能相关。

4.3 结果

作者使用 CommonCrawl 数据集，预训练了 0.3B 到 7B 参数量的多个模型，并在多种任务上评估，包括上下文学习、阅读理解、对先前上下文的忠实度、长上下文推理和检索增强等。与使用标准方法预训练的模型相比，In-context Pretraining 方法训练出的模型（ICLM）显示出显著的性能提升。

如下图 Figure 3 所示，本文 ICLM 训练的模型获得了更低的困惑度：

如下图 Table 1 所示，基于 ICLM 训练的模型在下游 In-context Learning 任务上也获得了更好的效果：

五、智谱 LongAlign

5.1 概述

智谱 AI 在 [2401.18058] LongAlign: A Recipe for Long Context Alignment of Large Language Models 中讨论了部分 Sample Packing 相关问题。如下图 Figure 3 左图所示，Sequence 的长度各不相同，从 0 - 60K，如果采用 Naive Batching 方式，会导致明显的 Bubble 问题（虽然 NoPadding 技术可以避免重复计算，但是如果采用 Data Parallelism 方式，比较快的设备需要等待比较慢的设备计算完成）。为了解决效率和效果问题，作者提出了 3 种解决方案：Packing、Loss Weighting 和 Sorted Batching。

5.2 Packing

如 Figure 3 右上图所示，就是我们之前介绍的 Sample Packing：将不同的 Sample 拼接在一个 Sequence 里，并且保证尽可能接近 Max Sequence Length，末尾的部分 Token 进行 Padding。然后通过 Block Diagonal Attention Mask 来区别不同的 Sample，以避免 Sample 之间的交叉污染，也就是 Document Level Attention。

PS：作者介绍，这里同样是使用了 FlashAttention2 的 Varlen 特性。

5.3 Loss Weighting

假设训练时的 Batch Size 为 K，总共包含 M 个 Sample，第 i 个 Sample 的 Token 数为 Ni，则对应的 Loss 如下图所示：

然而，增加 Sample Packing 之后会引入一个问题，如下图所示，一个 Sequence 中的不同 Sample 会被看成一个 Sample 来计算损失。当有些 Sample 比较长，其对应的 Token 很多，那么这个 Sample 对 Loss 的贡献就更大，模型可能会在训练时更倾向于优化长 Sample 的表现，进而可能会导致对短 Sample 的学习有所欠缺。

为了解决这个问题，作者提出 Loss Weighting，也就是对不同 Sample 的 Loss 加权。如下图所示，保证其和上述公式（2）等价。作者声称可以在下游任务上带来 10% 左右的效果提升。

PS：不过这里还会引入另外一个问题，因为采用了 Sample Packing，那么实际上不同 Step 中 Sample 的个数在不断变化。比如一个 Batch 里都是短 Sample，那么对应的 M 会比较大；如果一个 Batch 里都是长 Sample，相应的 M 会比较小。这样可能引入两个问题，1. 相当于一个 Batch Size 在不断变化；2. 同样一个 Sample 可能会因为顺序的原因被赋予不同的权重。因此需要尽量保证Batch 中 Sample 的平均个数比较稳定。 

5.4 Sorted Batching

如上图 Figure 3 右下图所示，可以将所有 Sample 进行排序，在组 Batch 时尽量保证一个 Batch 中的 Sample 长度相同（没有 Packing）。这样可以保证不同设备的计算尽可能的均衡。然而，这种方式也不可避免地引入不同 Batch 的数据分布的偏差，有些 Batch 都是长序列，有些 Batch 都是短序列，对于 SGD 优化来说可能并不友好。不过作者发现这种方式可以显著加快训练速度，而不会对效果产生明显的负面影响。这可能是因为使用了大的梯度累加（Micro Batch 中长度类似，但整个 Batch 中包含各种长度的 Sample）。

PS：这种方式也就对应 Transformer 等工作中的 LengthGroupedSampler，如下图所示，排序后可以有效降低无效计算（图片参考 数据分组— XTuner 0.1.23 文档）。

5.5 结果

如下图 Figure 5 所示，作者在 8xA800 GPU 上进行速度对比，可以看成，Packing 和 Sorted Batching 相比 Naive Batching 都有 2x-3x 的加速：

如下图 Table 3 所示，作者基于 ChatGLM3-6B-64K 和 LLaMA-2-7B-64K 进行了相关效果验证。可以看出，Loss Weighting 在 LongBench-Chat 上能带来 5%-10% 的提升，但是在其他任务上并不明显，并且这些方法看着都不是特别鲁棒。

5.6 Packing 负载均衡

智谱 AI 在训练 GLM4（GLM Long: Scaling Pre-trained Model Contexts to Millions | by ChatGLM | Medium） 模型时进一步解决了 Packing 带来的负载均衡问题。如下图所示，虽然都 Packing 到了相同的长度，但是由于其中的 Sample 个数、长度不同，导致其稀疏度差距很大，计算量也相应差距很大。如果它们在不同的设备上执行，同样会存在计算不均导致的 Bubble 问题。

如下表所示，我们在上一篇文章中的实验也能说明这个问题，随着 Sequence 中 Sample 分布的不同，计算的耗时甚至可能差 10x：

如下图所示，作者发现训练中每个 Step 的时间存在较大幅度的波动，这种现象在短文本的 Packing SFT 中并不明显。这是因为短文本时 Attention 的计算占比并不高，而超长文本训练中会尤其明显。

为了解决上述问题，作者进一步提出了 Sorted Packing。具体来说，作者在构建 Batch 数据时考虑了计算复杂度，以确保每个 Batch 中计算复杂度相似，从而减少 Bubble 时间。（PS：这里需要注意，计算复杂度不等于执行速度，如果能针对预估计算速度来打包也许能获得更优的效果）

作者指出也使用了 layer accumulation 技术（PS：上述介绍的梯度累加？）来避免排序导致的偏置问题。

六、Amazon Fewer Truncations

6.1 概述

在 [2404.10830] Fewer Truncations Improve Language Modeling 中，作者探讨了数据截断问题对模型效果的影响。作者指出，截断会损害数据完整性，从而阻碍模型学习基于完整上下文撰写逻辑连贯且事实一致的内容的能力。

为了解决这个问题，作者提出了 Best-fit Packing，通过长度感知组合优化来 Packing，可以完全消除不必要的截断，同时基本不影响训练效率。通过文本和代码预训练的实验结果表明，提出的方法可以在阅读理解上相对提升 4.7%，上下文跟随提升 16.8%，程序生成提升 9.2%。此外，也可以有效减少 58.3% 的封闭域幻觉问题。

PS：论文中对比实验时 Baseline 的 Concatenation 方案中有截断，并且没有使用 Block Diagonal Mask；而 Best-Fit Packing 使用了 Block Diagonal Mask，且没有截断。

6.2 方案

如下图 Figure 1 所示为本文 Best-Fit Packing 与传统方案的对比。

• 右图所示为传统方案：可以理解为将所有样本排成一行，然后按照 Max Sequence Length 进行切割，会导致大量 Sample 被截断。
• 左图为本文的方案：首先将所有 Document 按照 Max Sequence Length 截断，然后使用 Best-Fit Decreasing 算法来进行组合优化（Bin Packing 优化）。​

如下图 Table 2 所示，使用本文的 Best-Fit Packing 可以保证只增加不到 0.003% Sequence 数量，对训练效率的影响也就微乎其微。

6.3 结果

如下图 Table 3 所示，作者训练了 3 个不同规模、序列长度的模型：

如下图 Table 4 所示，提出方案训练出的模型的阅读理解能力有明显提升：

如下图 Table 9 所示，幻觉问题也可以明显降低：

如下图 Table 10 所示，作者也针对 Attention Mask 进行了相关消融实验，可以看出原始 Concatenation 方案加入 Block Diagonal Mask 后也有一定的提升。可以证明 Attention Mask 和截断都会对效果有一定影响，不过 Packing（避免截断） 的影响似乎更大一些。

七、IBM Packing with FlashAttention2

7.1 概述

在 [2407.09105] Enhancing Training Efficiency Using Packing with Flash Attention 中，作者总结了不同 Packing 策略、Mask 方式及与 FlashAttention 结合的优势。此外，作者也将相关工作提交到了 Huggingface Transformer 中，提供了新的 DataCollatorWithFlattening，具体可以参考：通过打包Flash Attention 来提升Hugging Face 训练效率。

7.2 相关方案

如下图 Table 1 所示，作者分析了不同的 Packing 方案以及它们的影响，具体包含如下几种方式：

• RandomSampling + Padding：最传统的随机采样，然后 Padding 的方式。存在冗余计算，并且占比很高。
• GroupByLength+Padding：先排序，然后尽量保证每个 Batch 中的序列长度接近。可以减少 Padding 的占比。
• RandomSampling + PosID：随机采样，但是不 Padding，而是通过 PosID 支持变长序列。几乎没有冗余计算，但可能存在明显的负载不均衡（计算量）。
• FixedLengthPacking：随机采样，随机 Packing，并且最后一个Sample 可能截断，保证填满 Max Sequence Length。没有区分不同 Sample，也就是 Causal Mask，没有冗余计算，并且负载很均衡。
• FixedLengthPacking + PosID：相比FixedLengthPacking多了 PosID，也就是可以区分不同 Sample，对应 Block Diagonal Mask。但依然会存在末尾截断，并且可能负载不均衡。
• MultiPack + PosID：使 Sequence 中的数据尽量接近 Batch 的 Max Sequence Length，降低 Sequence 中的长度不均衡，可以参考GitHub - imoneoi/multipack_sampler: Multipack distributed sampler for fast padding-free training of LLMs。需要对数据进行排序。
• SortedPacking + PosID：通过排序，使同一个 Batch 中的计算复杂度尽量接近。可以尽可能降低计算负载不均衡问题。
• RandomPacking + PosID：与FixedLengthPacking + PosID相比主要的区别就是最后一个 Sample不截断，可能存在部分 Bubble。​

7.3 结果

作者通过微调任务对比了一系列模型使用不同方案的效果和速度，其中 Max Sequence Length（msl）为 4096，每个 GPU 的 Mini-Batch Size 为 4。对应配置如下：

• no：表示最原始的RandomSampling + Padding，可以作为基线，冗余计算比较多，速度最慢，但是效果有保障。
• yes：表示FixedLengthPacking，存在交叉污染。
• flat：表示训练前 Offline Packing 好，也就是引入了排序，并且使用 PosID 实现 Document Level Mask。
• mini：表示训练中 mini batch 的 Online Packing，和 Random 类似，并且使用 PosID 实现 Document Level Mask。

如下图所示（PS：这里只是部分结果，全量请参考论文，结论基本一致，可以看出 yes 和 flat 都会对精度（VLoss）有比较大的影响，但速度（Tok/s）确实快了很多，可以达到 Baseline 的 3x-4x；而 mini 可以在保证精度的情况实现实现 2x 左右加速。

作者使用 FLAN_20k 数据集在 Mistral-7B 上针对之前提到的几种方案做了更多实验（PS：这里只是部分，gas 表示梯度累加次数），看起来 Packing 的方案都能获得不错的速度，但是精度的影响因素就比较多。但整体来说 RandomPacking + PosID 的方案还不错，看着似乎也不能有太多的梯度累加。

八、参考链接

1. ​​https://arxiv.org/abs/2107.02027​​
2. ​​https://arxiv.org/abs/2310.10638​​
3. ​​https://arxiv.org/abs/2401.18058​​
4. ​​https://xtuner.readthedocs.io/zh-cn/latest/acceleration/length_grouped_sampler.html#length-grouped-sampler​​
5. ​​https://medium.com/@ChatGLM/glm-long-scaling-pre-trained-model-contexts-to-millions-caa3c48dea85​​
6. ​​https://arxiv.org/abs/2404.10830​​
7. ​​https://arxiv.org/abs/2407.09105​​
8. ​​https://huggingface.co/blog/zh/packing-with-FA2​​
9. ​​https://github.com/imoneoi/multipack_sampler​​

本文转载自 ​​AI闲谈​​，作者： AI闲谈