Quando Precisão Encontra Posição: BFloat16 Derruba RoPE no Treinamento de Longo Contexto

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

November 20, 2024

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

cs.AI

Resumo

Ampliar o tamanho da janela de contexto permite que grandes modelos de linguagem (LLMs) processem sequências mais longas e lidem com tarefas mais complexas. A Incorporação Posicional Rotativa (RoPE) tornou-se o padrão de facto devido às suas propriedades de codificação posicional relativa que beneficiam o treinamento de longo contexto. No entanto, observamos que o uso do RoPE com o formato BFloat16 resulta em problemas numéricos, fazendo com que se desvie de sua codificação posicional relativa pretendida, especialmente em cenários de longo contexto. Esse problema surge da precisão limitada do BFloat16 e se acumula à medida que o comprimento do contexto aumenta, com o primeiro token contribuindo significativamente para esse problema. Para resolver isso, desenvolvemos o AnchorAttention, um método de atenção plug-and-play que alivia os problemas numéricos causados pelo BFloat16, melhora as capacidades de longo contexto e acelera o treinamento. O AnchorAttention reduz cálculos de atenção desnecessários, mantém a coerência semântica e aumenta a eficiência computacional tratando o primeiro token como uma âncora compartilhada com um ID de posição consistente, tornando-o visível para todos os documentos dentro do contexto de treinamento. Experimentos em três tipos de LLMs demonstram que o AnchorAttention melhora significativamente o desempenho de longo contexto e reduz o tempo de treinamento em mais de 50\% em comparação com mecanismos de atenção completos padrão, preservando as capacidades originais do LLM em tarefas gerais. Nosso código está disponível em 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.