条纹注意力：面向因果 Transformer 的更快 Ring Attention
机器学习
2023-11-17 v1 计算与语言
摘要
为帮助满足 transformer 模型对更长序列长度日益增长的需求，Liu 等人近期提出了 Ring Attention，一种能够将自注意力分布到多个设备上从而克服单设备内存瓶颈的精确注意力算法。在本文中，我们研究了在因果 transformer 模型这一重要特例下 Ring Attention 的性能特征，并识别出因因果注意力计算的三角结构所导致的关键工作负载不平衡。我们提出了对 Ring Attention 的一个简单扩展，称之为 Striped Attention，以修正该不平衡。不同于各设备拥有连续子序列，每个设备拥有均匀分布于整个序列中的令牌子集，我们证明这带来了更均衡的工作负载。在 A100 GPU 和 TPUv4 上运行 Striped Attention 的实验中，在 256k 序列长度的因果 transformer 训练上，我们相比原始 Ring Attention 算法实现了最高 1.45 倍的端到端吞吐提升。此外，在 16 块 TPUv4 芯片上，我们在 786k 序列长度下实现了 1.65 倍的加速。我们将实验代码开源。
引用
@article{arxiv.2311.09431,
title = {Striped Attention: Faster Ring Attention for Causal Transformers},
author = {William Brandon and Aniruddha Nrusimha and Kevin Qian and Zachary Ankner and Tian Jin and Zhiye Song and Jonathan Ragan-Kelley},
journal= {arXiv preprint arXiv:2311.09431},
year = {2023}
}