Wenn Präzision auf Position trifft: BFloat16 überwindet RoPE im Training mit langem Kontext.

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

November 20, 2024

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

cs.AI

Zusammenfassung

Die Erweiterung der Kontextfenstergrößen ermöglicht es großen Sprachmodellen (LLMs), längere Sequenzen zu verarbeiten und komplexere Aufgaben zu bewältigen. Die Rotierende Positionale Einbettung (RoPE) hat sich aufgrund ihrer relativen Positionscodierungseigenschaften, die das Training mit langem Kontext begünstigen, als de facto Standard etabliert. Wir beobachten jedoch, dass die Verwendung von RoPE im BFloat16-Format zu numerischen Problemen führt, die dazu führen, dass es von seiner beabsichtigten relativen Positionscodierung abweicht, insbesondere in Szenarien mit langem Kontext. Dieses Problem entsteht durch die begrenzte Genauigkeit von BFloat16 und häuft sich mit zunehmender Kontextlänge an, wobei das erste Token maßgeblich zu diesem Problem beiträgt. Um dies zu lösen, haben wir AnchorAttention entwickelt, eine Plug-and-Play-Aufmerksamkeitsmethode, die numerische Probleme, die durch BFloat16 verursacht werden, mildert, die Fähigkeiten im langen Kontext verbessert und das Training beschleunigt. AnchorAttention reduziert unnötige Aufmerksamkeitsberechnungen, bewahrt semantische Kohärenz und steigert die Rechenleistung, indem das erste Token als gemeinsamer Anker mit einer konsistenten Positions-ID behandelt wird, sodass es für alle Dokumente innerhalb des Trainingskontexts sichtbar ist. Experimente mit drei Arten von LLMs zeigen, dass AnchorAttention die Leistung im langen Kontext signifikant verbessert und die Trainingszeit im Vergleich zu Standard-Voll-Aufmerksamkeitsmechanismen um über 50\% reduziert, während die ursprünglichen Fähigkeiten des LLMs bei allgemeinen Aufgaben erhalten bleiben. Unser Code ist unter https://github.com/haonan3/AnchorContext verfügbar.

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.