Quand la Précision Rencontre la Position : BFloat16 Démonte RoPE dans l'Entraînement à Long Contexte

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

Résumé

L'extension des tailles de fenêtre de contexte permet aux grands modèles de langage (LLM) de traiter des séquences plus longues et de gérer des tâches plus complexes. L'incorporation de position rotative (RoPE) est devenue la norme de facto en raison de ses propriétés d'encodage de position relative bénéfiques pour l'entraînement à long contexte. Cependant, nous observons que l'utilisation de RoPE avec le format BFloat16 entraîne des problèmes numériques, le faisant dévier de son encodage de position relative prévu, en particulier dans des scénarios de long contexte. Ce problème découle de la précision limitée du BFloat16 et s'accumule à mesure que la longueur du contexte augmente, le premier jeton contribuant de manière significative à ce problème. Pour remédier à cela, nous avons développé AnchorAttention, une méthode d'attention plug-and-play qui atténue les problèmes numériques causés par le BFloat16, améliore les capacités de long contexte et accélère l'entraînement. AnchorAttention réduit les calculs d'attention inutiles, maintient la cohérence sémantique et améliore l'efficacité computationnelle en traitant le premier jeton comme une ancre partagée avec un ID de position cohérent, le rendant visible à tous les documents dans le contexte d'entraînement. Des expériences sur trois types de LLM montrent qu'AnchorAttention améliore significativement les performances en long contexte et réduit le temps d'entraînement de plus de 50\% par rapport aux mécanismes d'attention complets standard, tout en préservant les capacités originales du LLM sur des tâches générales. Notre code est disponible sur 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.