Efficiënte Training van Taalmodellen met Lange Context door Ontbinding van Kern-Attentie

Samenvatting

We presenteren kern-aandacht-disaggregatie (CAD), een techniek die de training van grote taalmodellen met lange contexten verbetert door de kern-aandachtsberekening, softmax(QK^T)V, los te koppelen van de rest van het model en deze uit te voeren op een aparte pool van apparaten. In bestaande systemen bevindt de kern-aandacht zich op dezelfde locatie als andere lagen; bij lange contextlengtes veroorzaakt de kwadratische rekenkundige groei in vergelijking met de bijna-lineaire groei van andere componenten een onbalans in de belasting en vertragingen in data- en pijplijn-parallelle groepen. CAD wordt mogelijk gemaakt door twee observaties. Ten eerste is kern-aandacht toestandsloos: het heeft geen trainbare parameters en slechts minimale tijdelijke data, waardoor balansverstoring neerkomt op het plannen van rekenintensieve taken. Ten tweede is het samenstelbaar: moderne aandachtskernels behouden een hoge efficiëntie bij het verwerken van samengevoegde batches van token-level shards met willekeurige lengtes. CAD verdeelt de kern-aandacht in token-level taken en wijst deze toe aan toegewijde aandachtsservers, die taken dynamisch hergroeperen om de rekenkundige belasting gelijk te verdelen zonder de kernefficiëntie op te offeren. We implementeren CAD in een systeem genaamd DistCA, dat een ping-pong-uitvoeringsschema gebruikt om communicatie volledig te overlappen met berekening en in-place uitvoering op aandachtsservers om het geheugengebruik te verminderen. Op 512 H200 GPU's en contextlengtes tot 512k tokens verbetert DistCA de end-to-end trainingsdoorvoer met maximaal 1,35x, elimineert het vertragingen in data- en pijplijn-parallelle groepen, en bereikt het een bijna perfecte balans in rekenkracht en geheugen.

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.

Efficiënte Training van Taalmodellen met Lange Context door Ontbinding van Kern-Attentie

Efficient Long-context Language Model Training by Core Attention Disaggregation

Samenvatting

Support