大模型訓練推理神作,又更新了!
主流大模型都在用的FlashAttention,剛剛升級第三代。
時隔一年,FlashAttention-3已經全方位升級。
訓練速度提升1.5-2倍,FP16下計算吞吐量高達740TFLOPs/s,達理論最大吞吐量75%,更充分利用計算資源,此前只能做到35%。
FP8下速度接近1.2PFLOPs/s!
同時誤差也進一步減小,FP8下的誤差比標準Attention減少2.6倍。
而且這一次,不再是一作Tri Dao單打獨鬥,FlashAttention-3直接和輝達、Meta、Google等合作,針對最強晶片H100專門做最佳化。
輝達CUTLASS團隊和cuDNN團隊,都直接為該研究提供支援。
同時和前作一樣,FlashAttention-3也將開源,PyTorch和Hugging Face中都整合。
作者之一Vijay Thakkar激動表示:
曾經在FA2發佈時,我就說過這句話。今天,我想再說一次:
看到CUTLASS和CuTe被用來開讓Tensor Core大顯身手的新演算法,真的泰褲辣。
前Stable Diffusion老闆Emad也非常關注這一進展,他推測使用FlashAttention-3,能將4090的FP8計算吞吐量推升到700+TFLOPs。
自初代發佈以來,FlashAttention已經使大模型速度提高了4-8倍,但還有一個遺憾:尚未充分利用現代 GPU。
針對輝達H100倍後的Hopper架構新特性,三代進行了專門最佳化。
整個系列的核心思路,是IO感知最佳化和分塊處理。
作者認為,傳統的注意力機制效率低的原因,在處理長序列時,會出現記憶體訪問操作頻繁,以及演算法複雜度指數級暴增這兩大問題。
FlashAttention通過IO感知最佳化將資料從較大但緩慢的高頻寬記憶體(HBM)載入到較小但更快的片上記憶體(SRAM),在SRAM中執行計算,減少了記憶體讀寫操作的次數。
分塊處理則是將輸入序列分成若干小塊,每次只處理一個小塊的資料。這種方法使得每次處理的資料量減少,從而降低了記憶體使用和計算複雜度。
這樣一來,兩個關鍵問題就得到瞭解決,這兩大核心思想也在本次的FlashAttention-3中得到了繼承。
△第一代FlashAttention原理圖
但是,第一代的FlashAttention也遺留下了平行性不夠強、工作分區劃分不合理,以及非矩陣乘法較多(GPU計算單元處理矩陣乘法比非矩陣速度更快)的問題。
針對這一問題,第二代FlashAttention通過重寫softmax,減少了重新縮放操作、邊界檢查和因果遮蔽操作的次數,使得大部分計算集中在矩陣乘法上。
另外,FlashAttention-2引入了序列長度維度上的平行化,並針對工作線上程塊之間的分配進行了最佳化,GPU利用效率更高了。
可以說前兩代當中,作者一直堅持著充分利用硬體特點這一思路,但站在今天的視角來看,對硬體的挖掘仍然不夠充分。
到了這次的FlashAttention-3,由於是直接和輝達官方合作,對輝達Hopper架構特點的理解更加透徹,軟硬體之間的協同進一步增強了。
FlashAttention-3的技術報告顯示,為了充分匹配Hopper架構,團隊主要做了三方面的技術升級。
首先,Hopper架構的一個重要特點是Tensor Core的非同步性,FlashAttention-3針對性地提出了一種非同步方式。
具體來說,FlashAttention-3引入了一種“生產者(Producer)-消費者(Consumer)”的程式設計模型,將注意力的計算劃分為兩個角色。
為了實現角色的劃分,作者引入了warp專門化技術,用不同的warp分別匹配生產者和消費者,讓兩者可以平行執行。
這其中利用了Hopper架構的動態warp暫存器分配特性,通過setmaxnreg指令最佳化了暫存器資源的利用。
為了進一步提高GPU的利用率,作者又提出了一種“乒乓調度”策略,讓一個warp組執行矩陣乘法時,另一個warp組執行softmax,從而實現計算的重疊。
具體講,FlashAttention-3使用CUDA的同步原語控制不同warp組之間的執行順序,讓不同warp組分別執行兩種運算,然後像乒乓球一樣交替運行。
第二大技術特點,是warp組內部GEMMs和softmax的重疊,核心奧義是重新安排計算的執行順序以提高GPU利用率。
與乒乓調度不同,這裡的計算重排處理的是warp組內部的重疊,而乒乓調度更關注組間協調。
實現方式上,FlashAttention-3提出了一種兩階段GEMM-softmax流水線方案,以打破不同操作之間的資料依賴。
通過引入額外的暫存器和共用記憶體緩衝區,FlashAttention-3實現了跨迭代的資料傳遞和重用。
在每個迭代中,Q·K^T的結果首先儲存在名為S_cur的緩衝區中,用於當前迭代的softmax計算,同時非同步執行下一個迭代的Q·K^T矩陣乘法,結果儲存在名為S_next的緩衝區中。
在執行當前迭代的P·V矩陣乘法時,非同步執行下一個迭代的softmax操作,並更新S_cur和S_next緩衝區。
第三項更新,是用更低的FP8精度替代FP16。
實際上,降低數值精度是一種常見的最佳化策略,可以顯著提高GPU的計算吞吐量和能效,Hopper GPU也引入了FP8精度的Tensor Core支援。
但是,直接將注意力計算從FP16轉換為FP8可能會引入較大的精度損失。
另外,FP8 Tensor Core對輸入資料的佈局也有特定的要求(K維度連續),不幸的是,注意力計算中的輸入資料儲存格式(頭維度連續)並不符合這樣的要求。
所以FlashAttention-3首先引入了一系列記憶體佈局轉換技術,動態轉置V矩陣的塊,改變其連續方式,從而適配FP8 Tensor Core的佈局要求。
在此基礎之上,為了獲得更高的計算精度,FlashAttention-3又採用了分塊量化和非相干處理技術。
傳統的量化方法通常對整個矩陣使用一個統一的縮放因子(per-tensor quantization),無法很好地適應不同區域的數值範圍。
FlashAttention-3則採用了分塊量化(block-wise quantization)的策略,為每個塊單獨設定縮放因子,更好地捕捉局部的數值分佈。
非相干處理(incoherent processing)技術則是通過隨機正交矩陣對輸入資料進行旋轉,破壞不同塊之間的相乾性,減少量化誤差的傳播。
這兩項技術的結合使得FlashAttention-3在FP8精度下取得了更高的計算精度,顯著優於傳統的量化方法。
結果,與基於傳統量化方法的FP8實現相比,FlashAttention-3的使得精度提高了2.6倍。
以上就是FlashAttention-3在充分研究Hopper架構特點後做出的三大更新,針對更新後的表現,作者主要進行了3方面測試。
首先來看注意力基準測試。
通過改變序列長度(512、1k、……16k),並設定批大小以確保總token數為16k。研究人員將隱藏維度設定為2048,頭維度設定為64、128或258,計算前向傳播、後向傳播。
對比標準Attention、FlashAttention-2、Triton、cuDNN和FlashAttention-3,在H100 80GB SXM5上FP16的執行階段間。
FlashAttention-3的前向傳播比FlashAttention-2快1.5-2倍,後向傳播快1.5-1.75倍。
與標準Attention相比,FlashAttention-3的速度快了3-16倍。
對於中長序列(1k以上),FlashAttention-3甚至超過了專門為H100最佳化的cuDNN。
在消融實驗中,通過對非因果FP16 FlashAttention-3進行了2階段WGMMA-softmax流水線和warp特殊化的消融研究,參數固定為{batch, seqlen, nheads, hdim} = {4, 8448, 16, 128}。
結果證實,FlashAttention-3改進帶來了顯著加速,從570提升到661。
另外,因為對FlashAttention的數值誤差感興趣,研究團隊還將FlashAttention-2、FlashAttention-3和標準Attention進行了比較。
為了模擬LLMs中的異常特徵和啟動,研究團隊生成了Q、K、V的條目,分佈為:N(0,1)+N(0,100)⋅Bernoulli(0.001)
也就是說,每個條目都服從均值為0、標準差為1的常態分配,但對於0.1%的條目,增加了一個獨立的項,其標準差為10。然後測量均方根誤差(RMSE)。
結果顯示,在FP16中,由於中間結果(softmax)保留在FP32中,FlashAttention-2和FlashAttention-3的RMSE比標準Attention減少1.7倍。
FP8的標準Attention使用每個張量的縮放,matmul累加器在FP32中,中間softmax結果保留在FP16中。由於塊量化和非相干處理,FP8中的FlashAttention-3比這個基線更準確2.6倍。
最後,論文還表示目前工作專注於Hopper架構,後續將推廣到其他硬體。
除了輝達為研究提供了技術支援外,Meta、Together AI和普林斯頓大學為研究提供了計算支援。 (量子位)