30行代碼，500萬長文本推理提速8倍！「樹注意力」讓GPU越多省的越多

夢晨 2024-08-12 12:27:09 來源：量子位

夢晨 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

跨GPU的注意力并行，最高提速8倍，支持512萬序列長度推理。

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

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

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

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

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

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

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

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

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

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

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

只要保證計算時間比數(shù)據(jù)傳輸時間長，這個過程就不會造成額外開銷。

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

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

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

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

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

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

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

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

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

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

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

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

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

不過，英偉達(dá)在H100上實驗性地支持了SM間點對點的指令，這為未來單卡注意力優(yōu)化帶來了新的想象空間。

最被低估的AI實驗室之一

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

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

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

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

參考鏈接：
[1]https://x.com/ryu0000000001/status/1822043300682985642
[2]https://www.zyphra.com/post/tree-attention-topology-aware-decoding-for-long-context-attention-on-gpu-clusters