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

AIGC動態(tài)歡迎閱讀

原標題：30行代碼，500萬長文本推理提速8倍！「樹注意力」讓GPU越多省的越多
關鍵字：注意力,能量,梯度,函數(shù),拓撲
文章來源：量子位
內容字數(shù)：0字

內容摘要：

夢晨 發(fā)自 凹非寺量子位 | 公眾號 QbitAI跨GPU的注意力并行，最高提速8倍，支持512萬序列長度推理。
環(huán)注意力（Ring Attention）后繼者——樹注意力（Tree Attention）來了。
最關鍵之處在于，通信步數(shù)隨設備數(shù)量成對數(shù)增長，而不是線性增長。
換句話說，樹注意力的優(yōu)勢隨著設備數(shù)量增大會更加明顯。實驗中，在128卡、512萬序列長度設置時達到最高8倍加速。
與環(huán)注意力相比，峰值內存占用也能節(jié)省不少。
相關代碼已經開源，基于谷歌jax框架，已和Flash Attention整合，實現(xiàn)起來只需要30行代碼。
論文一公布，就被業(yè)界評價為“對高推理需求的大型公司很重要”。
這下和黃仁勛的GPU“買的越多，省的越多”論對上了，英偉達再次贏麻。
注意力機制的能量視角首先簡單回顧一下這次被拿來對比的環(huán)注意力，由UC伯克利大牛Pieter Abeel團隊提出。
環(huán)注意力被認為是讓上一波大模型紛紛擴展到百萬上下文的關鍵，從谷歌Gemini 1.5到后來的Llama 3.1系列都用了它的某種變體。
簡單來說，環(huán)注意力的核心思想是將長序列分成多個Block，每個GPU處理一個。

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

聯(lián)系作者

文章來源：量子位
作者微信：
作者簡介：