30行代码,500万长文本推理提速8倍!「树注意力」让GPU越多省的越多

人工智能 新闻
环注意力(Ring Attention)后继者——树注意力(Tree Attention)来了。

跨GPU的注意力并行,最高提速8倍,支持512万序列长度推理。

环注意力(Ring Attention)后继者——树注意力(Tree Attention)来了。

图片

最关键之处在于,通信步数随设备数量成对数增长,而不是线性增长

换句话说,树注意力的优势随着设备数量增大会更加明显。实验中,在128卡、512万序列长度设置时达到最高8倍加速。

图片

与环注意力相比,峰值内存占用也能节省不少

图片

相关代码已经开源,基于谷歌jax框架,已和Flash Attention整合,实现起来只需要30行代码

图片

论文一公布,就被业界评价为“对高推理需求的大型公司很重要”。

图片

这下和黄仁勋的GPU“买的越多,省的越多”论对上了,英伟达再次赢麻。

图片

注意力机制的能量视角

首先简单回顾一下这次被拿来对比的环注意力,由UC伯克利大牛Pieter Abeel团队提出。

环注意力被认为是让上一波大模型纷纷扩展到百万上下文的关键,从谷歌Gemini 1.5到后来的Llama 3.1系列都用了它的某种变体。

图片

简单来说,环注意力的核心思想是将长序列分成多个Block,每个GPU处理一个。在拓扑意义上相当于所有GPU排成一个圆环,将Key-Value信息传下去,同时从上一个GPU接收信息。

只要保证计算时间比数据传输时间长,这个过程就不会造成额外开销。

同时与之前的近似方法不同,环注意力不会损失精度,保持了完整的注意力计算。

图片

最新的树注意力,在分块计算、跨设备并行、保持精度特性的基础上,提出了一种自注意力的能量函数,通过计算梯度利用树形拓扑优化多GPU间的通信

传统上,人们把注意力看作Query向量与Key向量的相似度匹配,再对Value向量做加权求和。

树注意力团队在Hopfield网络等基于能量的模型相关研究基础上,将注意力解释为一个能量函数对某变量的梯度。

存在一个标量能量函数F,它依赖于Key、Query、Value以及一个辅助变量ζ,而注意力的结果恰好等于F对ζ的梯度在ζ=0处的值。

图片

结合自动微分等技术,从能量和梯度的视角看待自注意力,暗示了只要能高效计算F就能高效计算自注意力

具体到语言模型中基于KV缓存的解码,能量函数可以表示成:

图片

由于logsumexp和max运算操作都满足结合律,可以按任意顺序进行,而不会影响最终结果。

图片

在此前提下,团队设计了新的并行化算法,先在各GPU上并行计算局部能量函数,再通过树状的Allreduce汇总各处结果,最后用自动微分取梯度,即可得到注意力的输出。

全过程仅需与计算能量函数相同的时间开销,而显存占用也几乎没有额外负担。

树注意力在设计上还充分利用了GPU集群的两级拓扑特点——即同节点内使用高速NVLink,而节点间则依赖IB或以太网等。

相比之下,环形注意力天然不适应这种拓扑,难以将通信与计算很好地重叠,终会被最慢的互联带宽所制约。

最后值得一提的是,虽然理论上单GPU内部也可用类似策略提速,但当前硬件的流式处理器(SM)间通信还是共享内存,优势并不明显。

不过,英伟达在H100上实验性地支持了SM间点对点的指令,这为未来单卡注意力优化带来了新的想象空间。

最被低估的AI实验室之一

树注意力团队主要成员来自Zyphra,一家新兴的AI创业公司,被评价为“当前最被低估的AI实验室之一”

图片

Zyphra重点关注边缘AI、端侧AI,曾发布基于Mamba架构的基础模型Zamba

创始人Krithik Puthalath以及树注意力共同一作Vasudev Shyam、Jonathan Pilault都有数学和理论物理学术背景

图片

论文地址:https://arxiv.org/abs/2408.04093

责任编辑:张燕妮 来源: 量子位
相关推荐

2024-07-01 12:17:54

2023-10-14 15:22:22

2023-02-14 09:00:26

2023-12-11 15:40:32

PyTorch代码大模型

2024-11-08 15:00:00

AI模型

2023-11-13 18:19:54

模型训练

2018-03-28 14:10:10

GoPython代码

2020-01-21 22:23:18

代码开发AI

2021-12-31 13:56:42

Transformer数据人工智能

2024-02-01 12:43:16

模型数据

2018-02-07 08:08:11

2017-10-02 08:09:04

2020-09-07 06:59:44

Kafka分布式场景

2022-03-09 14:57:53

Numbapython

2023-11-24 12:36:00

模型训练

2024-06-03 10:56:53

2024-09-19 10:07:41

2023-09-01 18:18:32

2018-07-27 09:32:18

Python代码数据

2016-10-08 16:02:37

WIFIMegaMIMO系统
点赞
收藏

51CTO技术栈公众号