arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2408.04093cs.LGcs.CL
收藏

树注意力:GPU集群上长上下文注意力的拓扑感知解码

Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters

Zyphra(齐夫拉公司);EleutherAI(EleutherAI 研究机构)

Vasudev Shyam, Jonathan Pilault, Emily Shepperd, Quentin Anthony, Beren Millidge

AI总结:

本文提出Tree Attention算法,通过树形归约在多个GPU上并行化精确注意力计算,相比Ring Attention解码速度最高提升8倍,通信量和峰值内存显著降低,并在Llama 3.1-8B及多种硬件上验证了效果。

摘要:

我们的公式表明,沿序列轴的归约可以通过树形归约高效并行计算。我们提出的算法称为Tree Attention,用于在多个GPU上并行化精确注意力计算,使跨设备解码的渐近速度显著快于Ring Attention等最先进方法(实验中最高快8倍),同时通信量大幅减少,峰值内存占用降低2倍。我们证明,Tree Attention在Llama 3.1-8B上将解码速度提升最高4倍,并可应用于多种硬件和网络配置,如H100 DGX节点、AMD MI300x节点以及PCIe连接的NVIDIA RTX 4090。我们的代码已公开:https://github.com/Zyphra/tree_attention

英文摘要:

Our formulation reveals that the reduction across the sequence axis can be efficiently computed in parallel through a tree reduction. Our algorithm, called Tree Attention, for parallelizing exact attention computation across multiple GPUs enables cross-device decoding to be performed asymptotically faster (up to 8x faster in our experiments) than state-of-the-art approaches such as Ring Attention, while also requiring significantly less communication volume and incurring 2x less peak memory. We demonstrate that Tree Attention speeds up decoding up to 4x on Llama 3.1-8B and can be applied to a variety of hardware and networking setups such as H100 DGX nodes, AMD MI300x nodes, and PCIe connected NVIDIA RTX 4090s. Our code is publicly available here: https://github.com/Zyphra/tree_attention

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑