Treinamento Eficiente de Modelos de Linguagem de Contexto Longo por meio da Desagregação da Atenção Central

Resumo

Apresentamos a técnica de desagregação de atenção central (CAD, do inglês Core Attention Disaggregation), que melhora o treinamento de modelos de linguagem de grande contexto ao desacoplar o cálculo da atenção central, softmax(QK^T)V, do restante do modelo e executá-lo em um pool separado de dispositivos. Nos sistemas existentes, a atenção central é colocalizada com outras camadas; em contextos longos, seu crescimento quadrático de computação em comparação com o crescimento quase linear de outros componentes causa desequilíbrio de carga e gargalos entre grupos de paralelismo de dados e pipeline. A CAD é viabilizada por duas observações. Primeiro, a atenção central é sem estado: ela não possui parâmetros treináveis e apenas dados transitórios mínimos, portanto, o equilíbrio se reduz ao agendamento de tarefas limitadas por computação. Segundo, ela é composável: kernels de atenção modernos mantêm alta eficiência ao processar lotes fundidos de fragmentos em nível de token com comprimentos arbitrários. A CAD particiona a atenção central em tarefas em nível de token e as distribui para servidores de atenção dedicados, que reúnem dinamicamente as tarefas para equalizar a computação sem sacrificar a eficiência do kernel. Implementamos a CAD em um sistema chamado DistCA, que utiliza um esquema de execução ping-pong para sobrepor completamente a comunicação com a computação e execução in-place nos servidores de atenção para reduzir o uso de memória. Em 512 GPUs H200 e comprimentos de contexto de até 512 mil tokens, o DistCA melhora a taxa de transferência de treinamento de ponta a ponta em até 1,35x, elimina gargalos de paralelismo de dados e pipeline e alcança um equilíbrio quase perfeito de computação e memória.

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.

Treinamento Eficiente de Modelos de Linguagem de Contexto Longo por meio da Desagregação da Atenção Central

Efficient Long-context Language Model Training by Core Attention Disaggregation

Resumo

Support