©作者 | 陳醒濠
單位 | 華為諾亞方舟實驗室
大型語言模型最近成為人工智慧領域的研究熱點，然而它們的訓練過程耗費巨大的計算和儲存資源。因此，高效壓縮儲存模型的checkpoint檔案顯得尤為關鍵。我們提出了一種創新的模型 checkpoint 壓縮方案（ExCP），該方案能夠在保持效能幾乎不受損失的前提下，顯著降低訓練過程中的儲存開銷。
首先，我們透過計算相鄰 checkpoint 的殘差值來捕獲關鍵但稀疏的資訊，從而提高壓縮比。接下來，我們實行權重和最佳化器動量的聯合壓縮，利用最佳化器儲存的動量資訊實現更高效的壓縮，並透過壓縮最佳化器動量進一步去除 checkpoint 中的冗餘引數。
最終，我們採用非均勻量化及編碼壓縮策略進一步減小 checkpoint 檔案的儲存尺寸。在從 410M 到 7B 不同規模的模型上測試我們的壓縮框架，特別是在 Pythia-410M 模型上，我們實現了近乎無損的效能下達到約70倍的壓縮比例。
論文標題：
ExCP: Extreme LLM Checkpoint Compression via Weight-Momentum Joint Shrinking
論文地址：
https://arxiv.org/abs/2406.11257
程式碼地址：
https://github.com/Gaffey/ExCP
方法
1.1 Checkpoint殘差
在訓練過程中，當前的引數可以看作上一個checkpoint儲存的權重加上逐次迭代時梯度更新的總和，這部分是相對稀疏的，包含的資訊量較少，因此我們對這一殘差進行壓縮，可以獲得更好的壓縮比例。
而與此相反的，最佳化器中儲存的動量是梯度一階矩和二階矩的滑動平均值，對於一階矩來說，它的滑動平均預設的引數是0.9，在數百到數千個迭代之後與上一次checkpoint儲存的內容已經沒有太大的關聯，所以我們對於最佳化器直接壓縮其本身的值而非殘差。最終待壓縮的checkpoint我們表示為:
1.2 權重-最佳化器動量聯合壓縮
目前已有的模型壓縮相關的工作中，大家一般只關注於模型的推理效能，或者是模型最終儲存checkpoint的大小，而不關注模型在整個訓練過程中對儲存空間的開銷。因而已有工作只對權重進行壓縮，而忽略了Adam等常見最佳化器中實際上儲存了兩倍於權重數量的動量。我們的工作一方面將兩者一起進行了壓縮，顯著提升了整體的壓縮比例；另一方面也利用了權重和最佳化器動量的關聯性，進一步提升彼此的壓縮比例。
權重剪枝：由於我們剪枝的權重是殘差值，最佳化器動量的二階矩可以大致表示在過去一段時間內權重殘差值的變化幅度，所以我們可以使用最佳化器動量的二階矩作為指標來確定不同層的剪枝比例。我們的剪枝策略如下文公式所示：
式中， 和 分別表示權重和二階矩。
最佳化器動量剪枝：對於動量剪枝，我們可以使用一階矩作為指示器來進行剪枝，論文中有關於可收斂性的一個簡要證明。同時，如果一個位置的權重已經被剪枝，那麼對應位置的最佳化器動量也應該同步被處理，所以我們的剪枝策略如下文公式所示：
式中， 表示一階矩。
1.3 整體壓縮流程
我們的整體壓縮流程如Algorithm 1所示，依次進行計算權重殘差/聯合壓縮/非均勻量化/編碼壓縮等步驟，得到最終的壓縮結果。
而恢復出checkpoint完整檔案的流程則如Algorithm 2所示，進行解壓縮之後，首先從非均勻量化後儲存的碼本和下標中恢復出浮點結果，然後再與基準權重（上一個checkpoint的原始權重或恢復出的重建權重）相加，得到checkpoint完整檔案。
而恢復出整個訓練流程中的checkpoint檔案的流程如Algorithm 3所示，我們在完成訓練後只儲存初始化權重的隨機種子和每個checkpoint儲存的壓縮結果，然後依次對checkpoint進行恢復以得到完整的checkpoint序列，以供從其中選擇某個或多個checkpoint恢復訓練/進行測試等。
實驗結果
2.1 大語言模型
我們在 Pythia 和 PanGu- 上驗證了我們的壓縮效果，其中Pythia-410M上我們選取了Pile資料集的一個1/20的子集進行基準模型訓練和帶壓縮的模型訓練以進行對比，而 PanGu- 上我們遵從了原論文的訓練策略。
可以看到，我們的結果對比原始模型幾乎沒有損失，同時實現了25-70倍的高倍率整體壓縮。同時，我們也提供了在Pythia-410M上的訓練loss曲線和checkpoint檔案尺寸變化曲線。可以看到帶壓縮的訓練loss曲線與無壓縮的訓練基本保持一致，損失只有非常微小的增加。
2.2 視覺模型
此外，為了方便與現有方法進行對比，我們在ViT-L32模型上進行了實驗，同時比較了其他SOTA方法以及我們的不同剪枝策略。為了進行公平比較，我們根據權重和最佳化器動量的尺寸大小估算了其他方法的整體壓縮率。但需要注意的是，即使是用其他方法在權重上的壓縮率與我們的整體壓縮率比較，我們仍然要高出很多。表裡的消融也證明了我們聯合剪枝策略的有效性。
2.3 消融實驗
我們進一步在Pythia-410M上進行了消融實驗，嘗試去除壓縮流程中的不同步驟以評估各部分的效果。結果表明，直接對權重而非權重殘差進行壓縮時，精度損失會比較嚴重，聯合使用殘差計算、剪枝和量化能夠取得最好的壓縮比例和精度。
總結
我們提出了一個大模型Checkpoint壓縮框架（ExCP），透過Checkpoint殘差、權重與最佳化器動量的聯合剪枝，以及非均勻量化等策略，在保持效能幾乎不受損失的前提下，顯著降低訓練過程中的儲存開銷。我們在大型語言模型和視覺模型上均對此方法進行了評估，證實了方法的有效性。
更多閱讀
#投 稿 通 道#
讓你的文字被更多人看到
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋樑，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。
稿件基本要求：• 文章確係個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標註
• 稿件建議以 markdown 格式撰寫，文中配圖以附件形式傳送，要求圖片清晰，無版權問題
• PaperWeekly 尊重原作者署名權，並將為每篇被採納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算
投稿通道：• 投稿郵箱：[email protected]
• 來稿請備註即時聯絡方式（微信），以便我們在稿件選用的第一時間聯絡作者
• 您也可以直接新增小編微信（pwbot02）快速投稿，備註：姓名-投稿
△長按新增PaperWeekly小編
現在，在「知乎」也能找到我們了
進入知乎首頁搜尋「PaperWeekly」
點選「關注」訂閱我們的專欄吧
·
·
·
