跨GPU的注意力并行，最高提速8倍，支持512萬(wàn)序列長(zhǎng)度推理。

環(huán)注意力（Ring Attention）后繼者——樹(shù)注意力（Tree Attention）來(lái)了。

最關(guān)鍵之處在于，通信步數(shù)隨設(shè)備數(shù)量成對(duì)數(shù)增長(zhǎng)，而不是線性增長(zhǎng)。

換句話說(shuō)，樹(shù)注意力的優(yōu)勢(shì)隨著設(shè)備數(shù)量增大會(huì)更加明顯。實(shí)驗(yàn)中，在128卡、512萬(wàn)序列長(zhǎng)度設(shè)置時(shí)達(dá)到最高8倍加速。

與環(huán)注意力相比，峰值內(nèi)存占用也能節(jié)省不少。

相關(guān)代碼已經(jīng)開(kāi)源，基于谷歌jax框架，已和Flash Attention整合，實(shí)現(xiàn)起來(lái)只需要30行代碼。

論文一公布，就被業(yè)界評(píng)價(jià)為“對(duì)高推理需求的大型公司很重要”。

這下和黃仁勛的GPU“買(mǎi)的越多，省的越多”論對(duì)上了，英偉達(dá)再次贏麻。

注意力機(jī)制的能量視角

首先簡(jiǎn)單回顧一下這次被拿來(lái)對(duì)比的環(huán)注意力，由UC伯克利大牛Pieter Abeel團(tuán)隊(duì)提出。

環(huán)注意力被認(rèn)為是讓上一波大模型紛紛擴(kuò)展到百萬(wàn)上下文的關(guān)鍵，從谷歌Gemini 1.5到后來(lái)的Llama 3.1系列都用了它的某種變體。

簡(jiǎn)單來(lái)說(shuō)，環(huán)注意力的核心思想是將長(zhǎng)序列分成多個(gè)Block，每個(gè)GPU處理一個(gè)。在拓?fù)湟饬x上相當(dāng)于所有GPU排成一個(gè)圓環(huán)，將Key-Value信息傳下去，同時(shí)從上一個(gè)GPU接收信息。

只要保證計(jì)算時(shí)間比數(shù)據(jù)傳輸時(shí)間長(zhǎng)，這個(gè)過(guò)程就不會(huì)造成額外開(kāi)銷。

同時(shí)與之前的近似方法不同，環(huán)注意力不會(huì)損失精度，保持了完整的注意力計(jì)算。

最新的樹(shù)注意力，在分塊計(jì)算、跨設(shè)備并行、保持精度特性的基礎(chǔ)上，提出了一種自注意力的能量函數(shù)，通過(guò)計(jì)算梯度利用樹(shù)形拓?fù)鋬?yōu)化多GPU間的通信。

傳統(tǒng)上，人們把注意力看作Query向量與Key向量的相似度匹配，再對(duì)Value向量做加權(quán)求和。

樹(shù)注意力團(tuán)隊(duì)在Hopfield網(wǎng)絡(luò)等基于能量的模型相關(guān)研究基礎(chǔ)上，將注意力解釋為一個(gè)能量函數(shù)對(duì)某變量的梯度。

存在一個(gè)標(biāo)量能量函數(shù)F，它依賴于Key、Query、Value以及一個(gè)輔助變量ζ，而注意力的結(jié)果恰好等于F對(duì)ζ的梯度在ζ=0處的值。

結(jié)合自動(dòng)微分等技術(shù)，從能量和梯度的視角看待自注意力，暗示了只要能高效計(jì)算F就能高效計(jì)算自注意力。

具體到語(yǔ)言模型中基于KV緩存的解碼，能量函數(shù)可以表示成：

由于logsumexp和max運(yùn)算操作都滿足結(jié)合律，可以按任意順序進(jìn)行，而不會(huì)影響最終結(jié)果。

在此前提下，團(tuán)隊(duì)設(shè)計(jì)了新的并行化算法，先在各GPU上并行計(jì)算局部能量函數(shù)，再通過(guò)樹(shù)狀的Allreduce匯總各處結(jié)果，最后用自動(dòng)微分取梯度，即可得到注意力的輸出。

全過(guò)程僅需與計(jì)算能量函數(shù)相同的時(shí)間開(kāi)銷，而顯存占用也幾乎沒(méi)有額外負(fù)擔(dān)。

樹(shù)注意力在設(shè)計(jì)上還充分利用了GPU集群的兩級(jí)拓?fù)涮攸c(diǎn)——即同節(jié)點(diǎn)內(nèi)使用高速NVLink，而節(jié)點(diǎn)間則依賴IB或以太網(wǎng)等。

相比之下，環(huán)形注意力天然不適應(yīng)這種拓?fù)?，難以將通信與計(jì)算很好地重疊，終會(huì)被最慢的互聯(lián)帶寬所制約。

最后值得一提的是，雖然理論上單GPU內(nèi)部也可用類似策略提速，但當(dāng)前硬件的流式處理器（SM）間通信還是共享內(nèi)存，優(yōu)勢(shì)并不明顯。

不過(guò)，英偉達(dá)在H100上實(shí)驗(yàn)性地支持了SM間點(diǎn)對(duì)點(diǎn)的指令，這為未來(lái)單卡注意力優(yōu)化帶來(lái)了新的想象空間。

最被低估的AI實(shí)驗(yàn)室之一

樹(shù)注意力團(tuán)隊(duì)主要成員來(lái)自Zyphra，一家新興的AI創(chuàng)業(yè)公司，被評(píng)價(jià)為“當(dāng)前最被低估的AI實(shí)驗(yàn)室之一”。

Zyphra重點(diǎn)關(guān)注邊緣AI、端側(cè)AI，曾發(fā)布基于Mamba架構(gòu)的基礎(chǔ)模型Zamba。

創(chuàng)始人Krithik Puthalath以及樹(shù)注意力共同一作Vasudev Shyam、Jonathan Pilault都有數(shù)學(xué)和理論物理學(xué)術(shù)背景。

論文地址：https://arxiv.org/abs/2408.04093