普林斯顿：高效训练长文本LLM实践

📖标题：How to Train Long-Context Language Models (Effectively)
🌐来源：arXiv, 2410.02660

摘要

🔸我们研究了语言模型（LM）的持续训练和监督微调（SFT），以有效利用长上下文信息。我们首先建立了一个可靠的评估协议来指导模型开发——我们使用了一组广泛的长上下文任务，而不是困惑或简单的NIAH测试，我们在SFT后用指令数据评估模型，因为这更好地揭示了长上下文能力。在我们稳健评估的支持下，我们进行了彻底的实验，以确定继续预训练的数据组合、指令调优数据集和许多其他设计选择。
🔸我们发现：（1）代码库和书籍是很好的长数据来源，但将它们与高质量的短数据结合起来至关重要；（2） 序列长度超过评估长度的训练可以提高长上下文性能；（3） 对于SFT，仅使用短指令数据集在长上下文任务上具有很强的性能。我们的最终模型ProLong-8B由Llama-3初始化，并在40B令牌上训练，在长度为128K的类似大小的模型中展示了最先进的长上下文性能。ProLong在大多数长上下文任务上的表现优于Llama-3.1-8B-Instruct，尽管在长上下文训练中只看到了5%的令牌。此外，ProLong可以有效地处理多达512K个令牌，这是公开可用的LM中最长的上下文窗口之一。已开源在https://github.com/princeton-nlp/ProLong

🛎️文章简介

🔸研究问题：如何有效地训练能够处理极长上下文（如512K）的语言模型？
🔸主要贡献：论文提出了ProLong模型，该模型在128K上下文长度下表现优异，并且具有512K的最大上下文长度。

📝重点思路

🔺相关工作

🔸调整现有的LM以适应长上下文：许多作品探索通过位置外推法以最少的训练扩展LM上下文窗口，Llama 3.1在标准预训练和监督微调之间添加长上下文持续训练阶段来实现长上下文能力。
🔸高效的长上下文架构：如线性注意力以及替代注意力架构，通常需要从头开始训练，并且许多在长上下文回忆方面具有固有的局限性，最近的工作探索了混合模型。
🔸长上下文评估：已经提出了许多用于长上下文评估的基准，探索位置偏差、上下文学习、和整本书的摘要等。

🔺论文方案

🔸数据选择与混合：论文从直觉出发，认为数据应包含长文档和短文档的混合。长数据由64K长度的单文档块组成，而短数据则通过打包文档直到每个序列达到64K来构建。
🔸监督微调：研究如何最好地使长上下文语言模型遵循指令，专注于指令数据集上的监督微调，并将强化学习和偏好优化留作未来工作。
🔸训练配方：总结了ProLong的训练配方，包括从Llama-3-8B-Instruct模型开始，在64K序列长度上训练20B的token，然后在512K序列长度上进一步训练20B的token，最终通过SFT在UltraChat上获得ProLong模型。
🔸数据处理与训练堆栈：使用Mosaic-streaming加载和混合数据，使用FlashAttention-2进行高效的注意力实现，并基于DeepSpeed-Ulysses实现序列并行。

🔎分析总结

🔸ProLong性能：在长上下文评估中优于所有10B规模的模型，特别是在除总结外的所有类别中优于Llama-3.1-8B-Instruct。
🔸数据混合的影响：使用标准短上下文指令数据在当前设置中达到最佳长上下文结果。
🔸长上下文数据的重要性：仅使用0.5B的tokens做领域平衡、长度上采样的数据不足以获得长上下文召回能力，需要更全面的数据处理。
🔸训练效率：通过智能批处理算法和变长注意力实现，提高了长上下文训练的速度和效率。

💡个人观点

论文提出了ProLong模型，并分析了数据选择、配比和训练方法的重要性。

附录