白交 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI
Flash is all you need！
最近，一個(gè)超快且省內(nèi)存的注意力算法FlashAttention火了。
通過感知顯存讀取/寫入，F(xiàn)lashAttention的運(yùn)行速度比PyTorch標(biāo)準(zhǔn)Attention快了2-4倍，所需內(nèi)存也僅是其5%-20%。
而它的表現(xiàn)還不止于此。
- 訓(xùn)練BERT速度相較于MLPerf訓(xùn)練記錄提升15%；
- 訓(xùn)練GPT-2的速度提高3.5倍；
- 訓(xùn)練Transformer的速度比現(xiàn)有基線快。
網(wǎng)友們紛紛表示驚嘆：Great Job！這項(xiàng)工作對(duì)我來說很有用。
來看看這是一項(xiàng)什么樣的研究~
FlashAttention
本文提出了一種IO感知精確注意力算法。
隨著Transformer變得越來越大、越來越深，但它在長序列上仍然處理的很慢、且耗費(fèi)內(nèi)存。（自注意力時(shí)間和顯存復(fù)雜度與序列長度成二次方）
現(xiàn)有近似注意力方法，在試圖通過去犧牲模型質(zhì)量，以降低計(jì)算復(fù)雜度來解決該問題。
但存在一定的局限性，即不能提升運(yùn)行時(shí)的訓(xùn)練速度。
研究者認(rèn)為，應(yīng)該讓注意力算法具有IO感知，即考慮顯存級(jí)間的讀寫，比如大但慢的HBM（High Bandwidth Memory）技術(shù)與小但快的SRAM。
基于這樣的背景，研究人員提出了FlashAttention，具體有兩種加速技術(shù)：按塊遞增計(jì)算即平鋪、并在后向傳遞中重新計(jì)算注意力，將所有注意力操作融合到CUDA內(nèi)核中。
FlashAttention使用平鋪來防止大的×注意力矩陣（虛線框）在GPU HBM上物化(materialization)。在外部循環(huán)中（紅色箭頭），F(xiàn)lashAttention循環(huán)通過K和V矩陣的塊，并將其加載到SRAM。
在每個(gè)區(qū)塊中，F(xiàn)lashAttention 循環(huán)Q矩陣的區(qū)塊（藍(lán)色箭頭）將其加載到 SRAM，并將注意力計(jì)算的輸出寫回 HBM。
這樣就產(chǎn)生了一種注意力算法，在實(shí)際耗時(shí)（wall-clock time）內(nèi)，其內(nèi)存效率和速度都很高，相比于標(biāo)準(zhǔn)的注意力算法可以更少地訪問HBM。
結(jié)果比現(xiàn)有注意力算法都快
研究人員評(píng)估了FlashAttention來訓(xùn)練Transformer的影響，包括訓(xùn)練時(shí)間、模型準(zhǔn)確性，以及注意力運(yùn)行時(shí)間和內(nèi)存效率。
首先在訓(xùn)練速度上。FlashAttention比MLPerf 1.1的BERT速度記錄高出15%。
在實(shí)現(xiàn)GPT-2上，比HuggingFace速度高出3倍，比Megatron的標(biāo)準(zhǔn)Transformer速度高出1.8倍，F(xiàn)lashAttention將LRA（long-range arena）的基準(zhǔn)速度提高了2.4倍。
在模型質(zhì)量，F(xiàn)lashAttention將Transformer擴(kuò)展到更長的序列，并且質(zhì)量更好。
長上下文的語言建模。
如圖所示，使用FlashAttention可以讓GPT-2上下文長度增加4倍的情況下，訓(xùn)練時(shí)間還比Megatron-LM優(yōu)化實(shí)現(xiàn)快30%，同時(shí)也獲得了0.7的困惑度（困惑度越低，說明語言模型越好）。
長文檔分類
對(duì)較長序列的Transformer訓(xùn)練可以提高M(jìn)IMIC-III和ECtHR數(shù)據(jù)集的性能，比如序列長度為16K在MIMIC上比長度512多出4.3分。
MIMIC-III：包含重癥監(jiān)護(hù)室病人的出院總結(jié)，每個(gè)都有多個(gè)標(biāo)簽注釋；ECtHR：包含歐洲人權(quán)法案的法律案件；兩個(gè)數(shù)據(jù)集都包含很長的文本文件。
此外，還完成了第一個(gè)能在Path-X和Path-256任務(wù)中實(shí)現(xiàn)非隨機(jī)性能的Transformer模型。
之后，研究人員還完成了基準(zhǔn)測(cè)試，測(cè)量FlashAttention和塊狀稀疏（Block-Sparse）FlashAttention的運(yùn)行時(shí)間和內(nèi)存性能，并與帶有40GB HBM的A100 GPU上的各種注意力基線進(jìn)行了比較。
結(jié)果顯示，F(xiàn)lashAttention的運(yùn)行時(shí)間，比PyTorch注意力實(shí)現(xiàn)快3倍；在短序列情況下，F(xiàn)lashAttention在短序列中仍比近似和稀疏注意力運(yùn)行得快；至于塊狀稀疏的FlashAttention，在所有的序列長度上都比現(xiàn)有注意力實(shí)現(xiàn)都快。
至于在顯存效率方面，F(xiàn)lashAttention比PyTorch注意力基線高20倍。
在64k序列長度、其他所有算法都已經(jīng)耗盡顯存的情況下，F(xiàn)lashAttention的效率仍比Linformer高2倍。
斯坦福博士一作
這篇研究來自斯坦福大學(xué)計(jì)算機(jī)系以及紐約州立大學(xué)布法羅分校。共同一作是兩位斯坦福計(jì)算機(jī)博士生Tri Dao和Dan Fu。
感興趣的朋友，可戳下方論文鏈接了解更多~
論文鏈接：
https://arxiv.org/abs/2205.14135
GitHub鏈接：
https://github.com/HazyResearch/flash-attention
參考鏈接：
https://twitter.com/tri_dao/status/1531437619791290369
