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