該論文的作者均來自于華為諾亞實(shí)驗(yàn)室，第一作者為李文碩，通訊作者為王云鶴和陳醒濠。相關(guān)團(tuán)隊(duì)團(tuán)隊(duì)近年來在ICML、CVPR、NeurIPS、ICCV、ECCV等頂會(huì)上有多項(xiàng)代表性工作發(fā)表，在高效大語言模型、視覺模型等領(lǐng)域都有豐富的成果產(chǎn)出，和知名高校和科研機(jī)構(gòu)合作廣泛。

大模型作為當(dāng)下 AI 工業(yè)界和學(xué)術(shù)界當(dāng)之無愧的「流量之王」，吸引了大批學(xué)者和企業(yè)投入資源去研究與訓(xùn)練。隨著規(guī)模越做越大，系統(tǒng)和工程問題已經(jīng)成了大模型訓(xùn)練中繞不開的難題。例如在 Llama3.1 54 天的訓(xùn)練里，系統(tǒng)會(huì)崩潰 466 次，平均 2.78 小時(shí)一次！

那么，頻繁存儲(chǔ)檢查點(diǎn)就顯得十分必要。但存儲(chǔ)檢查點(diǎn)本身也是一個(gè)大工程。

Meta 做了很多努力來加速存儲(chǔ)檢查點(diǎn)時(shí)間和增加存儲(chǔ)頻率，來對(duì)抗頻繁出現(xiàn)的系統(tǒng)失效。但頻繁存儲(chǔ)也意味著大量的存儲(chǔ)資源開銷，其訓(xùn)練集群配備了 240PB 的 SSD 來應(yīng)對(duì)這一挑戰(zhàn)，光存儲(chǔ)這一項(xiàng)的耗費(fèi)就要億元！

華為諾亞的 ExCP 方法也就應(yīng)運(yùn)而生，為了應(yīng)對(duì)存儲(chǔ)帶來的巨大開銷，他們提出了極致壓縮檢查點(diǎn)技術(shù)，能夠無損壓縮模型 70 倍，大幅降低訓(xùn)練中的存儲(chǔ)開銷。

代碼目前已經(jīng)開源，在 Apache 2.0 框架下發(fā)布，issue 中已經(jīng)有小伙伴成功復(fù)現(xiàn)了結(jié)果。

• 文章地址：https://arxiv.org/abs/2406.11257
• 倉庫地址：https://github.com/Gaffey/ExCP

方法也很有創(chuàng)新性，文章中提到了兩個(gè)重要的概念，一個(gè)是利用訓(xùn)練中檢查點(diǎn)的殘差信息，通過時(shí)間序列上信息的稀疏性實(shí)現(xiàn)更高的剪枝比例；另一個(gè)是將優(yōu)化器和權(quán)重聯(lián)合起來進(jìn)行壓縮，實(shí)現(xiàn)整體的高壓縮率。

具體方法

1. 檢查點(diǎn)殘差

在訓(xùn)練過程中，當(dāng)前的參數(shù)可以看作上一個(gè)檢查點(diǎn)存儲(chǔ)的權(quán)重加上逐次迭代時(shí)梯度更新的總和，這部分是相對(duì)稀疏的，包含的信息量較少，因此對(duì)這一殘差進(jìn)行壓縮，可以獲得更好的壓縮比例。而與此相反的，優(yōu)化器中存儲(chǔ)的動(dòng)量是梯度一階矩和二階矩的滑動(dòng)平均值，對(duì)于一階矩來說，它的滑動(dòng)平均默認(rèn)的參數(shù)是 0.9，在數(shù)百到數(shù)千個(gè)迭代之后與上一次檢查點(diǎn)存儲(chǔ)的內(nèi)容已經(jīng)沒有太大的關(guān)聯(lián)，所以對(duì)于優(yōu)化器直接壓縮其本身的值而非殘差。最終待壓縮的檢查點(diǎn)表示為

2. 權(quán)重 - 優(yōu)化器動(dòng)量聯(lián)合壓縮

目前已有的模型壓縮相關(guān)的工作一般只關(guān)注于模型的推理性能，或者是模型最終存儲(chǔ)檢查點(diǎn)的大小，而不關(guān)注模型在整個(gè)訓(xùn)練過程中對(duì)儲(chǔ)存空間的開銷。因而已有工作只對(duì)權(quán)重進(jìn)行壓縮，而忽略了 Adam 等常見優(yōu)化器中實(shí)際上存儲(chǔ)了兩倍于權(quán)重?cái)?shù)量的動(dòng)量。這一工作一方面將兩者一起進(jìn)行了壓縮，顯著提升了整體的壓縮比例；另一方面也利用了權(quán)重和優(yōu)化器動(dòng)量的關(guān)聯(lián)性，進(jìn)一步提升彼此的壓縮比例。

權(quán)重剪枝：由于剪枝的權(quán)重是殘差值，優(yōu)化器動(dòng)量的二階矩可以大致表示在過去一段時(shí)間內(nèi)權(quán)重殘差值的變化幅度，所以可以使用優(yōu)化器動(dòng)量的二階矩作為指標(biāo)來確定不同層的剪枝比例。剪枝策略如下文公式所示

式中，W 和分別表示權(quán)重和二階矩。

優(yōu)化器動(dòng)量剪枝：對(duì)于動(dòng)量剪枝，可以使用一階矩作為指示器來進(jìn)行剪枝，論文中有關(guān)于可收斂性的一個(gè)簡(jiǎn)要證明。同時(shí)，如果一個(gè)位置的權(quán)重已經(jīng)被剪枝，那么對(duì)應(yīng)位置的優(yōu)化器動(dòng)量也應(yīng)該同步被處理，所以剪枝策略如下文公式所示

式中， 表示一階矩。

3. 整體壓縮流程

整體壓縮流程如 Algorithm 1 所示，依次進(jìn)行計(jì)算權(quán)重殘差 / 聯(lián)合壓縮 / 非均勻量化 / 編碼壓縮等步驟，得到最終的壓縮結(jié)果。

而恢復(fù)出檢查點(diǎn)完整文件的流程則如 Algorithm 2 所示，進(jìn)行解壓縮之后，首先從非均勻量化后存儲(chǔ)的碼本和下標(biāo)中恢復(fù)出浮點(diǎn)結(jié)果，然后再與基準(zhǔn)權(quán)重（上一個(gè)檢查點(diǎn)的原始權(quán)重或恢復(fù)出的重建權(quán)重）相加，得到檢查點(diǎn)完整文件。而恢復(fù)出整個(gè)訓(xùn)練流程中的檢查點(diǎn)文件的流程如 Algorithm 3 所示，在完成訓(xùn)練后只保存初始化權(quán)重的隨機(jī)種子和每個(gè)檢查點(diǎn)存儲(chǔ)的壓縮結(jié)果，然后依次對(duì)檢查點(diǎn)進(jìn)行恢復(fù)以得到完整的檢查點(diǎn)序列，以供從其中選擇某個(gè)或多個(gè)檢查點(diǎn)恢復(fù)訓(xùn)練 / 進(jìn)行測(cè)試等。

實(shí)驗(yàn)結(jié)果

文章中不僅對(duì)于大語言模型做了評(píng)估，在 ViT-L32 這樣較大的視覺模型上這一方法也能取得很好的效果。

從消融實(shí)驗(yàn)里也可以看出，采用殘差剪枝的方法大大減少了剪枝帶來的損失。

文章中還提供了大語言模型壓縮前后問答的樣例，可以看到壓縮本身對(duì)于模型的問答能力也沒有造成損害。