Cuando la Precisión se Encuentra con la Posición: BFloat16 Descompone RoPE en el Entrenamiento de Contexto Largo

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

Resumen

Ampliar el tamaño de la ventana de contexto permite a los modelos de lenguaje grandes (LLMs) procesar secuencias más largas y abordar tareas más complejas. La Incrustación Posicional Rotativa (RoPE) se ha convertido en el estándar de facto debido a sus propiedades de codificación posicional relativa que benefician el entrenamiento con un contexto largo. Sin embargo, observamos que al utilizar RoPE con el formato BFloat16 se producen problemas numéricos, lo que provoca que se desvíe de su codificación posicional relativa prevista, especialmente en escenarios de largo contexto. Este problema surge de la precisión limitada de BFloat16 y se acumula a medida que aumenta la longitud del contexto, siendo el primer token el que contribuye significativamente a este problema. Para abordar esto, desarrollamos AnchorAttention, un método de atención plug-and-play que alivia los problemas numéricos causados por BFloat16, mejora las capacidades de largo contexto y acelera el entrenamiento. AnchorAttention reduce cálculos de atención innecesarios, mantiene la coherencia semántica y aumenta la eficiencia computacional al tratar el primer token como un anclaje compartido con un ID de posición consistente, haciéndolo visible para todos los documentos dentro del contexto de entrenamiento. Experimentos con tres tipos de LLMs demuestran que AnchorAttention mejora significativamente el rendimiento en contextos largos y reduce el tiempo de entrenamiento en más del 50\% en comparación con los mecanismos estándar de atención completa, preservando al mismo tiempo las capacidades originales del LLM en tareas generales. Nuestro código está disponible en 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.