Meta FAIR 聯(lián)合哈佛優(yōu)化大規(guī)模機(jī)器學(xué)習(xí)時(shí)產(chǎn)生的資料偏差，提供了新的研究架構(gòu)。

據(jù)所周知，大語(yǔ)言模型的訓(xùn)練常常需要數(shù)月的時(shí)間，使用數(shù)百甚至上千個(gè)GPU。以LLaMA2 70B模型為例，其訓(xùn)練總共需要1,720,320個(gè)GPU小時(shí)。由於這些工作負(fù)載的規(guī)模和複雜性，導(dǎo)致訓(xùn)練大模型存在著獨(dú)特的系統(tǒng)性挑戰(zhàn)。

最近，許多機(jī)構(gòu)在訓(xùn)練SOTA生成式AI模型時(shí)報(bào)告了訓(xùn)練過(guò)程中的不穩(wěn)定情況，它們通常以損失尖峰的形式出現(xiàn)，例如Google的PaLM模型訓(xùn)練過(guò)程中出現(xiàn)了多達(dá)20次的損失尖峰。

數(shù)值偏差是造成這種訓(xùn)練不準(zhǔn)確性的根因，由於大語(yǔ)言模型訓(xùn)練執(zhí)行成本極高，如何量化數(shù)值偏差儼然成為關(guān)鍵問(wèn)題。

在最新的一項(xiàng)工作中，來(lái)自 Meta、哈佛大學(xué)的研究者開(kāi)發(fā)了一個(gè)原則性定量方法來(lái)理解訓(xùn)練優(yōu)化中的數(shù)值偏差。以此評(píng)估不同的最新最佳化技術(shù)，並確定它們?cè)谟渺队?xùn)練大模型時(shí)是否可能引入意外的不穩(wěn)定性。 研究者發(fā)現(xiàn)，儘管現(xiàn)有的最佳化方法在一些任務(wù)上表現(xiàn)出色，但在大型模型上應(yīng)用時(shí)，會(huì)出現(xiàn)一些數(shù)值偏差。這種數(shù)值偏差可能會(huì)在訓(xùn)練過(guò)程中產(chǎn)生不穩(wěn)定性，導(dǎo)致模型的表現(xiàn)下降。 為了解決這個(gè)問(wèn)題，研究者提出了一種基於原則性量化方法的最佳化

• 論文標(biāo)題：Is Flash Attention Stable？
• 論文連結(jié)：https://arxiv.org/pdf/2405.02803

結(jié)果發(fā)現(xiàn)，在單獨(dú)的前向傳遞過(guò)程中，F(xiàn)lash Attention 的數(shù)值偏差比BF16 的Baseline Attention 大一個(gè)數(shù)量級(jí)。

具體而言，該方法包括兩個(gè)階段，包括：

• #開(kāi)發(fā)一個(gè)微基準(zhǔn)來(lái)擾動(dòng)給定最佳化中的數(shù)值精度；
• 透過(guò)基於Wasserstein 距離的資料驅(qū)動(dòng)分析評(píng)估數(shù)值偏差如何轉(zhuǎn)換為模型權(quán)重的變化。

研究者分析了 SOTA 最佳化技術(shù) Flash Attention，並量化了可可能引入的數(shù)值偏差。 Flash Attention 是一種廣泛用於加速注意力機(jī)制的技術(shù)，通常被認(rèn)為是 Transformer 模型中的系統(tǒng)瓶頸。 Flash Attention 在提高速度和減少記憶體存取量的同時(shí)，也依賴(lài)演算法最佳化，而演算法最佳化有可能導(dǎo)致數(shù)值偏差的增加。

研究者假設(shè)添加重新縮放因子（rescaling factors ）可能會(huì)引入無(wú)意的近似，導(dǎo)致數(shù)值折衷，這可能會(huì)在後續(xù)影響訓(xùn)練穩(wěn)定性。

他們?cè)诙嗄B(tài)文字到影像工作負(fù)載的背景下分析了 Flash Attention，以確定 Flash Attention 與其基準(zhǔn)之間數(shù)值偏差的潛在重要性。最終，他們引入了一個(gè)框架來(lái)量化訓(xùn)練優(yōu)化的數(shù)值偏差及其下游影響。

研究者在數(shù)值偏差量化上主要做出了以下兩點(diǎn)貢獻(xiàn)：

（1）設(shè)計(jì)了一個(gè)微基準(zhǔn)來(lái)分離數(shù)值精度對(duì)數(shù)值偏差的影響。

研究者所設(shè)計(jì)的微基準(zhǔn)作為一種技術(shù)，用於衡量和量化傳統(tǒng)黑盒最佳化（如 Flash Attention）所導(dǎo)致的數(shù)值偏差。透過(guò)擾動(dòng)通常在提供的內(nèi)核中不可用的方面，他們開(kāi)創(chuàng)性地發(fā)現(xiàn)在低數(shù)值精度（BF16）下，與 Baseline Attention 相比，F(xiàn)lash Attention 的數(shù)值偏差大約高出一個(gè)數(shù)量級(jí)。

（2）基於 Wasserstein Distance 度量進(jìn)行了資料驅(qū)動(dòng)的分析。

透過(guò)此分析，研究者將觀察到的數(shù)值偏差置於上下文，並為其對(duì)下游模型屬性的影響形成一個(gè)上限（upper bound）。在研究者的案例研究中，他們能夠限制觀察到的數(shù)值偏差的影響，並發(fā)現(xiàn)：「Flash Attention 引入的模型權(quán)重偏差大約為低精度訓(xùn)練的1/2 至1/5 倍。」

這項(xiàng)研究強(qiáng)調(diào)了開(kāi)發(fā)一種原則性方法的重要性：「不僅要量化，而且要將訓(xùn)練優(yōu)化對(duì)數(shù)值偏差的影響置於上下文中?！雇高^(guò)建立代理（proxies）來(lái)將數(shù)值偏差置於上下文中，旨在推斷通常難以衡量的下游模型效果（即訓(xùn)練不穩(wěn)定性）的可能性。

實(shí)驗(yàn)方法

研究者首先發(fā)展了一個(gè)微基準(zhǔn)來(lái)分離並研究 Flash Attention 所造成的數(shù)值偏差。如圖 2 所示，他們透過(guò)對(duì) Flash Attention 進(jìn)行數(shù)值上的重新實(shí)現(xiàn)，以分析不同的數(shù)值精度，並在演算法的每個(gè)步驟應(yīng)用潛在的最佳化措施。

圖 2: 微基準(zhǔn)設(shè)計(jì)摘要。

這是必要的，因?yàn)?Flash Attention 核心目前僅支援 FP16 和 BF16 數(shù)值格式。該核心也是 CUDA 程式碼的包裝 API 調(diào)用，這使得擾動(dòng)演算法以檢查數(shù)值偏差的影響變得具有挑戰(zhàn)性。

相比之下，他們的微基準(zhǔn)設(shè)計(jì)允許在演算法內(nèi)部進(jìn)行精度輸入和修改。研究者將微基準(zhǔn)與原始的 Flash Attention kernel 進(jìn)行了驗(yàn)證。

他們進(jìn)一步設(shè)計(jì)了一種技術(shù)，以比較模型執(zhí)行過(guò)程中每個(gè)步驟的 Attention 矩陣的輸出。並修改了模型程式碼，每次呼叫注意力時(shí)都計(jì)算 Baseline Attention 和 Flash Attention，這允許對(duì)相同的輸入矩陣進(jìn)行精確的輸出矩陣比較。

為了將其置於上下文中，研究者也透過(guò)相同和獨(dú)立的訓(xùn)練運(yùn)行，使用 Max difference 和 Wasserstein Distance 度量來(lái)量化模型權(quán)重在整個(gè)訓(xùn)練過(guò)程中的差異。

對(duì)於訓(xùn)練實(shí)驗(yàn)，研究者則使用一種將文字輸入轉(zhuǎn)換為圖像的生成式 AI workload（即文字到圖像模型）。他們使用 Shutterstock 資料集重新訓(xùn)練模型，並在一組英偉達(dá) 80GB A100 GPU 叢集上執(zhí)行此實(shí)驗(yàn)。

透過(guò)微基準(zhǔn)量化數(shù)值偏差

#研究者首先分析了 Flash Attention 在前向傳遞過(guò)程中的影響。他們利用微基準(zhǔn)測(cè)試，在隨機(jī)初始化查詢(xún)、鍵、值向量相同的情況下，檢驗(yàn)不同數(shù)值精確度對(duì) Attention 計(jì)算的輸出矩陣的影響。

如圖3 所示，當(dāng)研究者使用從BF16 到FP64 變化的不同數(shù)值格式時(shí)，F(xiàn)lash Attention 和Baseline Attention 之間的數(shù)值偏差隨著尾數(shù)位數(shù)的增加而減小。這表明數(shù)值差異是由於較少的尾數(shù)位數(shù)所固有的近似造成的。

圖 3：數(shù)值格式對(duì)於 Flash Attention 的數(shù)值偏差所產(chǎn)生的效果。

之後，研究者為進(jìn)行標(biāo)準(zhǔn)比較，在FP64 數(shù)值格式下的Baseline Attention 設(shè)定了「黃金值」，然後將不同數(shù)值格式下的Attention 輸出與該值進(jìn)行了比較（如圖4 所示）。

圖 4:FP64 下 Baseline Attention「黃金值」的比較。

結(jié)果表明，F(xiàn)lash Attention 的數(shù)值偏差大約是 BF16 下 Baseline 的 10 倍。

為了進(jìn)一步分析這種觀察到的數(shù)值偏差，研究者保持 tile 大小和 SRAM 大小不變的同時(shí)，掃描了矩陣的序列長(zhǎng)度（如圖 5 所示）。

圖 5: 序列長(zhǎng)度對(duì) Flash Attention 數(shù)值偏差的影響。

如圖所示，隨著序列長(zhǎng)度的增加，無(wú)論是透過(guò)（a）最大差異上限的測(cè)量，或是透過(guò)（b）差異的平均值和標(biāo)準(zhǔn)差的測(cè)量，F(xiàn)lash Attention和Baseline Attention 之間的數(shù)值偏差都在增加。

除此之外，研究者也利用微基準(zhǔn)設(shè)計(jì)進(jìn)行不同最佳化的實(shí)驗(yàn)，以便更了解數(shù)值偏差的影響（如圖 6 所示）。

圖 6a 顯示了調(diào)換 block 維數(shù)的順序如何導(dǎo)致 Flash Attention 和 Baseline Attention 之間的數(shù)值差異增大。圖 6b 中的其他擾動(dòng)，例如限制 tile 大小為正方形，不會(huì)對(duì)數(shù)值偏差產(chǎn)生影響。圖 6c 顯示了 block/tile 大小越大，數(shù)值偏差越小。

圖 6: 演算法的改變及其對(duì)觀察到的數(shù)值偏差的影響。

透過(guò)權(quán)重差異來(lái)了解數(shù)值偏差

雖然在前向傳遞過(guò)程中，F(xiàn)lash Attention 可能會(huì)導(dǎo)致Attention 輸出的數(shù)值偏差，但這項(xiàng)研究的最終目標(biāo)是確定這是否會(huì)在模型訓(xùn)練過(guò)程中產(chǎn)生任何影響，以研究它是否會(huì)導(dǎo)致訓(xùn)練的不穩(wěn)定性。

因此，研究者希望量化 Flash Attention 是否在訓(xùn)練過(guò)程中改變了模型，即上文觀察到的 Attention 輸出差異是否反映在訓(xùn)練過(guò)程中更新的模型權(quán)重中。

研究者利用兩個(gè)指標(biāo)來(lái)測(cè)量使用 Baseline Attention 訓(xùn)練的模型與使用 Flash Attention 訓(xùn)練的模型之間的模型權(quán)重差異。首先計(jì)算最大差異，即找出權(quán)重矩陣之間差異的絕對(duì)值並取最大值，從而得出偏差的上限，如下所示：

雖然最大差值提供了數(shù)值偏差的上限，但它沒(méi)有考慮到每個(gè)矩陣的分佈。因此，研究者透過(guò) Wasserstein Distance 來(lái)量化權(quán)重差異，這是衡量張量之間相似性的常用度量。雖然在計(jì)算上稍微複雜，但 Wasserstein Distance 包含了張量分佈的形狀資訊以衡量相似性。計(jì)算公式概述如下：

數(shù)值越低，表示矩陣之間的相似度越高。

利用這兩個(gè)指標(biāo)，研究者隨後量化了在整個(gè)訓(xùn)練過(guò)程中與Baseline Attention 相比，F(xiàn)lash Attention 的模型權(quán)重是如何變化的：

#根據(jù)Wasserstein Distance 和Max Difference 這兩個(gè)指標(biāo)，在整個(gè)訓(xùn)練過(guò)程中，F(xiàn)lash Attention 的加入確實(shí)改變了模型權(quán)重，而且隨著訓(xùn)練的繼續(xù)，這種差異只會(huì)越來(lái)越大，這顯示了使用Flash Attention 訓(xùn)練的模型與使用Baseline Attention 訓(xùn)練的相同模型收斂到了不同的模型。

然而，訓(xùn)練是一個(gè)隨機(jī)過(guò)程，某些模型結(jié)構(gòu)的改變可能會(huì)在下游效應(yīng)和準(zhǔn)確性方面產(chǎn)生相似的結(jié)果。即使使用 Flash Attention 和 Baseline Attention 訓(xùn)練的模型權(quán)重不同，這也是值得關(guān)注的。

完全訓(xùn)練模型並評(píng)估準(zhǔn)確性是一項(xiàng)昂貴且資源密集的任務(wù)，特別是對(duì)於訓(xùn)練需要數(shù)月的大模型來(lái)說(shuō)。

研究者透過(guò)配置一個(gè) proxy 來(lái)探索：

(a) 這些權(quán)重變化的意義有多大？

(b) 能否將其與其他廣泛採(cǎi)用的訓(xùn)練優(yōu)化中的標(biāo)準(zhǔn)權(quán)重變化聯(lián)繫起來(lái)？

為了實(shí)現(xiàn)這一目標(biāo)，研究者設(shè)計(jì)了一系列實(shí)驗(yàn)來(lái)比較在不同場(chǎng)景下，訓(xùn)練過(guò)程中的權(quán)重差異是如何變化的。

除了對(duì)比使用 Flash Attention 和 Baseline Attention 的訓(xùn)練過(guò)程外，他們還量化了在訓(xùn)練開(kāi)始時(shí)權(quán)重被初始化為不同隨機(jī)值的相同訓(xùn)練過(guò)程中的權(quán)重差異。這提供了一個(gè)界限，因?yàn)殡S機(jī)權(quán)重初始化是一種常用的技術(shù)，並且通常會(huì)產(chǎn)生等效的結(jié)果。

此外，研究者也測(cè)量了使用不同精度訓(xùn)練的模型權(quán)重的變化。數(shù)值精確度（即 FP16 與 FP32）有可能導(dǎo)致下游變化，這作為確定了 Flash Attention 權(quán)重重要性的一個(gè)上限。

如圖8 所示，可以發(fā)現(xiàn)，使用Flash Attention 的模型權(quán)重偏差變化率與不同模型初始化的權(quán)重偏差變化率相當(dāng)或更?。ㄗ⒁饧t色和藍(lán)色曲線的斜率）。

此外，使用 FP16 與 FP32 時(shí)的權(quán)重變化率比不同模型初始化時(shí)的權(quán)重變化率更高，變化也更大。

這些結(jié)果提供了一個(gè)proxy，並表明：「雖然Flash Attention 會(huì)出現(xiàn)數(shù)值偏差，但它會(huì)被隨機(jī)模型初始化和低精度訓(xùn)練所限制。而且所引入的模型權(quán)重偏差大約是低精度訓(xùn)練時(shí)的1/2 至1/5 倍。相對(duì)權(quán)重差異。

更多研究細(xì)節(jié)，可參考原文。