斯坦福博士獨(dú)作！大模型訓(xùn)練速度再翻倍，還官宣加入明星創(chuàng)業(yè)公司當(dāng)首席科學(xué)家

豐色 2023-07-18 13:00:57 來源：量子位

豐色 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

現(xiàn)有大語言模型的訓(xùn)練和推理速度，還能再快一點(diǎn)——

快多少？2-4倍。

各種大模型都在用的FlashAttention今天正式發(fā)布第2代并開源，所有Transformer架構(gòu)的模型都可使用它來加速。

一代方法去年6月發(fā)布，無需任何近似即可加速注意力并減少內(nèi)存占用。

現(xiàn)在，F(xiàn)lashAttention-2將它再度升級(jí)，使其核心注意力操作的速度再提高2倍，端到端訓(xùn)練Transformer時(shí)的速度再提高1.3倍，并可在英偉達(dá)A100上訓(xùn)練時(shí)實(shí)現(xiàn)72%的模型FLOP利用率（一般模型都在50%上下）。

鑒于現(xiàn)在煉一個(gè)大語言模型的成本高達(dá)數(shù)千萬美元，F(xiàn)lashAttention-2這一系列操作直接就能幫我們省掉數(shù)百萬（美元）！

網(wǎng)友驚得臟話都出來了（狗頭）：

目前，這個(gè)項(xiàng)目已在GitHub上收獲4.4k標(biāo)星。

與此同時(shí)，我們注意到，它的一作已經(jīng)完成斯坦福博士學(xué)位并加盟大模型創(chuàng)業(yè)公司Together AI。

具體實(shí)現(xiàn)

據(jù)介紹，一代FlashAttention是一種對注意力計(jì)算重新排序的算法，它利用經(jīng)典方法如tiling（切片）來顯著加快計(jì)算速度，并將序列長度的內(nèi)存使用量從二次方減為線性。

其中tiling方法指的是將輸入塊從HBM（GPU內(nèi)存）加載到SRAM（快速緩存），然后對該塊進(jìn)行attention操作，再更新HBM中的輸出。

對HBM的反復(fù)讀寫就成了最大的性能瓶頸。

正是這種通過避免將大型中間注意力矩陣寫入HBM的方法，F(xiàn)lashAttention減少了內(nèi)存讀/寫量，從而帶來2-4倍的時(shí)鐘時(shí)間加速。

然而，這個(gè)算法仍然存在一些低效率的問題，導(dǎo)致它仍然不如優(yōu)化矩陣乘法?(GEMM)?運(yùn)算來得快，最終僅達(dá)到理論最大FLOPs/s的25-40%（例如在A100上最多124 TFLOPs/s）。

究其原因，還是因?yàn)椴煌€程塊之間的工作和GPU上的wrap劃分不理想。

在此，F(xiàn)lashAttention-2進(jìn)行了三方面的改進(jìn)。

首先，在基礎(chǔ)算法上，減少非matmul（矩陣乘法）?FLOP的數(shù)量。

一層原因是由于現(xiàn)代GPU具有專門的計(jì)算單元，matmul速度更快。例如A100上FP16/BF16 matmul的最大理論吞吐量為312TFLOPs/s，但非matmul FP32的理論吞吐量僅為19.5 TFLOPs/s。

另一層原因是價(jià)格考量，畢竟每個(gè)非matmul FLOP比matmul FLOP貴16倍。同時(shí)在matmul FLOP上花費(fèi)盡可能多的時(shí)間也能保持高吞吐量。

為此，作者重寫了FlashAttention中的softmax trick，無需更改輸出即可減少重新縮放操作的數(shù)量，以及邊界檢查和因果屏蔽操作（causal masking operation）。

其次，當(dāng)batch size較小時(shí)并行化以獲得更高的占用率。

FlashAttention一代在batch size和注意力頭數(shù)量上進(jìn)行并行化。

由于它使用1個(gè)線程塊來處理1個(gè)注意力頭，總共就有（batch_size*注意力頭數(shù)）個(gè)線程塊，每個(gè)線程塊被安排在流式多處理器?(SM)?上運(yùn)行。

當(dāng)在像A100這樣有108個(gè)SM處理器上操作時(shí)，如果線程塊很多比如>=80，這樣的調(diào)度安排就很有效。

而在長序列的情況下，也就是batch size和頭數(shù)量很少（?。r(shí)，就需要在序列長度維度上另外進(jìn)行并行化來更好地利用GPU上的多處理器了。

這個(gè)改進(jìn)也是FlashAttention-2速度顯著提升的一大原因。

最后，改進(jìn)工作分區(qū)。

在線程塊內(nèi)，我們必須確定如何在不同的warp之間劃分工作。通常是每個(gè)塊使用4或8個(gè)warp，現(xiàn)在，作者改進(jìn)了這一方式，來減少不同warp之間的同步和通信量，從而減少共享內(nèi)存讀寫操作。

如下圖左所示，F(xiàn)lashAttention一代的做法是將K和V分割到4個(gè)warp上，同時(shí)保持Q可被所有warp訪問。這樣的后果是所有warp都需要將其中間結(jié)果寫入共享內(nèi)存，然后進(jìn)行同步再將中間結(jié)果相加，非常低效，減慢了FlashAttention中的前向傳播速度。

而在FlashAttention-2中，作者將Q分為四個(gè)warp，同時(shí)保證所有warp都可訪問K和V。

每個(gè)warp執(zhí)行矩陣乘法獲得Q K^T的切片后，只需與V的共享切片相乘即可獲得相應(yīng)的輸出。也就是說warp之間不需要通信，那么共享內(nèi)存讀寫操作就少了很多，速度也就提上來了。

除了這三個(gè)大改進(jìn)，F(xiàn)lashAttention-2還有兩個(gè)小改動(dòng)：
一是注意力頭數(shù)從128增至256，這意味著GPT-J、CodeGen和CodeGen2以及StableDiffusion 1.x等模型都可以使用 FlashAttention-2來進(jìn)行加速和內(nèi)存節(jié)省了；

二是支持多查詢注意力（MQA）和分組查詢注意力（GQA）。

實(shí)驗(yàn)評(píng)估

作者在A100 80GB SXM4 GPU上對不同配置（有無causal mask，頭數(shù)量64或128）下的運(yùn)行時(shí)間進(jìn)行了測量。

結(jié)果發(fā)現(xiàn)：

FlashAttention-2比FlashAttention（包括xformers庫和Triton中的其他實(shí)現(xiàn)）快大約2倍，這也意味我們可以用與之前訓(xùn)練8k上下文模型相同的價(jià)格來訓(xùn)練具有16k上下文的模型了（也就是模型上下文長度加倍）。

而與PyTorch中的標(biāo)準(zhǔn)注意力實(shí)現(xiàn)相比，F(xiàn)lashAttention-2的速度最高可達(dá)9倍。

此外，有了FlashAttention-2，我們只需在H100 GPU上運(yùn)行相同的實(shí)現(xiàn)（不使用特殊指令利用TMA和第四代Tensor Core等新硬件功能），訓(xùn)練速度就可以跑到高達(dá)335TFLOPs/s的成績。

以及當(dāng)用于端到端訓(xùn)練GPT式模型時(shí)，F(xiàn)lashAttention-2還能在A100上實(shí)現(xiàn)高達(dá)225TFLOPs/s的速度（模型FLOPs利用率達(dá)72%）。這與已經(jīng)優(yōu)化程序足夠高的FlashAttention相比，速度再提高了1.3倍。

一作加入大模型創(chuàng)業(yè)公司

FlashAttention-2論文僅顯示一位作者：Tri Dao。他也是FlashAttention一代的兩位共同作者之一。

據(jù)了解，Tri Dao的研究方向?yàn)闄C(jī)器學(xué)習(xí)和系統(tǒng)的交叉領(lǐng)域，去年拿下ICML 2022杰出論文亞軍獎(jiǎng)。

最近他剛剛獲得斯坦福大學(xué)計(jì)算機(jī)科學(xué)博士學(xué)位，即將上升普林斯頓大學(xué)助理教授，并已宣布加盟生成式AI創(chuàng)業(yè)公司Together AI（該司主要目標(biāo)構(gòu)建一個(gè)用于運(yùn)行、訓(xùn)練和微調(diào)開源模型的云平臺(tái)）擔(dān)任首席科學(xué)家。

One More Thing

最后，有網(wǎng)友發(fā)現(xiàn)，除了FlashAttention-2，最近還有一系列類似成果，包括DeepSpeed的ZeRO++、馬薩諸塞大學(xué)de ReLoRA。

它們都是用于加速大型模型預(yù)訓(xùn)練和微調(diào)，這些研究成果讓他覺得：

未來在低vram低帶寬的消費(fèi)顯卡上訓(xùn)練大模型，似乎已不是在做夢了。

大家認(rèn)為呢？

論文地址：
https://tridao.me/publications/flash2/flash2.pdf
博文地址：
https://princeton-nlp.github.io/flash-atttention-2/
GitHub主頁：
https://github.com/Dao-AILab/flash-attention
參考鏈接：
[1]https://twitter.com/tri_dao/status/1680987577913065472?s=20
[2]https://twitter.com/togethercompute/status/1680994294625337344?s=20
[3]https://twitter.com/main_horse/status/1681041183559254017?s=20