壓縮技術再次創新，大模型有機會塞進手機裡了

圖片來源：由無界AI生成

大型語言模型（LLM），尤其是生成式預訓練Transformer（GPT）模型在許多複雜的語言任務中表現出了出色的表現。這項突破使人們希望在行動裝置上本地運行這些LLM，以保護用戶隱私。可是，即使是小型LLM 也太大，無法在這些裝置上運作。

舉例來說，小型LLaMA 有7B 參數，其FP16 版本大小為14GB，而行動裝置只有18GB 的DRAM。因此，透過訓練時間優化（如稀疏化、量化或權重聚類別）來壓縮LLM 是設備上LLM 部署的關鍵步驟。然而，由於模型大小和計算資源開銷，LLM 的訓練時間最佳化非常昂貴。權重聚類SOTA 演算法之一DKM，由於需要分析所有權重和所有可能的聚類選項之間的相互作用，其訓練時間可變權重聚類對計算資源的需求過高。

因此，許多現有的LLM 壓縮技術，如GTPQ 和AWQ，都依賴訓練後的最佳化。在本文中，研究者提出了記憶優化技術，以實現訓練時間權重聚類別及其在DKM 中的應用，也就是eDKM。

本文所使用的技術包括跨裝置張量編排和權重矩陣唯一化及分片。在使用eDKM 對LLaMA 7B 模型進行微調並將其壓縮為每個權重因子佔位3bit 時，研究者實現了解碼器堆疊約130 倍的記憶體佔用減少，優於現有的3bit 壓縮技術。

提高DKM 的記憶體效率

如圖1 所示，剪枝、量化和歸一化都是較為流行的權重優化技術，這些方法將原始權重W，優化後得到權重

，以優化推理延遲、精度或模型大小。在這些技術中，本文研究者主要關注的是權重聚類，特別權重聚類演算法DKM 。

權重聚類別是一種非線性權重離散化，權重矩陣被壓縮成一個查找表和查找表的低精度索引列表，現代推理加速器可以處理這些索引。 DKM 透過分析權重（以W 表示）和中心點（以C 表示）之間的相互作用來執行可微權重聚類，並在壓縮比和準確性之間做出權衡。

因此，使用DKM 進行LLM 壓縮會產生高品質的結果。然而，DKM 計算過程中產生的注意力圖較大，前向/ 後向傳遞的記憶體複雜度為O (|W||C|)（即圖1 中的矩陣），這對LLM 壓縮來說尤其困難。舉例來說，一個LLaMA 7B 模型只計算4 bit 權重聚類別的注意力圖就需要至少224GB 的記憶體。

*圖1：權重優化系統概覽。 DKM 中，系統內部創建了一個可微分權重聚類別的注意力圖譜。 *

因此，研究者需要利用CPU 記憶體來處理如此大的記憶體需求，也就是先將資訊儲存至到CPU 內存，然後在需要時再複製回GPU。然而，這將在GPU 和CPU 之間產生大量的流量（會因此減慢訓練速度），並需要龐大的CPU 記憶體容量。這意味著減少CPU 和GPU 之間的交易數量並最大限度地降低每次交易的流量至關重要。為了因應這些難題，研究者在PyTorch 中引進了兩種新型記憶體優化技術。

• 跨裝置的張量編排：追蹤跨裝置複製的張量，避免冗餘複製，從而減少記憶體佔用，加快訓練速度。
• 權重唯一化及分片處理：利用16 bit 權重僅有216 個唯一值這一事實來減少注意力圖（如圖1 所示）的表示，並進一步將其分割給多個學習模型。

跨裝置張量編排

PyTorch 用資料儲存來表示張量，資料儲存連結到實際的資料佈局和元數據，元資料用於保存張量的形狀、類型等。這種張量架構讓PyTorch 可以盡可能重複使用資料存儲，並有效減少記憶體佔用。然而，當一個張量移動到另一個裝置上時（如從GPU 到CPU），資料儲存就無法重複使用，需要建立一個新的張量。

表1 舉例說明了張量在PyTorch 裝置間移動時的記憶體佔用情況。在第0 行分配的張量x0 在GPU 上消耗了4MB。當其視圖在第1 行中改變時，由於底層資料儲存可以重複使用（即x0 和x1 實際上是相同的），因此不需要額外的GPU 記憶體。然而，當x0 和x1 如第2 行和第3 行那樣移動到CPU 時，儘管y0 和y1 可以在CPU 上共享相同的資料存儲，但CPU 記憶體消耗卻變成了8MB，這導致CPU 記憶體冗餘，並增加了GPU 到CPU 的流量。

*表1：LLM 微調可能需要使用CPU 記憶體來卸載GPU 上的記憶體佔用。缺乏跨裝置的張量管理會導致跨裝置的冗餘拷貝（尤其是當計算圖很複雜時），這對於LLM 的訓練時間優化尤為不利。例如，雖然x0 和x1 是相同的張量，只是視圖不同，但當複製到CPU 時，生成的張量y0 和y1 並不共享資料存儲，而在GPU 上x0 和x1 共享資料存儲。 *

為了解決這種低效率問題，研究者在圖2 (b) 中放置了一個編排層，其中黑色代表實際數據儲存和元數據，灰色僅表示元數據。圖2 (a) 展示了表1 中的範例，其中x1 與x0 共享資料佈局，但y0 和y1 在CPU 上擁有重複的資料儲存。如圖2 (b) 所示，透過插入編排層，研究者避免了這種冗餘，並減少了GPU 傳至CPU 的流量。研究者使用PyTorch 中的save-tensor-hook 來實作這樣的交換方案，檢查相同的資料儲存是否已經被複製。

然而，使用這樣的方案來檢查目標裝置上是否存在相同的張量是很昂貴的。在圖2 (b) 的範例中，研究者並沒有將x1 複製到CPU，而是簡單地回傳了y0 的引用以及x1 和y0 之間的視圖操作。

*圖2：將跨裝置張量編排應用於表1 的情況時，可以避免CPU 端的重複，進而節省記憶體及流量。 *

瀏覽計算圖會增加額外的計算週期，節省不必要的複製可以彌補此類開銷。研究者發現，4 hop 內的搜尋足以檢測原始DKM 實作中計算圖中的所有合格的案例。

權重唯一化及分片處理

在大多數LLM 的訓練中，權重普遍使用16 bit 儲存（如BF16 或FP16），這意味著雖然LLM 中有數十億個參數，但由於位寬的原因，只有216 個唯一係數。這就為大幅壓縮權重和中心點之間的注意力圖提供了機會，如圖3 所示。

圖3：權重唯一化及分片

實驗結果

LLM 準確率

本文將eDKM 與其他基於量化的壓縮方案進行了比較，包括：RTN、SmoothQuant、GPTQ 、AWQ 和LLM-QAT 。對於eDKM，研究者也對嵌入層進行了8 bit 壓縮。最終得出以下結論：

• eDKM 使3 bit 壓縮LLaMA 7B 模型優於所有其他3 bit 壓縮方案。
• eDKM 在3 bit 和4 bit 配置的ARC-e 基準測試中具有最佳精度。
• 在使用4 bit 壓縮模型的PIQA 和MMLU 基準測試中，eDKM 的效能極具競爭力。

消融實驗

在消融實驗中，研究者以LLaMA 7B 解碼器堆疊中的一個注意層為例，測量了記憶體佔用與3 bit 壓縮的前向後向速度之間的權衡。單是跨裝置張量編排就減少了2.9 倍的記憶體佔用，運行時開銷很小，而分片和唯一化模組則分別節省了23.5 倍和16.4 倍。當所有技術結合時，eDKM 可節省約130 倍。雖然這些步驟需要額外的運算和通訊開銷，但由於GPU 和CPU 之間的流量大幅減少，因此運行時的開銷微不足道。