Addestramento Efficiente di Modelli Linguistici a Lungo Contesto tramite Disaggregazione dell'Attenzione Principale

Efficient Long-context Language Model Training by Core Attention Disaggregation

October 20, 2025

Autori: Yonghao Zhuang, Junda Chen, Bo Pang, Yi Gu, Yibo Zhu, Yimin Jiang, Ion Stoica, Eric Xing, Hao Zhang

cs.AI

Abstract

Presentiamo la disaggregazione dell'attenzione centrale (CAD, Core Attention Disaggregation), una tecnica che migliora l'addestramento di modelli linguistici di grandi dimensioni con contesti lunghi, disaccoppiando il calcolo dell'attenzione centrale, softmax(QK^T)V, dal resto del modello e eseguendolo su un pool separato di dispositivi. Nei sistemi esistenti, l'attenzione centrale è collocata insieme ad altri strati; con contesti lunghi, la sua crescita quadratica del calcolo rispetto alla crescita quasi lineare degli altri componenti causa squilibri di carico e ritardi nei gruppi paralleli di dati e pipeline. La CAD è resa possibile da due osservazioni. In primo luogo, l'attenzione centrale è senza stato: non ha parametri addestrabili e solo dati transitori minimi, quindi il bilanciamento si riduce alla pianificazione di task legati al calcolo. In secondo luogo, è componibile: i kernel di attenzione moderni mantengono un'elevata efficienza quando elaborano batch fusi di frammenti a livello di token con lunghezze arbitrarie. La CAD suddivide l'attenzione centrale in task a livello di token e li invia a server di attenzione dedicati, che raggruppano dinamicamente i task per bilanciare il calcolo senza sacrificare l'efficienza del kernel. Implementiamo la CAD in un sistema chiamato DistCA, che utilizza uno schema di esecuzione ping-pong per sovrapporre completamente la comunicazione con il calcolo e l'esecuzione in-place sui server di attenzione per ridurre l'uso della memoria. Su 512 GPU H200 e lunghezze di contesto fino a 512k token, DistCA migliora la produttività end-to-end dell'addestramento fino a 1,35x, elimina i ritardi nei gruppi paralleli di dati e pipeline e raggiunge un bilanciamento quasi perfetto del calcolo e della memoria.

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.