[论文解读] Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance

Mario Michael Krell, Matej Kosec+2 more|arXiv (Cornell University)|Jun 29, 2021

Natural Language Processing Techniques被引用 6

一句话总结

本文提出了一种新颖的序列打包形式化方法，将其建模为装箱问题，以消除大语言模型训练中的填充标记，实现高达2倍的加速且不损失性能。通过使用高效且数学等价的打包算法，并修改注意力机制以防止跨序列污染，该方法在标准硬件上实现了高吞吐量训练，且无精度下降。

ABSTRACT

Effective training of today's large language models (LLMs) depends on large batches and long sequences for throughput and accuracy. To handle variable-length sequences on hardware accelerators, it is common practice to introduce padding tokens, so that all sequences in a batch have the same length. We show in this paper that the variation in sequence lengths in common NLP datasets is such that up to 50% of all tokens can be padding. In less common, but not extreme, cases (e.g. GLUE-cola with sequence length 128), the ratio is up to 89%. Existing methods to address the resulting inefficiency are complicated by the need to avoid cross-contamination in self-attention, by a reduction in accuracy when sequence ordering information is lost, or by customized kernel implementations only valid for specific accelerators. This paper introduces a new formalization of sequence packing in the context of the well-studied bin packing problem, and presents new algorithms based on this formulation which, for example, confer a 2x speedup for phase 2 pre-training in BERT. We show how existing models can be adapted to ensure mathematical equivalence between the original and packed models, meaning that packed models can be trained with existing pre-training and fine-tuning practices.

研究动机与目标

• 解决大语言模型可变长度序列批处理中因填充标记导致的效率低下问题。
• 克服现有方法存在的跨序列污染、精度损失或依赖硬件优化的局限性。
• 开发一种通用的、与加速器无关的序列打包框架，兼容标准训练与微调工作流。
• 确保打包模型与未打包模型在数学上等价，以保持收敛行为和模型性能。
• 实现接近最优的打包效率，同时避免使用分隔标记或复杂内核实现。

提出的方法

• 通过新颖的数学公式将序列打包形式化为装箱问题，以最小化填充标记。
• 提出两种基于成熟求解器（如非负最小二乘法）的确定性、高效算法，用于计算最优打包策略。
• 采用带残差惩罚的加权线性规划公式，以最小化填充标记并处理剩余序列。
• 修改自注意力机制，防止打包批处理中不同序列之间的跨序列污染，确保注意力仅在序列内部计算。
• 引入基于残差的校正步骤，通过动态添加新的打包策略来处理剩余序列。
• 通过调整注意力掩码和位置嵌入，确保模型等价性，以保持打包序列中的注意力模式。

实验结果

研究问题

• RQ1 能否将序列打包形式化为装箱问题，以最小化填充标记并保持模型性能？
• RQ2 依赖分隔标记或排序批处理的现有打包方法在性能与精度之间存在何种权衡？
• RQ3 能否设计一种既高效又兼容标准训练流水线、无需硬件特定优化的打包策略？
• RQ4 如何在不引入精度损失的前提下，防止打包序列中自注意力机制的跨序列污染？
• RQ5 所提出方法在多大程度上可提升训练吞吐量，同时保持收敛行为和模型精度？

主要发现

• 在序列长度为512时，Wikipedia预训练数据集中高达50%的标记为填充标记；在GLUE-cola数据集、长度为128时，填充标记占比高达89%。
• 所提出的打包方法通过消除填充标记，使BERT在第二阶段预训练中实现高达2倍的理论加速。
• BERT-large在打包数据集上的收敛行为与未打包数据集在数学上完全等价，性能无任何下降。
• 该方法将填充标记减少至接近零水平，在Wikipedia数据集上实现了98.7%的打包效率。
• 对于包含数百万条序列的数据集，该算法在不到一秒内即可计算出最优打包策略，支持实时预处理。
• 该方法与硬件无关，兼容标准训练与微调工作流，无需模型重训练或自定义内核。