Quando la Precisione Incontra la Posizione: BFloat16 Scompone RoPE nell'Addestramento a Lungo Contesto

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

November 20, 2024

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

cs.AI

Abstract

L'estensione delle dimensioni della finestra di contesto consente ai grandi modelli linguistici (LLM) di elaborare sequenze più lunghe e gestire compiti più complessi. L'Incorporamento Posizionale Rotativo (RoPE) è diventato lo standard de facto grazie alle sue proprietà di codifica posizionale relativa che favoriscono l'addestramento a lungo contesto. Tuttavia, osserviamo che l'utilizzo di RoPE con il formato BFloat16 comporta problemi numerici, causando una deviazione dalla sua codifica posizionale relativa prevista, specialmente in scenari di lungo contesto. Questo problema deriva dalla limitata precisione di BFloat16 e si accumula all'aumentare della lunghezza del contesto, con il primo token che contribuisce significativamente a questo problema. Per affrontare ciò, sviluppiamo AnchorAttention, un metodo di attenzione plug-and-play che allevia i problemi numerici causati da BFloat16, migliora le capacità di lungo contesto e accelera l'addestramento. AnchorAttention riduce i calcoli di attenzione non necessari, mantiene la coerenza semantica e aumenta l'efficienza computazionale trattando il primo token come un ancoraggio condiviso con un ID di posizione costante, rendendolo visibile a tutti i documenti all'interno del contesto di addestramento. Gli esperimenti su tre tipi di LLM dimostrano che AnchorAttention migliora significativamente le prestazioni a lungo contesto e riduce il tempo di addestramento di oltre il 50\% rispetto ai meccanismi di attenzione completa standard, preservando nel contempo le capacità originali del LLM su compiti generali. Il nostro codice è disponibile su 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.