Efficient Sequence Packing without Cross-Contamination: Accelerating Large Language Models without Impacting Performance

发表时间: 2021-07 · arXiv:2107.02027 (Graphcore)

作者/机构: Mario Michael Krell (Graphcore Inc.), Matej Kosec (Graphcore Inc.), Sergio P. Perez (Graphcore Inc.), Andrew Fitzgibbon (Graphcore Inc.)

速读

一句话结论 本文提出了一种高效且硬件无关的序列打包算法及配套的模型调整方案，在不影响模型收敛和下游任务性能的前提下，通过消除训练数据中的填充令牌，使大型语言模型预训练的吞吐量提升了近2倍。

要解决什么问题 大型语言模型在预训练时，通常需要处理高度倾斜的序列长度分布。为了对齐张量维度，业界普遍采用填充策略，这导致大量计算资源被浪费在无意义的填充令牌上。例如在标准的 Wikipedia 数据集中，当最大序列长度设为 512 时，填充令牌占比高达 50%，这意味着有一半的算力在做无用功。现有的优化路线均存在致命的机制卡点：第一种是排序批处理（SORT），即把长度相近的样本凑成一批，但这不仅破坏了数据的独立同分布假设，还会导致计算图尺寸不断变化，在 XLA 等针对固定张量优化的现代编译器上会引发极大的重新编译开销；第二种是去填充方法，依赖定制内核在计算时动态剔除填充，但这会让每个批次的处理时间随序列长度剧烈波动，在多加速器集群训练时会引发严重的设备间负载不平衡，导致所有设备都在等待最慢的节点，彻底拖垮分布式扩展性；第三种是贪心拼接，即直接把短序列拼成长序列，但如果不做隔离，不同序列的令牌会在自注意力计算中相互影响，产生交叉污染，从而导致模型精度下降。

怎么做的 本文的核心思路是将变长序列拼接成固定长度的包，在预处理阶段离线解决填充问题，并在模型层引入严格的隔离机制。这种做法既保持了张量尺寸和计算时间的绝对固定，完美绕开了动态编译开销和分布式负载不平衡的卡点，又在数学上保证了与原模型的等效性。整个方法由打包算法和模型改造两个关键部件构成。首先，作者将序列打包转化为经典的箱子打包问题，并提出了非负最小二乘直方图打包（NNLSHP）算法。该算法不处理单个样本，而是直接对整个数据集的序列长度直方图进行操作。算法预先枚举出所有总长度不超过最大限制且包含序列数不超过设定深度的打包策略，构建策略矩阵 $A$。随后求解加权非负最小二乘问题： $$wAx = wb$$ 其中 $b$ 是数据集中序列长度计数的直方图向量，$x$ 是每种策略的重复次数，$w$ 是残差权重（用于优先保证长序列被打包，允许极短序列作为填充残留）。求解出的 $x$ 经过取整后即可生成近乎零填充的全局打包方案。其次，为了让模型处理打包数据时与原版等效，作者对标准 BERT 架构进行了三项改造：第一，将位置编码从连续的偏置加法改为独立的嵌入查找，确保包内每个新序列的位置索引都能从零重新开始；第二，引入块对角注意力掩码，在自注意力 softmax 操作前直接掩蔽掉不属于同一序列的令牌，从根本上阻断了序列间的交叉污染；第三，在计算损失时，利用掩码权重张量将输出的 logits 和标签在内部解包，恢复成逐序列计算交叉熵并求和的模式，避免了因包内序列数量不同而导致的损失权重失衡。

效果如何 实验在 Graphcore IPU-M2000（代表静态编译架构）和 Habana Gaudi 加速器上进行，使用 BERT-large 和 BERT-base 模型，在包含大量填充的 Wikipedia 数据集上执行 MLPerf 预训练，并在 SQuAD 1.1 上测试下游微调。对比基线包括：NONE（不打包的标准实现）、SORT（代表按长度排序批处理的路线）、GREEDY（代表 T5 采用的贪心拼接路线）以及去填充方法（代表动态内核路线）。量化结果显示，NNLSHP 算法实现了 99.75% 的打包效率，在预训练阶段带来了 1.913 倍的实际吞吐量加速，极其逼近 2 倍的理论极限，而 SORT 基线因重编译开销反而比不打包更慢，GREEDY 基线仅获得 1.566 倍加速。在多设备扩展性测试中，当加速器数量增加到 2048 个时，本文方法的加速比稳定保持在 1.913 倍，而代表动态路线的去填充方法因负载不均，加速比从 2 倍骤降至 1.3 倍。在下游任务表现上，经过打包预训练的模型在 SQuAD 1.1 上的 F1 和 EM 分数与标准模型差异在 0.3% 以内，证明了方法的无损性。该方法的代价是模型改造会引入低于 5% 的计算开销，且由于打包增加了有效批次大小，需要对 LAMB 优化器的衰减参数进行指数级启发式调整才能完全对齐原有的收敛曲线。作者也承认，该方法的局限性在于：如果数据集本身已经被预处理成了无填充的满长序列，打包算法将失去优化空间。

A1 主要贡献

本文旨在解决大型语言模型（LLM）训练中因处理可变长度序列而普遍采用填充（padding）策略所导致的计算资源浪费问题。研究表明，在常见的自然语言处理（NLP）数据集中，填充令牌的比例可高达50%，在某些情况下（如GLUE-cola数据集，序列长度128）甚至达到89%。现有的解决方案存在诸多弊端：例如，在自注意力机制中需要避免序列间的“交叉污染”（cross-contamination）；丢失序列排序信息可能导致模型精度下降；或者依赖于仅适用于特定加速器的定制化内核实现，通用性差。

针对以上问题，本文的核心研究目标是开发一种高效且硬件无关的序列打包（sequence packing）方法，以加速LLM训练，同时确保与原始模型在数学上等效，不影响模型性能。

主要贡献如下：

1. 揭示填充问题的普遍性：通过对多种数据集（包括Wikipedia、GLUE、SQuAD等）的序列长度分布进行可视化分析，量化了填充令牌在训练数据中占据的大量比例，明确了通过序列打包可获得的巨大加速潜力（例如，在BERT第二阶段预训练中可实现2倍加速）。
2. 提出新颖高效的打包算法：将序列打包问题形式化为经典的箱子打包问题（bin packing problem），并基于此提出了两种全新的确定性高效打包算法：最短包优先直方图打包（SPFHP）和非负最小二乘直方图打包（NNLSHP）。这两种算法能在数秒内为包含数百万序列的数据集生成近乎最优的打包方案。
3. 解决交叉污染并保持模型等效性：详细阐述了“交叉污染”问题，即打包后不同序列的令牌在自注意力计算中相互影响，并提出了一系列模型调整方案来避免此问题，包括调整位置编码和引入块对角注意力掩码。这些调整确保了打包后模型的训练行为与原模型在数学上等效，从而可以直接沿用现有的预训练和微调流程。
4. 提供全面的实验验证：通过在BERT large模型上的实验，证明了所提出的打包算法能为Wikipedia预训练数据集带来近2倍的吞吐量提升，且收敛行为与未打包的基线模型相当。此外，实验还验证了模型调整的必要性，并证明打包训练不影响在下游任务（如SQuAD）上的最终性能。

A3 背景知识与关键观察

序列长度分布

本文的一个关键观察是，许多用于语言模型训练的数据集具有高度倾斜的序列长度分布，导致大量计算资源被浪费在处理填充令牌上。

• Wikipedia预训练数据集的观察: 该数据集是BERT预训练的事实标准。即使经过了句子级别的打包预处理，在序列长度为512时，填充令牌仍占总令牌数的50%。这意味着通过避免处理填充令牌，理论上可以获得2倍的训练速度提升。具体来说，长度为512的样本仅占数据集的23.5%，而序列长度范围从5到512不等。随着最大序列长度的增加，填充比例和理论加速比也随之增加：序列长度128时加速比约为1.2倍，384时为1.7倍，512时为2.0倍。据作者所知，这类直方图此前未被发表，可能导致业界低估了序列打包的加速潜力。

• 问题的普遍性: 这种倾斜分布不仅限于Wikipedia数据集。如图1所示，GLUE、SQuAD 1.1（理论加速2.2倍）、LibriSpeech文本标签和音频数据，甚至非文本领域的QM9分子数据（理论加速1.6倍），都表现出类似的特性。这表明序列打包是一种具有广泛适用性的优化技术。对于LibriSpeech音频数据，由于其分布偏向长序列，尽管理论最大加速比为1.6倍，但实际通过打包仅能实现1.3倍的加速。对于其他数据集，本文提出的算法能够实现接近最优的打包效果。

图 1: 不同数据集的序列长度分布。左上三个图表显示了基于2020年10月1日维基百科文章转储的BERT预训练数据集在不同最大序列长度下的序列长度直方图（不包括填充的令牌计数）。理论加速比指的是不使用任何填充令牌且没有处理不同长度的开销。右上：GLUE数据集。底部从左到右：SQuAD 1.1、LibriSpeech文本标签、LibriSpeech音频令牌序列，以及QM9分子（图中序列化的图）。

相关工作

• 排序批处理（SORT）: 这是减少填充的一种直接方法，即在批处理前按大小对样本进行分组，将长度相近的样本放在一个批次中。例如，Faster Transformer 【【21，Faster Transformer，2021，https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer/v1】】 、lingvo 【【28，Lingvo: a modular and scalable framework for sequence-to-sequence modeling，2019】】、fairseq 【【22，fairseq: A fast, extensible toolkit for sequence modeling，2019】】 和RoBERTa 【【16，RoBERTa: A Robustly Optimized BERT Pretraining Approach，2019】】 都采用了这种方法。尽管计算效率高，但存在多个缺点：首先，排序破坏了数据分布的独立同分布（i.i.d.）假设，可能降低大批量训练时的收敛速度。其次，处理短序列批次会造成计算资源未被充分利用。虽然GPU上可以通过动态调整批次大小来缓解，但这增加了复杂性。第三，现代NLP应用常使用XLA等编译器针对固定张量大小进行优化，更改序列长度需要重新编译，带来显著开销。

• “去填充”（Un-padding）: 这类更高级的方法依赖于定制的计算内核，在计算过程中动态地移除和恢复填充值，如Effective Transformer 【【5，Effective Transformer，2021，https://github.com/bytedance/effective_transformer】】。虽然能完全避免处理填充，但引入了显著的内核启动开销。更严重的是，每个批次的处理时间会根据其中的序列长度而波动。在多加速器环境中，这会导致设备间的负载不平衡，所有设备都必须等待最慢的设备完成，从而严重削弱了在大规模集群上的加速效果 。

• 其他模型的打包策略: 基于因果注意力（causal attention）的模型（如GPT-3 【【4，Language Models are Few-Shot Learners，2020，NeurIPS】】）的打包问题相对简单，因为它们可以任意截断序列而不影响前面令牌的计算。GPT-3将多个句子片段拼接成一个长序列，并在段落间不使用注意力掩码以追求计算效率。RoBERTa和T5 【【24，Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer，2019】】 也采用类似的贪心打包策略，但RoBERTa的消融研究表明，混合来自不同文档的句子会降低准确性。本文提出的方法通过注意力掩码解决了这一交叉污染问题，确保了原始数据集和模型特性的匹配。

• 现有技术水平: 尽管一些已弃用的库（如tensor2tensor）中可能存在类似的注意力掩码机制，但它们缺乏充分的文档、测试、评估和社区沟通，未被视作NLP研究的当前最优方法。据作者所知，没有其他研究工作系统地关注NLP中的序列打包问题。

关于打包的技术背景

• F.1 经典打包问题: 箱子打包问题（Bin Packing Problem）旨在将一组物品放入容量固定的箱子中，目标是最小化所用箱子的数量。这是一个经典的强NP完全问题 【【14，Combinatorial Optimization，2012】】，对于大规模问题（如Wikipedia数据集有1600万序列），求解精确最优解是不现实的，需要近似方法。

• F.2 近似箱子打包问题: 近似算法分为在线（online）和离线（offline）两类。在线算法逐个处理序列，而离线算法可以访问所有待打包样本。现有算法的时间复杂度至少为$O(n \log(n))$。本文提出的方法通过三个主要简化来降低复杂度：首先，算法操作于序列长度的直方图而非单个样本，使问题规模与样本数$n$解耦；其次，专注于固定的打包深度（如3-partition问题）；最后，通过非负最小二乘（NNLS）或在排序直方图上应用最差拟合（Worst-Fit）来近似求解。

• F.3 定义:

  • 包（Pack）/箱（Bin）: 可互换使用。
  • 最大打包深度: 一个包中允许容纳的最大序列数。
  • 打包策略: 一个序列长度的有序列表，其总长度不超过箱子容量$s_m$，且序列数量不超过最大打包深度。
  • 策略重复计数: 每种打包策略被重复使用的次数。

• F.4 非负最小二乘直方图打包（NNLSHP）:

  • F.4.1 枚举打包策略: 通过动态规划列出所有在最大打包深度内可行的打包方式。例如，对于深度3，策略可以是 [512]、[6, 506] 或 [95, 184, 233]。为避免重复，策略内的序列长度强制排序。
  • F.4.2 构建打包矩阵: 构建一个矩阵$A$，其行数等于最大序列长度，列对应每一种打包策略。矩阵$A$中的元素表示一个策略使用了多少个特定长度的序列。表4展示了一个示例。
  • F.4.3 NNLS问题求解: 将问题形式化为求解$A \cdot x \approx b$，其中$b$是序列长度直方图，$x$是策略的重复计数组。使用非负最小二乘（NNLS）求解器 【【3，A fast non-negativity-constrained least squares algorithm，1997，Journal of Chemometrics】】得到$x$。
    
  • F.4.4 将填充视为残差: 求解后，将浮点解$x$四舍五入为整数，并计算残差$r = b - A \cdot x_{rounded}$。残差的负值部分表示“短缺”的序列，需要用填充来补足；正值部分表示未被打包的“剩余”序列。图7展示了残差的分布，通常短序列短缺，长序列剩余。
    
    图 7: NNLS打包问题残差的可视化
  • F.4.5 残差加权: 为了优先打包长序列，可以对NNLS问题进行加权，即求解$\min ||w(Ax - b)||^2$。通过降低短序列（如长度小于8）残差的权重（例如设为0.09），可以鼓励算法使用更多短序列（即使需要生成填充）来打包长序列，从而减少长序列的剩余。如图8所示，加权后，残差几乎完全转移到短序列上。
    
    图 8: NNLS打包问题加权残差的可视化

• F.5 残差权重选择的讨论: 通过调整“偏移量”（offset，即权重变化的长度阈值）和“权重”（weight）两个参数来控制加权。实验发现，权重选择对最终的打包效率影响很小（小于1%），但可以有效地实现次要目标，如避免剩余长序列。例如，在SQuAD数据集上，通过网格搜索微调权重，效率从96.94%提升到98.767%。这表明该方法具有一定的灵活性。

A2 方法细节

本文的方法包含三个部分：首先，在预处理阶段高效地打包数据样本以充分利用最大序列长度$s_m$（第3.1节）；其次，对BERT模型进行一系列修改以保持其与原始实现的等效性，包括使用自注意力掩码防止序列间交叉attend（第3.2.2节）和调整位置编码（第3.2.1节）；最后，提供超参数调整建议，以使打包和未打包的BERT实现具有相似的收敛行为（第3.3节）。

3.1 打包算法

背景介绍。著名的箱子打包问题旨在将物品装入固定容量的箱子，以最小化使用的箱子数量。由于精确解是强NP完全的【【14，Combinatorial Optimization，2012】】，因此已提出了众多近似解【【12，Near-optimal bin packing algorithms，1973】；【15，A Simple On-Line Bin-Packing Algorithm，1985，Journal of the ACM (JACM)】；【13，A 7160 theorem for bin packing，1985，Journal of Complexity】；【36，A simple proof of the inequality MF F D(L) ≤ 71/60OP T (L) + 1, L for the MFFD bin-packing algorithm，1995】】。鉴于现有近似算法的复杂度至少为$O(n \log n)$，本文提出了两种新的启发式离线算法，它们针对NLP场景进行了优化，并应用于整个数据集。

3.1.1 最短包优先直方图打包 (SPFHP)

• 算法原理。最短包优先直方图打包（SPFHP）算法操作于序列长度直方图（箱子大小为1）而非单个样本。该算法按从长到短的顺序遍历直方图的箱子。在遍历过程中，应用最差拟合（worst-fit）算法【【12，Near-optimal bin packing algorithms，1973】；【36，A simple proof of the inequality MF F D(L) ≤ 71/60OP T (L) + 1, L for the MFFD bin-packing algorithm，1995】】，将当前处理的直方图箱子放入剩余空间最大（即“最短的包”）的“包”中。如果直方图箱子不能完全放入，则创建一个新的包。
• 实现细节。算法还限制了打包深度，即一个包中允许的最大序列数。因此，只有当一个现有包未达到最大打包深度时，才会向其添加序列。算法的详细代码在附录的清单3中提供（实际在清单2中）。
• 算法复杂度。该算法的时间和空间复杂度分别为$O(n + s_m^2)$和$O(s_m^2)$（详见附录G.2【1】）。

3.1.2 非负最小二乘直方图打包 (NNLSHP)

• 问题形式化。NNLSHP算法将打包问题重述为一个（加权的）非负最小二乘（NNLS）问题【【3，A fast non-negativity-constrained least squares algorithm，1997，Journal of Chemometrics】】，形式为$wAx = wb$，其中$x \ge 0$。向量$b$是数据集中所有序列长度计数的直方图。
• 打包矩阵A的构建。我们首先生成所有可能的序列长度组合（称为“策略”），这些组合的总和恰好等于最大序列长度。我们特别关注每个包最多包含3个序列的策略，并将每个策略编码为稀疏矩阵$A$的一列。例如，一个包含长度为128、128和256的序列的策略，表示为一个列向量，其在第128行值为2，第256行值为1，其他行均为0。
• 求解与加权。变量$x$描述了每种策略的非负重复次数。在无加权的情况下，$Ax=b$表示我们希望混合预定义的策略（$A$的列），使得样本数量与直方图$b$匹配。我们使用残差权重$w$来控制对不同序列长度上$Ax-b$残差的惩罚。启发式地，我们将长度为8或更短的序列的权重设置为0.09，因为它们被认为是可接受的填充序列，而所有其他序列长度的权重为1。
• 求解与后处理。在使用现成的求解器求解$wAx = wb$后，我们得到一个浮点解$\hat{x}$，然后将其四舍五入到最近的整数。实验发现，这种四舍五入引入的误差与数据集大小相比可以忽略不计。
• 算法复杂度。该算法的时间和空间复杂度分别为$O(n + s_m^5)$和$O(s_m^3)$。更多细节在附录F.4中提供。

3.2 packedBERT: 模型变更

本节描述了如何修改任何标准的BERT实现以处理打包序列，从而使模型的行为与处理未打包序列时相同。保持数学等效性对于确保现有BERT预训练和微调实践的有效性至关重要，也是MLPerf™【【17，MLPerf: An Industry Standard Benchmark Suite for Machine Learning Performance，2020】】等基准测试的要求。

3.2.1 调整位置嵌入

• 问题与原因。BERT模型使用令牌、段和位置三种嵌入。位置嵌入通常实现为偏置加法操作，因为位置索引对每个序列都是线性增加的。然而，当使用打包数据格式时，位置索引需要在每个新打包的序列开始时重置。例如，打包一个长度为2和另一个长度为3的序列，需要拾取的位置嵌入索引是 [0, 1, 0, 1, 2]。
• 解决方案。为了实现这一点，必须将偏置加法替换为嵌入查找（embedding look-up），以便为包中的每个令牌提取正确的位置嵌入。这还需要一个额外的输入来指定每个令牌在其序列中的位置。
• 影响。这一调整对绝对准确率/损失的影响很小（见第4.2和4.2.1节）。

3.2.2 调整注意力掩码

• 问题与原因（交叉污染）。为了保持与未打包版本的一致性，一个包内来自不同序列的令牌不应相互关注（attend to each other）。其他实现通常通过使用自定义注意力内核来解包序列，然后对每个序列单独进行注意力计算【【5，Effective Transformer，2021，https://github.com/bytedance/effective_transformer】】 。
• 解决方案。我们提出在注意力softmax之前，直接用一个块对角掩码来掩盖注意力矩阵。这在现代框架中很容易实现（见图2）。
• 实现与成本。虽然构建和应用掩码会产生一些开销，但这是保持准确性所必需的（见表1，第4.1节，第4.2节）。

图 2: 注意力掩码代码[左]，对应的0-1掩码[中]，以及序列损失的向量化解包[右]。白色矩形对应填充。

3.2.3 调整逐序列损失和准确率

• 问题与原因。BERT的典型实现是按令牌计算掩码语言模型（MLM）的交叉熵损失。然而，其他NLP任务（如SQuAD）以及下一句预测（NSP）任务，是在逐个序列的基础上计算损失和准确率的。如果直接将打包序列输入到相同的交叉熵实现中，会导致按包计算加权损失，即一个包的损失权重与其中包含的序列数无关，这将导致模型收敛到与未打包实现不同的最优点。
• 解决方案。为了恢复标准未打包BERT实现的逐序列平均行为，我们有效地“解包”输入的logits和标签。一旦序列被解包，我们就可以像往常一样分别计算每个序列的损失，然后将损失相加。为了最小化解包损失计算的延迟开销，我们并行计算所有索引，而不是循环遍历序列索引（见图2）。
• 实现示例。我们使用“masked lm weight”【【7，BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding，2019，https://github.com/google-research/bert】】输入张量来表示给定的掩码令牌属于哪个序列（0、1、2等），这与标准BERT实现一致。完整方法在附录清单5中详述，并可应用于其他分类或预训练任务 。

3.3 调整超参数

• 打包带来的影响。在收敛行为方面，打包的主要后果是有效批次大小（就序列数和真实令牌数而言）的增加，并且在不同迭代之间增加了一些变异。
• 直接解决方案。一个直接的解决方案是按打包因子（每个包的平均序列数）减少计算批次大小，并保持所有其他超参数不变。例如，如果打包因子为2，将梯度累积计数减半就足够了。这种策略的优点是不需要微调超参数，并且性能曲线具有可比性。但缺点是可能导致内存/计算资源未被充分利用。
• 优化硬件利用率的方案。为了保持批次大小并优化硬件利用率，我们额外提出了一个更新LAMB优化器【【35，Large Batch Optimization for Deep Learning: Training BERT in 76 minutes，2019】】衰减参数的近似启发式方法。对于一个打包因子为$p$的打包数据集，我们将衰减参数更新为：$\beta_1 := \beta_1^p$, $\beta_2 := \beta_2^p$。当$p=2$时，这对应于用相同的梯度更新两次时计算动量和速度的精确参数（见附录D）。
• 关于学习率的说明。虽然根据批次大小缩放学习率是一种常见方法，但我们的实验（第4.2节）表明，这样做会降低收敛速度。
• 启发式方法的局限性。由于这些调整只是启发式方法，模型的收敛将是可比较的但并非完全相同。特别是，仅调整超参数不太可能完全消除批次大小增加带来的影响。然而，通过这些调整，研究人员应该能够继续使用现有的配置。

A4 实验

实验环境

• 数据集:

  • Wikipedia: 用于BERT预训练的核心数据集。在序列长度512时，包含83.3亿令牌，其中41.7亿是填充令牌。
  • MLPerf™ v0.7: 用于BERT预训练的标准化基准测试，使用Wikipedia数据集中300万个最大长度为512的序列。
  • SQuAD 1.1: 用于下游任务微调和评估。
  • GLUE: 包含多个NLU任务的数据集，用于展示打包的普适性。
  • LibriSpeech: 音频和文本数据集，用于跨领域验证。
  • QM9: 分子数据集，用于非文本领域验证。

• 模型架构:

  • BERT-large: 用于MLPerf™基准测试和全流程预训练。
  • BERT-base: 用于全流程预训练。

• 硬件配置:

  • 主平台: Graphcore IPU-M2000，包含16个IPU芯片。选择IPU是因为其静态编译特性，可以凸显本文方法（避免代码重编译）的优势。
  • CPU: Intel(R) Xeon(R) Gold 6138 @ 2.00GHz，用于运行打包算法。
  • 对比平台: 实验结果也在Habana Gaudi加速器上得到复现。与GPU上的动态方法（如un-padding）进行了扩展性分析。

• 软件配置:

  • 实现: 模型修改在框架级别完成，无需定制内核。代码基于标准的BERT实现，并提及了TensorFlow和PopArt（用于IPU自定义操作）。
  • 代码库: 提供了打包算法和模型修改的参考实现，并已公开。

实验结果

4.1 箱子打包算法比较

本实验评估了不同打包算法在IPU上的性能，关键指标包括打包效率、开销（Overhead）和实际加速比。

• 算法: 比较了本文提出的SPFHP、NNLSHP与基线（NONE，即无打包）、排序批处理（SORT）和贪心打包（GREEDY）。

• 结果 (见表1):

  • NNLSHP (深度3) 表现最佳，实现了1.913倍的加速比，非常接近2.001的理论上限。其打包效率高达99.75%。
  • SPFHP 同样高效，且计算时间极短（约0.03秒）。
  • GREEDY（T5中使用的方法）效率较低，加速比仅为1.566倍。
  • SORT 虽然打包效率高，但在IPU这种静态图架构上，因计算图不断变化导致巨大开销，反而比无打包更慢。
  • 模型修改开销 (OH) 非常小，随着打包深度的增加略有上升，但远小于打包带来的收益。

• 结论: 本文提出的打包算法（特别是NNLSHP）能在极小的开销下实现接近理论极限的加速，且方法具有硬件和软件的独立性，已在Habana Gaudi加速器上得到验证。

Table 1: 提出的打包算法（SPFHP和NNLSHP）在IPU上的关键性能结果。

4.2 MLPerf™ 第二阶段预训练：学习曲线与超参数调整

本实验在MLPerf™ v0.7基准上比较了打包（NNLSHP深度3）与未打包（经典BERT）的收敛行为和实际加速效果。

• 实验设置:
  1. 相同有效批次大小: 通过将打包训练的梯度累积数减半，使平均批次中的序列数与未打包训练相同。
  2. 不同启发式调整: 保持计算批次大小不变，评估第3.3节提出的超参数调整启发式方法。
  3. 优化设置对比: 比较两种方法达到目标准确率所需的总时间。

• 结果 (见图3):

  • 设置1: 当有效批次大小相同时，打包和未打包的训练曲线几乎完全重合，证明了打包训练在数学上的等效性。
  • 设置2: 使用启发式调整后，打包训练因有效批次更大，早期收敛稍慢，但调整LAMB衰减参数后，后期性能与未打包相当。
  • 设置3: 实际总加速比超过了2倍。这是因为除了减少计算量，打包还降低了数据传输到设备的延迟。

• 结论: 打包训练不仅能达到与未打包训练相同的收敛效果，而且实际加速效果甚至超过了基于计算量减少的理论预期。

图 3: 打包和未打包处理的学习曲线比较。[左] 相同的有效批次大小（ebs是批次大小乘以打包因子），[中] 不同的超参数启发式调整，[右] 打包实现的实际加速比（超过期望的2倍）。

4.2.1 消融研究

本研究分析了模型修改的必要性。

• 结果 (见图4):
  • 无注意力掩码调整: 如果不进行注意力掩码调整以防止交叉污染，训练损失和准确率会急剧恶化。
  • 无位置编码调整: 如果不调整位置编码，准确率会停滞在71.8%，无法达到72.1%的目标。
  • NSP损失的影响: 在完整的训练设置中，移除NSP损失会导致下游SQuAD任务的F1分数下降1.31%，EM分数下降1.15%。

• 结论: 本文提出的注意力掩码和位置编码调整对于保持模型性能至关重要。同时，NSP损失对于某些任务仍然是必要的。

图 4: 在我们的打包BERT方法中，使用和不使用掩码或位置嵌入调整的学习曲线比较。要达到的灰色准确率基线是72.1%。

4.3 全程预训练和SQuAD微调

本实验旨在验证从头开始的打包预训练是否会影响下游任务的性能。

• 设置: 对BERT-base和BERT-large模型进行完整的两阶段预训练（阶段二序列长度384），然后在美国SQuAD 1.1数据集上进行微调。为了公平比较，打包训练通过减少梯度累积数来匹配未打包训练的有效批次大小。
• 结果:
  • 加速比 (见表2): 打包训练在预训练阶段实现了显著的加速，例如BERT-large在阶段二的加速比为1.63倍。
  • 下游任务性能 (见表3): 经过打包预训练的模型，在SQuAD 1.1上的F1和EM分数与使用标准预训练的模型相当，差异在0.3%以内，属于正常波动范围。

• 结论: 打包策略不会对模型的下游任务性能产生负面影响。

Table 2: BERT预训练中使用打包的实测加速比。

Table 3: 使用打包进行BERT预训练后在SQuAD 1.1上的得分。

4.4 扩展性分析：加速器数量的影响

本节分析了本文的打包方法与“去填充”（un-padding）方法在多加速器环境下的扩展性。

• 核心论点: “去填充”方法依赖动态内核，每个批次的处理时间不固定，导致多设备间存在负载不平衡。设备必须等待最慢的那个，从而降低了整体加速效果。而本文的打包方法，每个批次的处理时间是固定的，天然具有负载均衡特性。
• 结果 (见图5):
  • 打包方法: 加速比（1.913倍）不随加速器数量的增加而改变，表现出完美的扩展性。
  • 去填充方法: 随着加速器数量从1增加到2048，其理论加速比从超过2倍急剧下降到仅1.3倍。

• 结论: 在大规模分布式训练中，本文提出的打包方法相比“去填充”等动态方法具有显著的扩展性优势。

图 5: 随着加速器数量增加，理论加速比的比较。

A5 结论

本文系统地研究并解决了LLM训练中的序列打包问题。
首先，通过对多种数据集（涵盖语言、音频、分子领域）的序列长度分布进行可视化，揭示了填充令牌普遍存在导致计算资源大量浪费的现象，一些场景下通过移除填充可获得超过2倍的加速。
其次，本文提出了两种基于成熟求解器的高效打包算法（SPFHP和NNLSHP），它们能在数秒内为海量数据集生成近乎无填充的打包方案，远优于现有缓慢且次优的方法。
第三，本文论证了若不对序列处理算法（如BERT）进行相应调整，打包会导致预测性能下降。为此，提出了一系列必要的模型调整（如注意力掩码、位置编码调整），以保持预测性能。
最后，通过全面的实验证明，得益于这些调整，打包训练能够保持与未打包训练等效的预测性能，同时带来显著的速度提升（超过2倍），且模型修改引入的开销低于5%。实验还证实，打包不会影响下游任务的性能。

未来工作展望：

1. 扩展到计算机视觉领域：一个有趣的方向是将打包技术应用于不同尺寸的图像，以加速计算机视觉应用，特别是考虑到Transformer在视觉领域的广泛应用（如ViT）。
2. 改进其他模型：应用本文提出的避免交叉污染的方法来改进RoBERTa、GPT-3、T5等模型，解决它们在拼接不同文档的非连续片段时可能出现的性能下降问题。甚至BERT本身也可能从避免两个拼接段之间的污染中受益。

A6 附录

A 更广泛的影响

减少碳足迹。我们展示了在Wikipedia上预训练BERT时，处理填充令牌的计算开销约为50%。通过消除这部分浪费的计算时间，本文提出的方法为将基于BERT的模型的训练碳足迹减半铺平了道路。

提升技术可及性。此外，我们的方法避免了对自定义内核的需求，使得更广泛的NLP从业者能够轻松获得打包带来的好处。因此，我们希望这项研究能对NLP社区产生积极影响，并且我们认为使用这种方法没有任何缺点。

方法的适用前提与未来探索。我们算法的益处基于两个假设：训练数据集中存在倾斜的长度分布，以及硬件设置在固定批次大小上能高效训练。如果硬件能高效处理可变批次大小，那么像FasterTransformer和fairseq的排序批处理方法将产生相同甚至更大的益处。如果数据集生成方式不同，如GPT模型和RoBERTa（FULL-SENTENCES），所有序列都已满长，则无法拼接，打包也就没有益处。然而，达到满长序列的策略通常会混合来自不同无关文档来源的片段，这可能导致性能下降。我们的论文引入了一种避免序列间污染的方法，同样的方法也可以应用于避免段落间的污染，探索其在BERT预训练之外的益处是未来的工作。

跨语言与跨领域适用性。未来的工作需要研究打包在不同文化和语言生成的文本上的适用性。我们已经证明，使用我们的方法带来的加速不仅发生在Wikipedia上预训练BERT时，也发生在SQuAD和GLUE等其他数据集上。此外，原始英语文本的句子长度分布也显示出类似的特征。我们的研究使我们相信，可压缩的分布在语言任务及其他领域中自然产生，例如DNA序列长度、蛋白质长度和语音（M节）。许多此类序列建模工作负载都基于BERT/transformer架构的变体，因此很容易从我们的加速中受益。

潜在风险与规避。NLP的失败可能对社会产生重大影响。虽然我们的方法产生的错误可以避免，但一个潜在的错误来源是实现。注意力掩码和逐序列损失都需要修改以支持打包。这些更改比自定义内核所需的更改要小得多，但实现和调试仍可能耗时。为了帮助降低任何实现错误的风险，我们在附录中分享了所需更改的参考实现。

B 可复现性声明

所有打包算法的代码均在附录（U节）中提供，并直接链接到我们的GitHub页面，以简化下载和使用。我们甚至为不同变体提供了代码，以及为BERT训练或微调而分词的不同数据集的序列长度直方图。

为了生成学习曲线，可以使用我们向MLPerf™提交的公开结果，并且我们正在准备在其他框架中发布更多代码。为了鼓励在打包序列的模型中使用这些调整，我们还在TensorFlow中提供了详细的解释和代码片段。

本附录提供了详细的数学公式（E节和F节）、一个定理证明（D节）和复杂度计算（G节），以全面支持我们在论文中的主张。

D 关于LAMB超参数校正启发式的定理

问题背景。使用打包时，有效批量大小会改变，因此需要调整LAMB优化器【【35，Large Batch Optimization for Deep Learning: Training BERT in 76 minutes，2019】】的超参数。对于打包因子为$p$的打包数据集，我们按如下方式更新衰减参数：$\beta_1 := \beta_1^p$, $\beta_2 := \beta_2^p$。例如，如果未打包数据集的$\beta_1 = 0.81$，那么对于平均每个样本包含2个序列的打包数据集，应使用值$0.81^2 \approx 0.66$。

定理。假设梯度变化很小或没有变化，并且$p$是自然数，我们可以证明这个启发式方法是确保LAMB中的动量和速度不受打包影响的精确解。这可以通过数学归纳法证明。注意，根据定义$p \ge 1$。

定理 D.1: 对于任何$p \in N$，并假设在$b$个随机样本批次上的相应梯度（近似）相同，在LAMB优化器中选择超参数
可以确保经过$p$个独立更新步骤后的动量和速度与使用$p \times b$个样本进行一次打包更新步骤后的动量和速度相同。

证明:

• 基础情况: 当$p=1$时，等式左右两边相同，这与未打包的情况完全匹配。因此，定理对$p=1$成立。
• 归纳假设: 假设定理对所有直到某个$k$ ($k \ge 1$)的值$p$都成立。
• 归纳命题: 定理对$p = k + 1$成立。
• 归纳步骤证明: 让$l$为损失函数，$w_t$为$t$次