Wanneer Precisie Positie Ontmoet: BFloat16 Doorbreekt RoPE in Training met Lange Contexten

When Precision Meets Position: BFloat16 Breaks Down RoPE in Long-Context Training

November 20, 2024

Auteurs: Haonan Wang, Qian Liu, Chao Du, Tongyao Zhu, Cunxiao Du, Kenji Kawaguchi, Tianyu Pang

cs.AI

Samenvatting

Het vergroten van de contextvenstergroottes stelt grote taalmodellen (LLM's) in staat om langere sequenties te verwerken en complexere taken aan te kunnen. Rotary Positional Embedding (RoPE) is de facto standaard geworden vanwege zijn relatieve positionele coderingseigenschappen die gunstig zijn voor training met een lange context. We merken echter op dat het gebruik van RoPE met het BFloat16-formaat leidt tot numerieke problemen, waardoor het afwijkt van zijn beoogde relatieve positionele codering, vooral in scenario's met een lange context. Dit probleem ontstaat door de beperkte precisie van BFloat16 en stapelt zich op naarmate de contextlengte toeneemt, waarbij de eerste token aanzienlijk bijdraagt aan dit probleem. Om dit aan te pakken, hebben we AnchorAttention ontwikkeld, een plug-and-play aandachtsmethode die numerieke problemen veroorzaakt door BFloat16 verlicht, de mogelijkheden voor lange context verbetert en de training versnelt. AnchorAttention vermindert onnodige aandachtsberekeningen, behoudt semantische coherentie en verbetert de rekenkundige efficiëntie door de eerste token te behandelen als een gedeelde anker met een consistente positie-ID, waardoor het zichtbaar is voor alle documenten binnen de trainingscontext. Experimenten met drie soorten LLM's tonen aan dat AnchorAttention de prestaties op lange context aanzienlijk verbetert en de trainingsduur met meer dan 50% verkort in vergelijking met standaard volledige aandachtsmechanismen, terwijl de oorspronkelijke mogelijkheden van het LLM behouden blijven voor algemene taken. Onze code is beschikbaar op https://github.com/haonan3/AnchorContext.

English

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.