When Precision Meets Position: BFloat16 Breaks Down RoPE in Long-Context
  Training
Authors: name: Haonan Wang; name: Qian Liu; name: Chao Du; name: Tongyao Zhu; name: Cunxiao Du; name: Kenji Kawaguchi; name: Tianyu Pang
Published: 2024-11-20 17:22:31

Abstract: Extending context window sizes allows large language models (LLMs) to process
longer sequences and handle more complex tasks. Rotary Positional Embedding
(RoPE) has become the de facto standard due to its relative positional encoding
properties that benefit long-context training. However, we observe that using
RoPE with BFloat16 format results in numerical issues, causing it to deviate
from its intended relative positional encoding, especially in long-context
scenarios. This issue arises from BFloat16's limited precision and accumulates
as context length increases, with the first token contributing significantly to
this problem. To address this, we develop AnchorAttention, a plug-and-play
attention method that alleviates numerical issues caused by BFloat16, improves
long-context capabilities, and speeds up training. AnchorAttention reduces
unnecessary attention computations, maintains semantic coherence, and boosts
computational efficiency by treating the first token as a shared anchor with a
consistent position ID, making it visible to all documents within the training
context. Experiments on three types of LLMs demonstrate that AnchorAttention
significantly improves long-context performance and reduces training time by
over 50\% compared to standard full attention mechanisms, while preserving the
original LLM's capabilities on general tasks. Our code is available at
https://github.com/haonan3/AnchorContext.