讯享网

 <p><blockquote id="2UCAEKKA">梦晨 发自 凹非寺<br/>量子位 | 公众号 QbitAI<br/></blockquote></p><p id="2UCAEKIS">跨GPU的注意力并行，<strong>最高提速8倍</strong>，支持<strong>512万序列长度</strong>推理。</p><p id="2UCAEKIT">环注意力（Ring Attention）后继者——<strong>树注意力</strong>（Tree Attention）来了。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2Fa3d56781j00si38gc000od000hs003wm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKIV">最关键之处在于，<strong>通信步数随设备数量成对数增长，而不是线性增长</strong>。</p><p id="2UCAEKJ0">换句话说，树注意力的优势随着设备数量增大会更加明显。实验中，在128卡、512万序列长度设置时达到最高8倍加速。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2F6e3b0d31j00si38gc0016d000hs00ahm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJ2">与环注意力相比，<strong>峰值内存占用也能节省不少</strong>。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2F51799befj00si38gd004bd000hs00ipm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJ4">相关代码已经开源，基于谷歌jax框架，已和Flash Attention整合，实现起来<strong>只需要30行代码</strong>。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2Fbf86bf31j00si38gc002nd000hs00stm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJ6">论文一公布，就被业界评价为“对高推理需求的大型公司很重要”。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2Fj00si38gc000jd000hs0064m.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJ8">这下和黄仁勋的GPU<strong>“买的越多，省的越多”论</strong>对上了，英伟达再次赢麻。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2F405b49a9j00si38gc000qd000hs00a0m.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJA">注意力机制的能量视角</p><p id="2UCAEKJB">首先简单回顾一下这次被拿来对比的<strong>环注意力</strong>，由UC伯克利大牛Pieter Abeel团队提出。</p><p id="2UCAEKJC">环注意力被认为是让<strong>上一波大模型纷纷扩展到百万上下文的关键</strong>，从谷歌Gemini 1.5到后来的Llama 3.1系列都用了它的某种变体。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2Fb099adc0j00si38gc000vd000hs006im.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJE">简单来说，环注意力的核心思想是<strong>将长序列分成多个Block，每个GPU处理一个</strong>。在拓扑意义上相当于所有GPU排成一个圆环，将Key-Value信息传下去，同时从上一个GPU接收信息。</p><p id="2UCAEKJF">只要保证计算时间比数据传输时间长，这个过程就不会造成额外开销。</p><p id="2UCAEKJG">同时与之前的近似方法不同，环注意力不会损失精度，保持了完整的注意力计算。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2F55432ca5j00si38gc001cd000hs00hkm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJI">最新的<strong>树注意力</strong>，在分块计算、跨设备并行、保持精度特性的基础上，<strong>提出了一种自注意力的能量函数，通过计算梯度利用树形拓扑优化多GPU间的通信</strong>。</p><p id="2UCAEKJJ">传统上，人们把注意力看作Query向量与Key向量的相似度匹配，再对Value向量做加权求和。</p><p id="2UCAEKJK">树注意力团队在Hopfield网络等基于能量的模型相关研究基础上，将注意力解释为一个能量函数对某变量的梯度。</p><p id="2UCAEKJL">存在一个标量能量函数F，它依赖于Key、Query、Value以及一个辅助变量ζ，而注意力的结果恰好等于F对ζ的梯度在ζ=0处的值。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2F8ac99e53j00si38gc0006d000hs003bm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJN">结合自动微分等技术，从能量和梯度的视角看待自注意力，暗示了<strong>只要能高效计算F就能高效计算自注意力</strong>。</p><p id="2UCAEKJO">具体到语言模型中基于KV缓存的解码，能量函数可以表示成：</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2Fcf51bf0fj00si38gc0009d000hs0046m.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJQ">由于logsumexp和max运算操作都满足结合律，可以按任意顺序进行，而不会影响最终结果。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2F526a3ecdj00si38gc000jd000hs008cm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKJS">在此前提下，团队设计了新的并行化算法，先在各GPU上并行计算局部能量函数，再通过树状的Allreduce汇总各处结果，最后用自动微分取梯度，即可得到注意力的输出。</p><p id="2UCAEKJT">全过程仅需与计算能量函数相同的时间开销，而显存占用也几乎没有额外负担。</p><p id="2UCAEKJU">树注意力在设计上还<strong>充分利用了GPU集群的两级拓扑特点</strong>——即同节点内使用高速NVLink，而节点间则依赖IB或以太网等。</p><p id="2UCAEKJV">相比之下，环形注意力天然不适应这种拓扑，难以将通信与计算很好地重叠，终会被最慢的互联带宽所制约。</p><p id="2UCAEKK0">最后值得一提的是，虽然理论上单GPU内部也可用类似策略提速，但当前硬件的流式处理器（SM）间通信还是共享内存，优势并不明显。</p><p id="2UCAEKK1">不过，<strong>英伟达在H100上实验性地支持了SM间点对点的指令</strong>，这为未来单卡注意力优化带来了新的想象空间。</p><p id="2UCAEKK2">最被低估的AI实验室之一</p><p id="2UCAEKK3">树注意力团队主要成员来自<strong>Zyphra</strong>，一家新兴的AI创业公司，被评价为<strong>“当前最被低估的AI实验室之一”</strong>。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2Fcdj00si38gc000pd000hs006mm.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKK5">Zyphra重点关注边缘AI、端侧AI，<strong>曾发布基于Mamba架构的基础模型Zamba</strong>。</p><p id="2UCAEKK6">创始人Krithik Puthalath以及树注意力共同一作Vasudev Shyam、Jonathan Pilault<strong>都有数学和理论物理学术背景</strong>。</p><p class="f_center"><img src="https://nimg.ws.126.net/?url=http%3A%2F%2Fdingyue.ws.126.net%2F2024%2F0812%2F8cj00si38gc000md000hs0097m.jpg&thumbnail=660x&quality=80&type=jpg"/><br/><br/></p><p id="2UCAEKK8">论文地址：<br/>https://arxiv.org/abs/2408.04093</p><p id="2UCAEKK9">参考链接：<br/>[1]https://x.com/ryu0000000001/status/<br/>[2]https://www.zyphra.com/post/tree-attention-topology-aware-decoding-for-long-context-attention-on-gpu-clusters</p>