PyTorch官方認(rèn)可！斯坦福博士新作：長上下文LLM推理速度提8倍

豐色 2023-10-17 15:37:02 來源：量子位

豐色 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

這兩天，F(xiàn)lashAttention團(tuán)隊推出了新作：

一種給Transformer架構(gòu)大模型推理加速的新方法，最高可提速8倍。

該方法尤其造福于長上下文LLM，在64k長度的CodeLlama-34B上通過了驗證。

甚至得到了PyTorch官方認(rèn)可：

如果你之前有所關(guān)注，就會記得用FlashAttention給大模型加速效果真的很驚艷。

不過它僅限于訓(xùn)練階段。

因此，這一新成果一出，就有網(wǎng)友表示：

等推理加速等了好久，終于來了。

據(jù)介紹，這個新方法也是在FlashAttention的基礎(chǔ)之上衍生而出，主要思想也不復(fù)雜：

用并行操作盡快加載Key和Value緩存，然后分別重新縮放再合并結(jié)果，最終獲得推理速度上的大幅提升。
詳細(xì)來看。

提速8倍的長上下文推理方法來了

該方法被命名為Flash-Decoding。

背景與動機

根據(jù)作者介紹：

LLM的推理（即“解碼”）過程是迭代的，即一次生成一個token，組成一個完整句子需要n個token以及n次前向傳遞。

不過，由于我們可以緩存之前計算出來的token，所以單個生成步驟并不總是依賴于上下文長度。

但有一個操作例外：注意力?(attention)，它不能隨著上下文長度靈活擴展。

鑒于長上下文已成趨勢，比如目前最大的開源LLM已達(dá)100k（CodeLlama），我們不得不注意到attention在大模型推理過程中浪費了太多時間，時間就是金錢。

更別提attention在batch size上進(jìn)行擴展時，即使模型上下文相對較短，它也可能成為性能瓶頸（因為模型要讀取的內(nèi)存量與batch size成比例，而它僅取決于模型其余部分的大?。?。

怎么破解？

不可復(fù)用的FlashAttention優(yōu)化

模型在推理也就是解碼過程中，為了計算softmax(queries @keys.transpose)@values這兩個值，生成的每個新token都需要關(guān)注先前的所有token。

團(tuán)隊先前的工作FlashAttention，已經(jīng)在訓(xùn)練階段對此操作進(jìn)行了優(yōu)化。

當(dāng)時，F(xiàn)lashAttention解決的主要瓶頸是讀寫中間結(jié)果的內(nèi)存帶寬（例如，Q @ K^T）。

然而，在推理階段，我們要面對的瓶頸變了，導(dǎo)致FlashAttention所做的優(yōu)化并不能直接拿過來應(yīng)用。

具體而言：

在階段階段，F(xiàn)lashAttention在batch size和查詢長度維度上進(jìn)行并行化。

在推理階段，查詢長度通常為1，這意味著如果batch size小于GPU上的流式多處理器數(shù)量（例如，A100為108），該操作將僅使用GPU的一小部分。

這對于長上下文情況尤甚，因為長上下文需要較小的batch size才能適應(yīng)GPU內(nèi)存。

所以，結(jié)果就是，當(dāng)batch size為1時，F(xiàn)lashAttention將只占用不足1%的GPU，非常不劃算。

當(dāng)然，你可能會說，不用FlashAttention也行，用矩陣乘法原語來完注意力操作。

不過，作者指出，這種情況又會完全占用GPU，并啟動非常多的寫入和讀取中間結(jié)果的內(nèi)核，也不是最佳辦法。

Flash-Decoding誕生

最終，基于以上考量，作者在FlashAttention的基礎(chǔ)上，添加了一個新的并行化緯度：key和value序列長度。

這個方法（即Flash-Decoding）結(jié)合上述兩種方法的優(yōu)點：

與FlashAttention一樣，它在全局內(nèi)存中存儲的額外數(shù)據(jù)非常少，但只要上下文長度足夠大，即使batch size很小，它也可以充分利用GPU。

詳細(xì)來看，F(xiàn)lash-Decoding一共分為三個步驟：

1、先將key和value值分成更小的塊。

2、用FlashAttention并行計算每塊分割的查詢注意力。并為每行和每塊分割寫入一個額外標(biāo)量：注意力值的log-sum-exp。

3、最后，通過減少所有分割來計算實際輸出，使用log-sum-exp來scale每塊分割的貢獻(xiàn)。

作者指出，由于attention/softmax可以迭代計算，以上所有操作均可行。

并且在Flash-Decoding中，ttention/softmax既可以在分割塊內(nèi)，也可以跨分割塊來執(zhí)行最終的縮減，只不過后者可縮減的步驟很少。

而在實際操作中，步驟1不涉及任何GPU操作，因為key和value塊是完整的張量視圖。然后由2個獨立的內(nèi)核分別執(zhí)行步驟2和3。

最高提速8倍

驗證環(huán)節(jié)，作者在CodeLLaMa-34b（架構(gòu)與Llama 2相同）上對其解碼吞吐量進(jìn)行了基準(zhǔn)測試。

具體以tok/s為單位，測量了512到64k序列長度下的解碼速度（上限為從內(nèi)存中讀取整個模型以及KV緩存所需的時間），并和多種計算注意力的方法進(jìn)行對比，包括：

• Pytorch，使用純PyTorch原語運行注意力
• FlashAttention v2
• FasterTransformer：使用FasterTransformer注意力內(nèi)核

最終，F(xiàn)lash-Decoding最高可將長序列解碼速度提升8倍，并比其他方法具 有更好的擴展性（受長度影響較小）

此外，作者還在A100上對各種序列長度和batch size的縮放多頭注意力進(jìn)行了微基準(zhǔn)測試。

結(jié)果顯示，當(dāng)序列長度擴展到64k時，Flash-Decoding實現(xiàn)了幾乎恒定的運行時間。

如何使用？

以下是Flash-Decoding的獲取途徑，戳文末官方博客即可找到地址：

• FlashAttention包，2.2版本及以上
• xFormers包，0.0.22版本及以上

調(diào)度程序?qū)⒏鶕?jù)問題的大小自動使用Flash-Decoding或 FlashAttention方法。

團(tuán)隊介紹

目前Flash-Decoding還沒出論文，但作者團(tuán)隊已透露，這次不再是Tri Dao“單打獨斗”，不過一作仍然是他。

Tri Dao今年博士畢業(yè)于斯坦福，7月份加盟大模型創(chuàng)業(yè)公司Together AI擔(dān)任首席科學(xué)家。

明年9月將上任普林斯頓大學(xué)助理教授，他是FlashAttention v1和v2的主要作者。

剩下三位作者分別是：

Daniel Haziza，F(xiàn)acebook AI Research研究工程師，主要負(fù)責(zé)xformers（用于訓(xùn)練加速的開源框架）；

Francisco Massa，同F(xiàn)acebook AI Research研究工程師， 主要從事PyTorch相關(guān)工作；

Grigory Sizov，Meta機器學(xué)習(xí)工程師，主要工作是優(yōu)化GPU上的LLM推理和其他AI工作負(fù)載，為PyTorch生態(tài)做出過貢獻(xiàn)。

官方博客：
https://princeton-nlp.github.io/flash-decoding/
參考鏈接：
https://twitter.com/tri_dao/status/1712904220519944411?s=20