大模型訓(xùn)練loss突刺原因和解決辦法

最近閱讀了《A Theory on Adam Instability in Large-Scale Machine Learning 》這篇論文。比較全面的闡述了100B以上的大模型預(yù)訓(xùn)練中出現(xiàn)loss spike的原因(loss 突然大幅度上漲)，并介紹了一些可能的解決辦法。論文寫的非常精彩，但整體上有點(diǎn)散和深，我嘗試著站在工業(yè)立場上把它串一下

突刺是什么

首先介紹一下什么是loss spike：

loss spike指的是預(yù)訓(xùn)練過程中，尤其容易在大模型（100B以上）預(yù)訓(xùn)練過程中出現(xiàn)的loss突然暴漲的情況

如圖所示模型訓(xùn)練過程中紅框中突然上漲的loss尖峰loss spike的現(xiàn)象會(huì)導(dǎo)致一系列的問題發(fā)生，譬如模型需要很長時(shí)間才能再次回到spike之前的狀態(tài)（論文中稱為pre-explosion），或者更嚴(yán)重的就是loss再也無法drop back down，即模型再也無法收斂

PaLM和GLM130b之前的解決辦法是找到loss spike之前最近的checkpoint，更換之后的訓(xùn)練樣本來避免loss spike的出現(xiàn)。

突刺成因分析

這篇論文（以下稱本文）對(duì)loss spike的出現(xiàn)原因做了十分詳細(xì)的分析，最后認(rèn)為預(yù)訓(xùn)練使用的Adam優(yōu)化器是導(dǎo)致這個(gè)現(xiàn)象出現(xiàn)的重要原因之一

首先回顧一下Adam優(yōu)化器的結(jié)構(gòu)（這里介紹的是較為傳統(tǒng)的Adam優(yōu)化器，現(xiàn)在nlp任務(wù)更偏向于使用帶有正則化項(xiàng)的Adamw變體）：

其中均為超參數(shù)( ,防止除0)，表示第t次更新的梯度， 的初始值 為0。

本文首先對(duì)Adam的有效性做了論述，其本質(zhì)在于證明了Adam優(yōu)化過程是對(duì)牛頓下降法（二階導(dǎo)）的一個(gè)有效逼近，因此在收斂速度上大幅度領(lǐng)先傳統(tǒng)SGD(一階導(dǎo))，證明過程不做贅述，可以參考本文和Adam系列相關(guān)論文

Adam算法是牛頓下降法的一個(gè)迭代逼近

一切顯得十分完美，但是理想很豐滿，現(xiàn)實(shí)很骨感，收斂過程并不是一帆風(fēng)順的

首先我們想象一下這個(gè)更新參數(shù)的變化趨勢

在 的時(shí)候， ，， 集中在 附近。而在訓(xùn)練到最后，假設(shè)模型收斂到某個(gè)最優(yōu)點(diǎn)，此時(shí) 應(yīng)該集中在0附近。也就是說，我們似乎可以把更新參數(shù)的變化過程想象成為一個(gè)從兩端（非穩(wěn)態(tài)）向中間（穩(wěn)態(tài)）收攏的過程。實(shí)際的觀察現(xiàn)象也是如此：

非穩(wěn)態(tài)

中間態(tài)

穩(wěn)態(tài)

進(jìn)入正態(tài)分布的穩(wěn)態(tài)之后，理想的更新參數(shù)變化趨勢應(yīng)該是方差越來越小，所有更新參數(shù)逐漸向0靠近。這應(yīng)該是一個(gè)單向的過程，即穩(wěn)定的單峰狀態(tài)（unimodal）不會(huì)再次進(jìn)入非穩(wěn)定的雙峰狀態(tài)(bimodal)，但事實(shí)并非如此，更新參數(shù)會(huì)再次進(jìn)入非穩(wěn)定的雙峰狀態(tài)

本文在理論層面做了研究和解釋，從中心極限定理（可以結(jié)合道爾頓板實(shí)驗(yàn)理解）出發(fā)，認(rèn)為隨機(jī)事件的疊加進(jìn)入單峰的正態(tài)分布的必要條件之一是各個(gè)隨機(jī)事件事件之間應(yīng)該是相互獨(dú)立的，但是梯度變化以及更新參數(shù)的變化并不能特別好的滿足獨(dú)立性這一條件，而這一點(diǎn)恰恰是導(dǎo)致更新參數(shù)振蕩，loss spike出現(xiàn)以及l(fā)oss 不收斂的重要原因之一

造成梯度變化不獨(dú)立的原因（1、淺層參數(shù)長時(shí)間不更新2、batch太大，后期梯度更新趨于平穩(wěn)）上述的理論有些晦澀，本文作者可能也了解這一點(diǎn)，之后開始直接點(diǎn)題，結(jié)合實(shí)驗(yàn)觀察拋出了重要現(xiàn)象和結(jié)論

即訓(xùn)練過程中l(wèi)oss spike的出現(xiàn)與：梯度更新幅度， 大小，batch大小這三個(gè)條件密切相關(guān)

本文作者對(duì)loss spike出現(xiàn)時(shí)模型的前后變化做了仔細(xì)拆解，發(fā)現(xiàn)下列一系列連續(xù)現(xiàn)象的出現(xiàn)導(dǎo)致了loss spike：

當(dāng)前模型處在穩(wěn)態(tài)（健康狀態(tài)），即單峰的正態(tài)分布狀態(tài)，并且梯度值 ，此時(shí)loss平穩(wěn)，訓(xùn)練過程正常

2.模型淺層(embedding層)梯度 ，這一般是由于訓(xùn)練一段時(shí)間之后，淺層的語義知識(shí)表示此時(shí)一般已經(jīng)學(xué)習(xí)的較好。但此時(shí)深層網(wǎng)絡(luò)（對(duì)應(yīng)復(fù)雜任務(wù)）的梯度更新還是相對(duì)較大

3.一段時(shí)間淺層(embedding層)梯度 之后會(huì)導(dǎo)致 ， 。此時(shí)趨于0。因此導(dǎo)致淺層參數(shù)得不到更新（也對(duì)應(yīng)于上述參數(shù)更新事件不獨(dú)立的原因）

4.此時(shí)雖然淺層(embedding層)參數(shù)長時(shí)間不更新，但是深層的參數(shù)依然一直在更新。長時(shí)間這樣的狀態(tài)之后，batch之間的樣本分布變化可能就會(huì)直接導(dǎo)致淺層(embedding層)再次出現(xiàn)較大的梯度變化（可以想象成一個(gè)水壩蓄水太久終于被沖開了。至于小模型為什么不會(huì)出現(xiàn)這種情況，推測是小模型函數(shù)空間小，無法捕獲樣本的分布變化，越大規(guī)模的模型對(duì)樣本之間不同維度的特征分布變化越敏感），此時(shí) ， 再次集中在 附近（此時(shí) , ），變成雙峰的非穩(wěn)定狀態(tài)，本文提到了淺層(embedding層)這種突然的參數(shù)變化可能造成模型的連鎖反應(yīng)進(jìn)而出現(xiàn)loss spike的現(xiàn)象（這也對(duì)應(yīng)了更換樣本重新訓(xùn)練有可能會(huì)減少loss spike的出現(xiàn)頻率，實(shí)際上就是選擇分布變化較小的樣本，減小淺層梯度變換幅度）

5.這個(gè)階段模型處于非穩(wěn)態(tài)，梯度變化幅度較大，每一次的梯度變化和更新參數(shù)變化事件之間又出現(xiàn)了一定的獨(dú)立性，因此經(jīng)過一定的時(shí)間之后模型有可能再次進(jìn)入穩(wěn)態(tài)，loss再次drop back down（注意，本文著重提了這個(gè)再次drop back down并不是一定出現(xiàn)的，也很有可能loss長期處于flat狀態(tài)，再也無法收斂）

因此我們得出一些結(jié)論，loss spike的出現(xiàn)和淺層的梯度更新幅度， 大小密切相關(guān)（batch大小帶來的相關(guān)性問題倒是顯得沒那么大說服力），實(shí)際上就是淺層網(wǎng)絡(luò)參數(shù)突然進(jìn)入到了之前長時(shí)間不在的狀態(tài)與模型深層參數(shù)當(dāng)前的狀態(tài)形成了連鎖反應(yīng)造成了模型進(jìn)入非穩(wěn)態(tài)。同時(shí)一般情況即使出現(xiàn)loss spike也會(huì)自動(dòng)回復(fù)到正常狀態(tài)，但也有可能再也不會(huì)

突刺解法

本文最后提到了防止loss spike出現(xiàn)的一些方法：

1.如之前提到的PaLM和GLM130B提到的出現(xiàn)loss spike后更換batch樣本的方法（常規(guī)方法，但是成本比較高）

2.減小learning rate，這是個(gè)治標(biāo)不治本的辦法，對(duì)更新參數(shù)的非穩(wěn)態(tài)沒有做改進(jìn)

3.減小 大小?；蛘咧苯影?設(shè)為0，重新定義

在等于0時(shí)候的值（這應(yīng)該是個(gè)值得嘗試的辦法）

值得一提的是智譜華章在本文發(fā)表之前，在去年的GLM130B訓(xùn)練時(shí)似乎也觀察到了淺層梯度變化和loss spike相關(guān)這一現(xiàn)象（GLM-130B: An Open Bilingual Pre-trained Model），他采取的是把淺層梯度直接乘以縮放系數(shù) 來減小淺層梯度更新值

出自130b

其實(shí)這塊我有個(gè)自己的想法，和是否也可以做衰減，隨著訓(xùn)練過程逐漸減小，來避免loss spike的現(xiàn)象

另外假設(shè)我們能一次性加載所有樣本進(jìn)行訓(xùn)練（實(shí)際上不可能做到），是否還會(huì)出現(xiàn)loss spike的現(xiàn)象

最后目前流行的fp8，fp16混合訓(xùn)練，如果upscale設(shè)置的過小，導(dǎo)致梯度在進(jìn)入優(yōu)化器之前就下溢，是不是會(huì)增加淺層梯度長時(shí)間不更新的可能性，進(jìn)而增加loss spike的出現(xiàn)的頻率。（這么看來似乎提升upscale大小以及優(yōu)化 大小是進(jìn)一步提升模型效果的一個(gè)思路）

審核編輯：黃飛