Flash Attention穩(wěn)定嗎?Meta、哈佛發(fā)現(xiàn)其模型權(quán)重偏差呈現(xiàn)數(shù)量級波動
Meta FAIR 聯(lián)合哈佛優(yōu)化大規(guī)模機器學習時產(chǎn)生的數(shù)據(jù)偏差,提供了新的研究框架。
眾所周知,大語言模型的訓練常常需要數(shù)月的時間,使用數(shù)百乃至上千個 GPU。以 LLaMA2 70B 模型為例,其訓練總共需要 1,720,320 GPU hours。由于這些工作負載的規(guī)模和復雜性,導致訓練大模型存在著獨特的系統(tǒng)性挑戰(zhàn)。
最近,許多機構(gòu)在訓練 SOTA 生成式 AI 模型時報告了訓練過程中的不穩(wěn)定情況,它們通常以損失尖峰的形式出現(xiàn),比如谷歌的 PaLM 模型訓練過程中出現(xiàn)了多達 20 次的損失尖峰。
數(shù)值偏差是造成這種訓練不穩(wěn)定性的潛在原因,由于大語言模型訓練運行成本極高,如何量化數(shù)值偏差儼然成為關(guān)鍵問題。
在最新的一項工作中,來自 Meta、哈佛大學的研究者開發(fā)了一個原則性定量方法來理解訓練優(yōu)化中的數(shù)值偏差,以此評估不同的最新優(yōu)化技術(shù),并確定它們在用于訓練大模型時是否可能引入意外的不穩(wěn)定性。
- 論文標題:Is Flash Attention Stable?
- 論文鏈接:https://arxiv.org/pdf/2405.02803
結(jié)果發(fā)現(xiàn),在一次單獨的前向傳遞過程中,F(xiàn)lash Attention 的數(shù)值偏差比 BF16 的 Baseline Attention 大一個數(shù)量級。
具體而言,該方法包括兩個階段,包括:
- 開發(fā)一個微基準來擾動給定優(yōu)化中的數(shù)值精度;
- 通過基于 Wasserstein 距離的數(shù)據(jù)驅(qū)動分析評估數(shù)值偏差如何轉(zhuǎn)化為模型權(quán)重的變化。
研究者分析了 SOTA 優(yōu)化技術(shù) Flash Attention ,并量化了可能引入的數(shù)值偏差。Flash Attention 是一種廣泛用于加速注意力機制的技術(shù),通常被認為是 Transformer 模型中的系統(tǒng)瓶頸。Flash Attention 在提高速度和減少內(nèi)存訪問量的同時,也依賴于算法優(yōu)化,而算法優(yōu)化有可能導致數(shù)值偏差的增加。
研究者假設添加重新縮放因子(rescaling factors )可能會引入無意的近似,導致數(shù)值折衷,這可能會在后續(xù)影響訓練穩(wěn)定性。
他們在多模態(tài)文本到圖像工作負載的背景下分析了 Flash Attention,以確定 Flash Attention 與其基線之間數(shù)值偏差的潛在重要性。最終,他們引入了一個框架來量化訓練優(yōu)化的數(shù)值偏差及其下游影響。
研究者在數(shù)值偏差量化上主要作出了以下兩點貢獻:
(1)設計了一個微基準來分離數(shù)值精度對數(shù)值偏差的影響。
研究者所設計的微基準作為一種技術(shù),用于衡量和量化傳統(tǒng)黑盒優(yōu)化(如 Flash Attention)所導致的數(shù)值偏差。通過擾動通常在提供的內(nèi)核中不可用的方面,他們開創(chuàng)性地發(fā)現(xiàn)在低數(shù)值精度(BF16)下,與 Baseline Attention 相比,F(xiàn)lash Attention 的數(shù)值偏差大約高出一個數(shù)量級。
(2)基于 Wasserstein Distance 度量進行了數(shù)據(jù)驅(qū)動的分析。
通過該分析,研究者將觀察到的數(shù)值偏差置于上下文,并為其對下游模型屬性的影響形成一個上限(upper bound)。在研究者的案例研究中,他們能夠限制觀察到的數(shù)值偏差的影響,并發(fā)現(xiàn):「Flash Attention 引入的模型權(quán)重偏差大約為低精度訓練的 1/2 至 1/5 倍。」
這項研究強調(diào)了開發(fā)一種原則性方法的重要性:「不僅要量化,而且要將訓練優(yōu)化對數(shù)值偏差的影響置于上下文中?!雇ㄟ^構(gòu)建代理(proxies)來將數(shù)值偏差置于上下文中,旨在推斷通常難以衡量的下游模型效果(即訓練不穩(wěn)定性)的可能性。
實驗方法
研究者首先開發(fā)了一個微基準來分離并研究 Flash Attention 引起的數(shù)值偏差。如圖 2 所示,他們通過對 Flash Attention 進行數(shù)值上的重新實現(xiàn),以分析不同的數(shù)值精度,并在算法的每個步驟應用潛在的優(yōu)化措施。
圖 2: 微基準設計摘要。
這是必要的,因為 Flash Attention 內(nèi)核目前僅支持 FP16 和 BF16 數(shù)值格式。該內(nèi)核還是 CUDA 代碼的包裝 API 調(diào)用,這使得擾動算法以檢查數(shù)值偏差的影響變得具有挑戰(zhàn)性。
相比之下,他們的微基準設計允許在算法內(nèi)部進行精度輸入和修改。研究者將微基準與原始的 Flash Attention kernel 進行了驗證。
他們進一步設計了一種技術(shù),以比較模型執(zhí)行過程中每個步驟的 Attention 矩陣的輸出。并修改了模型代碼,每次調(diào)用注意力時都計算 Baseline Attention 和 Flash Attention,這允許對相同的輸入矩陣進行精確的輸出矩陣比較。
為了將其置于上下文中,研究者還通過相同和獨立的訓練運行,使用 Max difference 和 Wasserstein Distance 度量來量化模型權(quán)重在整個訓練過程中的差異。
對于訓練實驗,研究者則使用一種將文本輸入轉(zhuǎn)換為圖像的生成式 AI workload(即文本到圖像模型)。他們使用 Shutterstock 數(shù)據(jù)集重新訓練模型,并在一組英偉達 80GB A100 GPU 集群上運行此實驗。
通過微基準量化數(shù)值偏差
研究者首先分析了 Flash Attention 在前向傳遞過程中的影響。他們利用微基準測試,在隨機初始化查詢、鍵、值向量相同的情況下,檢驗不同數(shù)值精度對 Attention 計算的輸出矩陣的影響。
正如圖 3 所示,當研究者使用從 BF16 到 FP64 變化的不同數(shù)值格式時,F(xiàn)lash Attention 和 Baseline Attention 之間的數(shù)值偏差隨著尾數(shù)位數(shù)的增加而減小。這表明數(shù)值差異是由于較少的尾數(shù)位數(shù)所固有的近似造成的。
圖 3:數(shù)值格式對于 Flash Attention 的數(shù)值偏差所產(chǎn)生的效果。
之后,研究者為進行標準比較,在 FP64 數(shù)值格式下的 Baseline Attention 設置了「黃金值」,然后將不同數(shù)值格式下的 Attention 輸出與該值進行了比較(如圖 4 所示)。
圖 4:FP64 下 Baseline Attention「黃金值」的比較。
結(jié)果表明,F(xiàn)lash Attention 的數(shù)值偏差大約是在 BF16 下 Baseline 的 10 倍。
為了進一步分析這種觀察到的數(shù)值偏差,研究者保持 tile 大小和 SRAM 大小不變的同時,掃描了矩陣的序列長度(如圖 5 所示)。
圖 5: 序列長度對 Flash Attention 數(shù)值偏差的影響。
如圖所示,隨著序列長度的增加,無論是通過(a)最大差異上限的測量,還是通過(b)差異的平均值和標準差的測量,F(xiàn)lash Attention 和 Baseline Attention 之間的數(shù)值偏差都在增加。
除此之外,研究者還利用微基準設計進行不同優(yōu)化的實驗,以便更好地了解數(shù)值偏差的影響(如圖 6 所示)。
圖 6a 顯示了調(diào)換 block 維數(shù)的順序如何導致 Flash Attention 和 Baseline Attention 之間的數(shù)值差異增大。圖 6b 中的其他擾動,比如限制 tile 大小為正方形,不會對數(shù)值偏差產(chǎn)生影響。圖 6c 表明了 block/tile 大小越大,數(shù)值偏差越小。
圖 6: 算法的改變及其對觀察到的數(shù)值偏差的影響。
通過權(quán)重差異來了解數(shù)值偏差
雖然在前向傳遞過程中,F(xiàn)lash Attention 可能會導致 Attention 輸出的數(shù)值偏差,但這項研究的最終目標是確定這是否會在模型訓練過程中產(chǎn)生任何影響,以研究它是否會導致訓練的不穩(wěn)定性。
因此,研究者希望量化 Flash Attention 是否在訓練過程中改變了模型,即上文觀察到的 Attention 輸出差異是否反映在訓練過程中更新的模型權(quán)重中。
研究者利用兩個指標來衡量使用 Baseline Attention 訓練的模型與使用 Flash Attention 訓練的模型之間的模型權(quán)重差異。首先計算最大差異,即找出權(quán)重矩陣之間差異的絕對值并取最大值,從而得出偏差的上限,如下所示:
雖然最大差值提供了數(shù)值偏差的上限,但它沒有考慮到每個矩陣的分布情況。因此,研究者通過 Wasserstein Distance 來量化權(quán)重差異,這是衡量張量之間相似性的常用度量。雖然在計算上稍顯復雜,但 Wasserstein Distance 包含了張量分布的形狀信息以衡量相似性。計算公式概述如下:
數(shù)值越低,表明矩陣之間的相似度越高。
利用這兩個指標,研究者隨后量化了在整個訓練過程中與 Baseline Attention 相比,F(xiàn)lash Attention 的模型權(quán)重是如何變化的:
根據(jù) Wasserstein Distance 和 Max Difference 這兩個指標,在整個訓練過程中,F(xiàn)lash Attention 的加入確實改變了模型權(quán)重,而且隨著訓練的繼續(xù),這種差異只會越來越大,這表明了使用 Flash Attention 訓練的模型與使用 Baseline Attention 訓練的相同模型收斂到了不同的模型。
然而,訓練是一個隨機過程,某些模型結(jié)構(gòu)的改變可能會在下游效應和準確性方面產(chǎn)生相似的結(jié)果。即使使用 Flash Attention 和 Baseline Attention 訓練的模型權(quán)重不同,這也是值得關(guān)注的。
完全訓練模型并評估準確性是一項成本昂貴且資源密集的任務,特別是對于訓練需要數(shù)月的大模型來說。
研究者通過配置一個 proxy 來探尋:
(a) 這些權(quán)重變化的意義有多大?
(b) 能否將其與其他廣泛采用的訓練優(yōu)化中的標準權(quán)重變化聯(lián)系起來?
為了實現(xiàn)這一目標,研究者設計了一系列實驗來比較在不同場景下,訓練過程中的權(quán)重差異是如何變化的。
除了對比使用 Flash Attention 和 Baseline Attention 的訓練過程外,他們還量化了在訓練開始時權(quán)重被初始化為不同隨機值的相同訓練過程中的權(quán)重差異。這提供了一個界限,因為隨機權(quán)重初始化是一種常用的技術(shù),并且通常會產(chǎn)生等效的結(jié)果。
此外,研究者還測量了使用不同精度訓練的模型權(quán)重的變化。數(shù)值精度(即 FP16 與 FP32)有可能導致下游變化,這作為確定了 Flash Attention 權(quán)重重要性的一個上限。
如圖 8 所示,可以發(fā)現(xiàn),使用 Flash Attention 的模型權(quán)重偏差變化率與不同模型初始化的權(quán)重偏差變化率相當或更小(注意紅色和藍色曲線的斜率)。
此外,使用 FP16 與 FP32 時的權(quán)重變化率比不同模型初始化時的權(quán)重變化率更高,變化也更大。
這些結(jié)果提供了一個 proxy,并表明:「雖然 Flash Attention 會出現(xiàn)數(shù)值偏差,但它會被隨機模型初始化和低精度訓練所限制。而且所引入的模型權(quán)重偏差大約是低精度訓練時的 1/2 至 1/5 倍?!?/span>
圖 8: 使用 Wasserstein Distance metric 測量的訓練過程中的相對權(quán)重差異。
更多研究細節(jié),可參考原論文。