长上下文注意力基准:从内核效率到分布式上下文并行

发表
Jingwei XuJingwei Xu 提交
作者: Tao Bu, Qiangang Wang, Bowen Zeng, Hanwen Sun, Yunpeng Huang, Chun Cao, Jingwei XuJingwei Xu

摘要

AI 生成总结
统一基准评估了基于 Transformer 的大型语言模型中的注意力机制,重点关注不同注意力掩码模式和序列长度下的效率、可扩展性和性能。
基于 Transformer 的大型语言模型 (LLM) 取得了显著成功,但其标准注意力机制在序列长度方面会产生二次方的计算和内存成本,这对长上下文训练构成了主要瓶颈。以前的工作沿着两个方向解决了这一挑战:(1) 内核级优化,加速密集和稀疏注意力操作;(2) 模块级策略,通常称为分布式注意力或上下文并行训练,它在多个设备上扩展注意力。然而,系统评估仍然有限:操作级比较通常不完整,而上下文并行策略通常是框架特定的,在不同上下文下的性能分析不明确。为了解决这些空白,我们提出了一个统一的基准,它集成了具有模块化和可扩展评估接口的代表性注意力内核和上下文并行机制。该基准从两个关键维度评估方法:(1) 注意力掩码模式,这强烈影响效率、可伸缩性和可用性;(2) 序列长度和分布式规模,这决定了极端长上下文训练下的性能。通过在多达 96 个 GPU 集群上的全面实验,我们的基准实现了可重现的比较,突出了特定方法的权衡,并为长上下文 LLM 训练中注意力机制的设计和部署提供了实用指导。
查看 arXiv 页面查看 PDF

评论

Jingwei XuJingwei Xu
论文作者
论文提交者

测试