Эффективное обучение языковых моделей с длинным контекстом за счет разделения основного механизма внимания

Efficient Long-context Language Model Training by Core Attention Disaggregation

October 20, 2025

Авторы: Yonghao Zhuang, Junda Chen, Bo Pang, Yi Gu, Yibo Zhu, Yimin Jiang, Ion Stoica, Eric Xing, Hao Zhang

cs.AI

Аннотация

Мы представляем метод разделения основного механизма внимания (Core Attention Disaggregation, CAD), который улучшает обучение языковых моделей с длинным контекстом за счет отделения основного вычисления внимания, softmax(QK^T)V, от остальной части модели и выполнения его на отдельном пуле устройств. В существующих системах основной механизм внимания размещается вместе с другими слоями; при длинных контекстах его квадратичный рост вычислительной сложности по сравнению с почти линейным ростом других компонентов приводит к дисбалансу нагрузки и задержкам в группах данных и параллельных конвейеров. CAD основан на двух ключевых наблюдениях. Во-первых, основной механизм внимания не имеет состояния: он не содержит обучаемых параметров и лишь минимальные временные данные, поэтому балансировка сводится к планированию вычислительно-зависимых задач. Во-вторых, он композируем: современные ядра внимания сохраняют высокую эффективность при обработке объединенных пакетов сегментов токенов произвольной длины. CAD разделяет основной механизм внимания на задачи уровня токенов и распределяет их на выделенные серверы внимания, которые динамически перегруппировывают задачи для выравнивания вычислений без потери эффективности ядер. Мы реализовали CAD в системе под названием DistCA, которая использует схему выполнения "пинг-понг" для полного перекрытия коммуникаций и вычислений, а также выполнение на месте на серверах внимания для снижения использования памяти. На 512 GPU H200 и длинах контекста до 512k токенов DistCA повышает сквозную пропускную способность обучения до 1.35x, устраняет задержки в данных и параллельных конвейерах и достигает почти идеального баланса вычислений и памяти.

English

We present core attention disaggregation (CAD), a technique that improves long-context large language model training by decoupling the core attention computation, softmax(QK^T)V, from the rest of the model and executing it on a separate pool of devices. In existing systems, core attention is colocated with other layers; at long context lengths, its quadratic compute growth compared to the near-linear growth of other components causes load imbalance and stragglers across data and pipeline parallel groups. CAD is enabled by two observations. First, core attention is stateless: it has no trainable parameters and only minimal transient data, so balancing reduces to scheduling compute-bound tasks. Second, it is composable: modern attention kernels retain high efficiency when processing fused batches of token-level shards with arbitrary lengths. CAD partitions core attention into token-level tasks and dispatches them to dedicated attention servers, which dynamically rebatch tasks to equalize compute without sacrificing kernel efficiency. We implement CAD in a system called DistCA, which uses a ping-pong execution scheme to fully overlap communication with computation and in-place execution on attention servers to reduce memory use. On 512 H200 GPUs and context lengths up to 512k tokens, DistCA improves end-to-end training throughput by up to 1.35x, eliminates data and pipeline parallel stragglers, and achieves near-perfect compute and memory balance.