[论文解读] Packing: Towards 2x NLP BERT Acceleration

Matej Kosec, Sheng Fu+1 more|arXiv (Cornell University)|Jun 29, 2021

VLSI and FPGA Design Techniques参考文献 6被引用 4

一句话总结

本文提出基于确定性直方图的打包算法，以消除 BERT 预训练中的填充，通过将多个序列合并为单个样本，实现每秒序列数 2 倍的加速。所提方法 SPFHP 与 NNLSHP 基于序列长度直方图运行，实现线性时间效率，实现近乎最优的打包效果，对模型性能影响极小，并在 BERT-Large 预训练中实现完全收敛。

ABSTRACT

We find that at sequence length 512 padding tokens represent in excess of 50% of the Wikipedia dataset used for pretraining BERT (Bidirectional Encoder Representations from Transformers). Therefore by removing all padding we achieve a 2x speed-up in terms of sequences/sec. To exploit this characteristic of the dataset, we develop and contrast two deterministic packing algorithms. Both algorithms rely on the assumption that sequences are interchangeable and therefore packing can be performed on the histogram of sequence lengths, rather than per sample. This transformation of the problem leads to algorithms which are fast and have linear complexity in dataset size. The shortest-pack-first histogram-packing (SPFHP) algorithm determines the packing order for the Wikipedia dataset of over 16M sequences in 0.02 seconds. The non-negative least-squares histogram-packing (NNLSHP) algorithm converges in 28.4 seconds but produces solutions which are more depth efficient, managing to get near optimal packing by combining a maximum of 3 sequences in one sample. Using the dataset with multiple sequences per sample requires additional masking in the attention layer and a modification of the MLM loss function. We demonstrate that both of these changes are straightforward to implement and have relatively little impact on the achievable performance gain on modern hardware. Finally, we pretrain BERT-Large using the packed dataset, demonstrating no loss of convergence and the desired 2x speed-up.

研究动机与目标

• 解决在序列长度为 512 的维基百科数据集上 BERT 预训练中因过度填充（超过 50%）导致的效率低下问题。
• 开发快速、确定性的打包算法，通过基于序列长度直方图的处理避免逐样本处理。
• 在保持模型收敛性和性能的前提下，实现接近最优的打包效率。
• 证明注意力掩码和针对打包序列修改的 MLM 损失在现代硬件上轻量且高效。

提出的方法

• 将打包问题从逐序列处理转化为基于直方图的优化，将复杂度降低至与数据集大小呈线性关系。
• 设计最短打包优先直方图打包（SPFHP）算法，优先处理较短序列，实现对 1600 万条序列的打包仅耗时 0.02 秒。
• 实现非负最小二乘直方图打包（NNLSHP）算法，以优化深度效率，支持每份打包样本最多合并三条序列。
• 在自注意力层中引入序列特定的掩码机制，防止来自不同原始序列的标记在打包样本中相互注意。
• 修改掩码语言建模（MLM）损失函数，以考虑部分标记因打包而被掩码，而非随机损坏。
• 使用打包后的数据集对 BERT-Large 进行预训练，以验证模型性能得以保持，且实现 2 倍的加速。

实验结果

研究问题

• RQ1 基于确定性直方图的打包方法能否在不牺牲模型收敛性的前提下，消除 BERT 预训练中超过 50% 的填充？
• RQ2 与逐样本方法相比，基于直方图的打包算法在运行时间和打包深度效率方面表现如何？
• RQ3 在打包的 BERT 训练中，注意力掩码和修改后的 MLM 损失对性能的影响如何？
• RQ4 能否在下游任务性能下降可忽略的前提下，实现每秒序列数 2 倍的加速？

主要发现

• 打包将序列长度为 512 时的填充比例降低至 50% 以下，从而在 BERT 预训练中实现每秒序列数 2 倍的加速。
• SPFHP 算法仅用 0.02 秒即可完成对整个 1600 万条序列的维基百科数据集打包，证明了其线性时间可扩展性。
• NNLSHP 算法实现了更高的打包深度效率，通过每份样本合并最多三条序列，接近最优打包效果。
• 修改后的注意力掩码和 MLM 损失函数对模型性能影响极小，确保了 BERT-Large 预训练过程中的收敛性。
• 使用打包数据集进行 BERT-Large 预训练时，其收敛行为与标准训练完全一致，验证了该方法的可行性。
• 该方法在现代硬件上具有实用性，打包算法与推理阶段的修改均易于实现。