70倍極致壓縮！大模型的檢查點再多也不怕

AIxiv專欄是本站發布學術、技術內容的欄位。過去數年，本站AIxiv專欄接收通報了2,000多篇內容，涵蓋全球各大專院校與企業的頂尖實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或聯絡報道。投稿信箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

該論文的作者皆來自於華為諾亞實驗室，第一作者為李文碩，通訊作者為老櫞。相關團隊團隊近年來在ICML、CVPR、NeurIPS、ICCV、ECCV等頂會上有多項代表性工作發表，在高效大語言模型、視覺模型等領域都有豐富的成果產出，和知名大學和科學研究機構合作廣泛。

大模型作為當下 AI 工業界和學術界當之無愧的“流量之王”，吸引了大批學者和企業投入資源去研究與訓練。隨著規模越做越大，系統和工程問題就成了大模型訓練中繞不開的難題。例如在 Llama3.1 54 天的訓練裡，系統會崩潰 466 次，平均 2.78 小時一次！

那麼，頻繁儲存檢查點就顯得十分必要。但儲存檢查點本身也是一個大工程。

Meta 做了很多努力來加速儲存檢查點時間和增加儲存頻率，來對抗頻繁出現的系統失效。但頻繁儲存也意味著大量的儲存資源開銷，其訓練叢集配備了 240PB 的 SSD 來應對這項挑戰，光儲存這項的耗費就要億元！

華為諾亞的 ExCP 方法也就應運而生，為了應對存儲帶來的巨大開銷，他們提出了極致壓縮檢查點技術，能夠無損壓縮模型 70 倍，大幅降低訓練中的存儲開銷。

程式碼目前已經開源，在 Apache 2.0 框架下發布，issue 中已經有小夥伴成功復現了結果。

• 文章地址：https://arxiv.org/abs/2406.11257
• 倉庫地址：https://github.com/Gaffey/Ex創新性，文章中提到了兩個重要的概念，一個是利用訓練中檢查點的殘差信息，通過時間序列上信息的稀疏性實現更高的剪枝比例；另一個是將優化器和權重聯合起來進行壓縮，實現整體的高壓縮率。

具體方法

1. 檢查點殘差

化檢查迭代時梯度更新的總和，這部分是相對稀疏的，包含的資訊量較少，因此對這一殘差進行壓縮，可以獲得更好的壓縮比例。而與此相反的，優化器中儲存的動量是梯度一階矩和二階矩的滑動平均值，對於一階矩來說，它的滑動平均預設的參數是0.9，在數百到數千個迭代之後與上一次檢查點儲存的內容已經沒有太大的關聯，所以對於最佳化器直接壓縮本身的值而非殘差。最終待壓縮的檢查點表示為

2. 權重- 最佳化器動量聯合壓縮

目前已有的模型壓縮相關的工作一般只關注於模型的推理性能，或者是模型最終檢查點的大小，而不關注模型在整個訓練過程中對儲存空間的開銷。因而已有工作只對權重進行壓縮，而忽略了 Adam 等常見優化器中實際上儲存了兩倍於權重數量的動量。這項工作一方面將兩者一起進行了壓縮，顯著提升了整體的壓縮比例；另一方面也利用了權重和優化器動量的關聯性，進一步提升彼此的壓縮比例。

權重剪枝：由於剪枝的權重是殘差值，優化器動量的二階矩可以大致表示在過去一段時間內權重殘差值的變化幅度，所以可以使用優化器動量的二階矩作為指標來決定不同層的剪枝比例。剪枝策略如下文公式所示

式中，W 和分別表示權重和二階矩。

優化器動量剪枝：對於動量剪枝，可以使用一階矩作為指示器來進行剪枝，論文中有關於可收斂性的一個簡要證明。同時，如果一個位置的權重已經被剪枝，那麼對應位置的最佳化器動量也應該同步被處理，所以剪枝策略如下文公式所示

式中， 表示一階矩。

3. 整體壓縮流程

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

而恢復出檢查點完整文件的流程則如Algorithm 2 所示，進行解壓縮之後，首先從非均勻量化後儲存的碼本和下標中恢復出浮點結果，然後再與基準權重（上一個檢查點的原始權重或恢復出的重建權重）相加，得到檢查點完整檔案。而恢復出整個訓練流程中的檢查點檔案的流程如Algorithm 3 所示，在完成訓練後只保存初始化權重的隨機種子和每個檢查點儲存的壓縮結果，然後依序對檢查點進行恢復以獲得完整的檢查點序列，以供從其中選擇某個或多個檢查點恢復訓練/ 進行測試等。

實驗結果

文章中不僅對於大語言模型做了評估，在 ViT-L32 這樣的視覺模型上這一方法也能取得較大的效果。

從消融實驗裡也可以看出，採用殘差剪枝的方法大大減少了剪枝帶來的損失。

文章中也提供了大語言模型壓縮前後問答的範例，可以看到壓縮本身對於模型的問答能力也沒有造成損害。

以上是70倍極致壓縮！大模型的檢查點再多也不怕的詳細內容。更多資訊請關注PHP中文網其他相關文章！