通过拆分注意力计算,实现长文本模型高效训练。
Efficient Long-context Language Model Training by Core Attention Disaggregation
- 将核心注意力计算独立部署到专用设备上调度执行
- 在512K上下文长度下训练吞吐提升1.35倍,消除计算瓶颈
- 适合大规模长文本模型训练,尤其对算力不均衡场景有效
我们提出核心注意力拆分(CAD),通过将核心注意力计算(softmax(QK^T)V)从模型其他部分解耦,并在独立设备池中执行,提升长上下文大语言模型的训练效率。现有系统中,核心注意力与其它层共置,其二次方计算增长远超其他组件的近线性增长,导致数据和流水并行组间负载失衡与延迟。CAD基于两个观察:第一,核心注意力无状态,无可训练参数且仅需少量临时数据,负载平衡可简化为调度计算密集型任务;第二,现代注意力核在处理任意长度的分片批时仍保持高效率。CAD将核心注意力拆分为令牌级任务,分发至专用注意力服务器,动态重批以均衡计算而不牺牲内核效率。我们实现了名为DistCA的系统,采用乒乓执行模式完全重叠通信与计算,并在注意力服务器上就地执行以减少内存占用。在512个H200 GPU上,针对最长512k令牌的上下文,DistCA将端到端训练吞吐提升最高达1.35倍,消除数据与流水并行中的慢节点,实现近乎完美的计算与内存平衡。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。