Entraînement efficace de modèles de langage à contexte long par désagrégation de l'attention principale

papers.abstract

Nous présentons la technique de désagrégation de l'attention centrale (CAD), qui améliore l'entraînement des modèles de langage à grand contexte en découplant le calcul de l'attention centrale, softmax(QK^T)V, du reste du modèle et en l'exécutant sur un pool distinct de dispositifs. Dans les systèmes existants, l'attention centrale est colocalisée avec d'autres couches ; pour des contextes longs, sa croissance quadratique en termes de calcul par rapport à la croissance quasi-linéaire des autres composants entraîne un déséquilibre de charge et des retardataires dans les groupes parallèles de données et de pipelines. CAD est rendu possible par deux observations. Premièrement, l'attention centrale est sans état : elle ne possède pas de paramètres entraînables et seulement des données transitoires minimales, ce qui réduit l'équilibrage à la planification de tâches liées au calcul. Deuxièmement, elle est composable : les noyaux d'attention modernes conservent une efficacité élevée lors du traitement de lots fusionnés de fragments au niveau des jetons avec des longueurs arbitraires. CAD partitionne l'attention centrale en tâches au niveau des jetons et les répartit sur des serveurs d'attention dédiés, qui réorganisent dynamiquement les tâches pour équilibrer le calcul sans sacrifier l'efficacité des noyaux. Nous implémentons CAD dans un système appelé DistCA, qui utilise un schéma d'exécution en ping-pong pour superposer entièrement la communication avec le calcul et une exécution sur place sur les serveurs d'attention pour réduire l'utilisation de la mémoire. Sur 512 GPU H200 et des longueurs de contexte allant jusqu'à 512k jetons, DistCA améliore le débit d'entraînement de bout en bout jusqu'à 1,35x, élimine les retardataires dans les parallélismes de données et de pipelines, et atteint un équilibre quasi parfait en termes de calcul et de mémoire.

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.

Entraînement efficace de modèles de langage à contexte long par désagrégation de l'attention principale

Efficient Long-context Language Model Training by Core Attention Disaggregation

papers.abstract

Support