引入Packed BERT以将自然语言处理的训练速度提高2倍

使用新的打包（Packing）算法，我们在训练BERT-Large时将自然语言处理（NLP）速度提高了2倍以上。我们新的打包技术去除了填充（Padding），显著提高了计算效率。

我们猜测这或许也可以应用于基因组学和蛋白质折叠模型以及其他具有倾斜长度分布的模型，以在不同行业和应用中产生更广泛的影响。

我们在一篇新论文^[1]中介绍了Graphcore高效的非负最小二乘直方图打包（Non-Negative Least Squares Histogram-Packing，NNLSHP）算法以及应用于打包序列的BERT算法。

序列填充导致的自然语言处理中的计算浪费

Computational Waste in NLP due to sequence padding

在处理我们最近的MLPerf™基准测试提交时，我们开始研究优化BERT训练的新方法，其目标是开发可以在真实世界应用中被轻松采用的、有用的优化。BERT是着重于此类优化的模型之一，因此成为自然之选，同时也是因为它已经被行业和我们的许多客户广泛采用。

当得知在我们自己使用维基百科数据集的BERT-Large训练应用中，数据集里50%的标记是填充时，我们真的很惊讶——这导致大量的计算被浪费了。

填充序列，以将它们全部对齐到相等的长度，这是GPU的常见方法。但我们认为值得尝试不同的方法。

序列的长度变化很大，原因有二：

维基百科的基础数据显示文档长度变化很大。
BERT预处理本身随机减小提取的文档的尺寸，这些文档组合在一起生成了训练序列。

将长度填充到最大长度512会导致50%的标记是填充标记。用真实数据替换50%的填充可能产生的结果是，相同的计算工作量处理的数据多50%，因此在最佳条件下速度提高2倍。

图1：维基百科数据集分布

这个情况是否是维基百科独有的呢？不是。

那么，这个情况是否只发生于特定语言呢？不是。

事实上倾斜的长度分布随处可见：在语言、基因组学和蛋白质折叠中都可以看到。图2和图3显示了SQuAD1.1数据集和GLUE数据集的分布。

图2：SQuAD 1.1 BERT预训练数据集序列长度直方图，最大序列长度为384。

图3：GLUE数据集序列长度直方图，最大序列长度为128。

我们如何在避免计算浪费的同时处理不同的长度呢？

针对不同的长度，当前的方法需要不同的计算内核，或者工程师以编程方式为每个注意力块和损失计算重复删除然后添加填充。通过放大代码并使其变得更复杂来节省计算的做法并不吸引人。所以，我们寻找更好的做法。我们难道不能将多个序列放在一个最大长度的包里一起处理吗？事实证明，我们可以。

该方法需要三个关键要素：

一种有效的算法来决定将哪些样本放在一起以尽可能减少剩余填充
调整BERT模型以处理包而非序列
调整超参数

打包

Packing

起初，似乎不太可能非常有效地打包像维基百科这样的大型数据集。这个问题通常被称为装箱。即使打包限制为三个或更少的序列，由此产生的问题仍然是NP完全（NP-Complete）的，缺乏有效的算法解决方案。现有的启发式打包算法并不乐观，因为它们的复杂度至少为\(O(n log(n))\)，其中n是序列的数量（维基百科约为16M）。我们对可以很好地扩展到数百万个序列的方法产生了兴趣。

两个技巧帮助我们大大降低了复杂性：

将一个包中的序列数量限制为3个（针对我们的第一种解决方案）。
仅对序列长度的直方图进行操作，每个出现的长度都有一个容器。

我们的最大序列长度是512。因此，转而着手于直方图的方法将维度和复杂性从1600万个序列减少到了512个长度计数。一个包中最多允许三个序列的做法将允许的长度组合数量减少到22K。这已经包含了要求序列按包中的长度排序的技巧。那么为什么不尝试4个序列呢？这将组合数量从22K增加到940K，对于我们的第一种建模方法来说太多了。此外，深度3已经实现了非常高的打包效率。

最初，我们认为在一个包中使用三个以上的序列会增加计算开销并影响训练过程中的收敛行为。但是，为了支持推理等需要更快、实时打包的应用程序，我们开发了高效的非负最小二乘直方图打包（Non-Negative Least Squares Histogram-Packing，NNLSHP）算法。

非负最小二乘直方图打包

Non-Negative Least Squares Histogram-Packing，NNLSHP

装箱问题通常被表述为一个数学优化问题。但是，对于1600万个序列（或更多），这是不切实际的。仅问题变量就超过了绝大多数机器的存储。基于直方图的方法的数学程序非常简洁。为简单起见，我们决定使用直方图向量\(b\)的最小二乘法(\(Ax=b\))。通过要求策略向量x为非负，并添加权重，我们将之扩展以允许较小的填充。

棘手的部分是策略矩阵。每列的最大总和为3，并对那些打包在一起以完全匹配所需总长度的序列进行编码；在我们的例子中是512。行对每个潜在组合进行编码，以达到总长度。策略向量\(x\)就是我们要寻找的，它描述了我们选择20k组合中的任何一个组合的频率。有趣的是，最后只选择了大约600个组合。要获得精确解，x中的策略计数必须是正整数，但我们意识到仅包含非负\(x\)的近似舍入解就足够了。对于近似解，可以使用简单的开箱即用求解器在30秒内获得结果。

图4：序列长度8和打包深度3的策略矩阵示例。行代表长度为1-8的序列打包在一起，列代表包中所有可能的长度组合，没有特定的排序。

最后，我们不得不修复一些没有被分配策略但最小的样本。我们还开发了一个变体求解器，强制每个序列都被打包，可能带有填充，并且具有依赖于填充的权重。这个方法花费的时间更长，但解决方案没有变得更好。

最短包优先的直方图打包

Shortest-Pack-First Histogram Packing，SPFHP

NNLSHP为我们提供了充分的打包方法。然而，我们想知道我们是否可以在理论上获得更快的在线能力，并消除只能将3个序列放在一起的限制。

因此，我们从现有的打包算法中汲取了一些灵感，但仍然专注于直方图。

我们的第一个算法，即最短包优先的直方图打包（Shortest-Pack-First Histogram Packing，SPFHP），有四个要素：

从最长序列到最短序列对直方图计数进行操作。
如果当前序列长度不匹配任何包，则开始一组新的包。
如果有多个匹配，取序列长度之和最短的包，分别修改计数。
再次检查剩余计数是否匹配。

这种方法实施起来最简单，只需要0.02秒。

一种变体是采用序列长度的最大总和而不是最短和分裂计数来获得更完美的匹配。但总体而言，这并没有明显改变效率，却大大增加了代码复杂性。

最短包优先直方图打包的工作原理

维基百科，SQuAD 1.1，GLUE打包结果

Wikipedia, SQuAD 1.1, GLUE packing results

表1、2和3显示了我们提出的两种算法的打包结果。打包深度描述了打包序列的最大数量。打包深度1是基线BERT实施。最大出现的打包深度，没有设置限制，用附加的“max”表示。包的数量描述了新打包数据集的长度。效率是打包数据集中真实标记的百分比。打包因子描述了与打包深度1相比产生的潜在加速。

我们有四个主要观察结果：

分布越倾斜，打包的好处就越大。
所有数据集都受益于打包。有些甚至超过2倍。
当打包深度不受限制时，SPFHP的效率更高。
对于最多3个打包序列，NNLSHP越复杂，效率越高（99.75 vs. 89.44）。

表1：维基百科上提议的打包算法（SPFHP和NNLSHP）的关键性能结果

表2：SQUaD1.1 BERT预训练的提议打包算法的性能结果

表3：GLUE数据集的提议打包算法的性能结果。只展示了基线和没有限制打包深度的SPFHP打包结果

BERT处理调整

BERT processing adjustment

BERT架构的有趣之处在于，大多数处理都发生在标记（Token）级别，这意味着它不会干扰我们的打包。只有四个组件需要调整：注意力掩码、MLM损失、NSP损失和准确性。

处理不同数量序列的所有四种方法的关键是向量化和使用可以连接的最大数量的序列。在注意力方面，我们已经有了一个掩码来解决填充问题。将其扩展到多个序列很简单，如下面的TensorFlow伪代码所示。这个概念在于，我们要确保注意力仅限于单独的序列，不能超出这个范围。

长按扫描识别左侧二维码

查看注意力掩码代码示例

图5：示例零一掩码

对于损失计算，原则上我们将序列解包并计算单独的损失，最终获得序列（而不是包）的损失平均值。

对于MLM损失，代码如下所示：

长按扫描识别左侧二维码

查看代码

对于NSP损失和准确性，原理是一样的。在我们的公共示例中，您可以使用我们内部的PopART框架^[2]找到相应的代码。

维基百科开销和加速估计

Wikipedia overhead and speedup estimate

关于对BERT的修改，我们有两个问题：

它带来了多少开销？
开销在多大程度上取决于放在一个包中的最大序列数？

由于BERT中的数据准备可能很麻烦，我们使用了一个捷径，针对多个不同的打包深度编译了代码，并对比了各自的（测量的）周期。结果在表4中显示。对于开销，我们对为了启用打包而进行的模型更改（例如用于注意力的掩蔽方案和更改的损失计算）导致的吞吐量百分比下降做出表示。实现的加速是由于打包（打包因子）导致的加速和由于开销导致的吞吐量下降的组合。

表4：维基百科上提议的打包算法（SPFHP和NNLSHP）的估计加速对比

有赖于向量化技术，开销小得惊人，而且将多个序列打包在一起也没有任何劣势。

超参数调整

Hyperparameter-adjustments

通过打包，我们将有效批尺寸（平均）增加了一倍。这意味着我们需要调整训练超参数。一个简单的技巧是将梯度累积计数减半，以保持与训练前相同的有效平均批尺寸。通过使用带有预训练检查点的基准设置，我们可以看到准确度曲线完全匹配。

图6：打包和非打包加工的学习曲线对比，打包方法的批尺寸减少。

准确率匹配：MLM训练损失在开始时可能略有不同，但很快就追赶上来。这种初始差异可能来自于注意力层的轻微调整，这些调整可能在之前的训练中偏向于短序列。

为避免速度变慢，一种有效的做法是，保持原始批尺寸不变，并将超参数调整为增加后的有效批尺寸（翻倍）。要考虑的主要超参数是beta参数和学习率。一种常见的方法是将批尺寸翻倍，而在我们的案例中这会降低性能。查看LAMB优化器的统计数据，我们可以证明，将beta参数提高到打包因子的幂，对应连续训练多个批次，以保持动量和速度的可比性。