70倍極致壓縮！大模型的檢查點再多也不怕
該論文的作者均來自于華為諾亞實驗室，第一作者為李文碩，通訊作者為王云鶴和陳醒濠。相關團隊團隊近年來在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/ExCP
方法也很有創新性，文章中提到了兩個重要的概念，一個是利用訓練中檢查點的殘差信息，通過時間序列上信息的稀疏性實現更高的剪枝比例；另一個是將優化器和權重聯合起來進行壓縮，實現整體的高壓縮率。
具體方法
1. 檢查點殘差
在訓練過程中，當前的參數可以看作上一個檢查點存儲的權重加上逐次迭代時梯度更新的總和，這部分是相對稀疏的，包含的信息量較少，因此對這一殘差進行壓縮，可以獲得更好的壓縮比例。而與此相反的，優化器中存儲的動量是梯度一階矩和二階矩的滑動平均值，對于一階矩來說，它的滑動平均默認的參數是 0.9，在數百到數千個迭代之后與上一次檢查點存儲的內容已經沒有太大的關聯，所以對于優化器直接壓縮其本身的值而非殘差。最終待壓縮的檢查點表示為
2. 權重 - 優化器動量聯合壓縮
目前已有的模型壓縮相關的工作一般只關注于模型的推理性能，或者是模型最終存儲檢查點的大小，而不關注模型在整個訓練過程中對儲存空間的開銷。因而已有工作只對權重進行壓縮，而忽略了 Adam 等常見優化器中實際上存儲了兩倍于權重數量的動量。這一工作一方面將兩者一起進行了壓縮，顯著提升了整體的壓縮比例；另一方面也利用了權重和優化器動量的關聯性，進一步提升彼此的壓縮比例。
權重剪枝：由于剪枝的權重是殘差值，優化器動量的二階矩可以大致表示在過去一段時間內權重殘差值的變化幅度，所以可以使用優化器動量的二階矩作為指標來確定不同層的剪枝比例。剪枝策略如下文公式所示
式中，W 和
分別表示權重和二階矩。
優化器動量剪枝：對于動量剪枝，可以使用一階矩作為指示器來進行剪枝，論文中有關于可收斂性的一個簡要證明。同時，如果一個位置的權重已經被剪枝，那么對應位置的優化器動量也應該同步被處理，所以剪枝策略如下文公式所示
式中，
表示一階矩。
3. 整體壓縮流程
整體壓縮流程如 Algorithm 1 所示，依次進行計算權重殘差 / 聯合壓縮 / 非均勻量化 / 編碼壓縮等步驟，得到最終的壓縮結果。
而恢復出檢查點完整文件的流程則如 Algorithm 2 所示，進行解壓縮之后，首先從非均勻量化后存儲的碼本和下標中恢復出浮點結果，然后再與基準權重（上一個檢查點的原始權重或恢復出的重建權重）相加，得到檢查點完整文件。而恢復出整個訓練流程中的檢查點文件的流程如 Algorithm 3 所示，在完成訓練后只保存初始化權重的隨機種子和每個檢查點存儲的壓縮結果，然后依次對檢查點進行恢復以得到完整的檢查點序列，以供從其中選擇某個或多個檢查點恢復訓練 / 進行測試等。
實驗結果
文章中不僅對于大語言模型做了評估，在 ViT-L32 這樣較大的視覺模型上這一方法也能取得很好的效果。
從消融實驗里也可以看出，采用殘差剪枝的方法大大減少了剪枝帶來的損失。
文章中還提供了大語言模型壓縮前后問答的樣例，可以看到壓縮本身對于模型的問答能力也沒有造成損害。
