当精度遇上位置：BFloat16 在长上下文训练中破坏 RoPE
计算与语言
2024-11-27 v2
摘要
扩展上下文窗口大小使得大语言模型（LLMs）能够处理更长的序列并应对更复杂的任务。旋转位置编码（RoPE）因其有利于长上下文训练的相对位置编码特性而成为事实上的标准。然而，我们观察到在 BFloat16 格式下使用 RoPE 会导致数值问题，使其偏离预期的相对位置编码，在长上下文场景中尤为明显。该问题源于 BFloat16 的有限精度，并随上下文长度的增加而累积，其中首个 token 对此问题贡献显著。为解决这一问题，我们开发了 AnchorAttention，一种即插即用的注意力方法，可缓解 BFloat16 引起的数值问题，提升长上下文能力并加速训练。AnchorAttention 通过将首个 token 视为具有一致位置 ID 的共享锚点，使其对训练上下文中的所有文档可见，从而减少不必要的注意力计算，保持语义连贯性，并提升计算效率。在三种类型的 LLMs 上的实验表明，与标准的全注意力机制相比，AnchorAttention 显著提升了长上下文性能，并将训练时间减少了 50% 以上，同时保留了原始 LLM 在通用任务上的能力。我们的代码可在 https://github.com/haonan3/AnchorContext 获取。
引用
@article{arxiv.2411.13476,
title = {When Precision Meets Position: BFloat16 Breaks Down RoPE in Long-Context Training},
author = {Haonan Wang and Qian Liu and Chao Du and Tongyao Zhu and Cunxiao Du and Kenji Kawaguchi and Tianyu Pang},
journal= {arXiv preprint arXiv:2411.13476},
year = {2024}
}