Когда точность встречает позицию: BFloat16 разрушает RoPE в обучении с длинным контекстом.

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

November 20, 2024

Авторы: Haonan Wang, Qian Liu, Chao Du, Tongyao Zhu, Cunxiao Du, Kenji Kawaguchi, Tianyu Pang

cs.AI

Аннотация

Увеличение размеров окна контекста позволяет крупным языковым моделям (LLM) обрабатывать более длинные последовательности и решать более сложные задачи. Поворотное позиционное вложение (RoPE) стало фактическим стандартом благодаря своим относительным свойствам кодирования позиций, которые полезны для обучения на длинных контекстах. Однако мы замечаем, что использование RoPE с форматом BFloat16 приводит к числовым проблемам, вызывая отклонение от заданного относительного позиционного кодирования, особенно в сценариях с длинным контекстом. Эта проблема возникает из-за ограниченной точности BFloat16 и накапливается с увеличением длины контекста, причем первый токен значительно способствует этой проблеме. Для решения этой проблемы мы разработали AnchorAttention, метод внимания, который смягчает числовые проблемы, вызванные BFloat16, улучшает возможности работы с длинным контекстом и ускоряет обучение. AnchorAttention снижает ненужные вычисления внимания, сохраняет семантическую связность и повышает вычислительную эффективность, рассматривая первый токен как общий якорь с постоянным идентификатором позиции, делая его видимым для всех документов в пределах обучающего контекста. Эксперименты на трех типах LLM показывают, что AnchorAttention значительно улучшает производительность на длинных контекстах и сокращает время обучения более чем на 50\% по сравнению со стандартными механизмами полного внимания, сохраняя при этом возможности исходной LLM по общим задачам. Наш код доступен по адресу 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.