機器之心報道
編輯：趙陽
大模型的內(nèi)容安全問題使得人們希望可以在終端設(shè)備上完成模型訓(xùn)練及運行。對于手機來說，大模型的權(quán)重數(shù)據(jù)顯得尤為龐大。
大型語言模型（LLM），尤其是生成式預(yù)訓(xùn)練 Transformer（GPT）模型在許多復(fù)雜的語言任務(wù)上表現(xiàn)出了出色的性能。這一突破使人們希望在移動設(shè)備上本地運行這些 LLM，以保護(hù)用戶隱私�？墒�，即使是小型 LLM 也太大，無法在這些設(shè)備上運行。
舉例來說，小型 LLaMA 有 7B 參數(shù)，其 FP16 版本大小為 14GB，而移動設(shè)備只有 18GB 的 DRAM。因此，通過訓(xùn)練時間優(yōu)化（如稀疏化、量化或權(quán)重聚類）來壓縮 LLM 是設(shè)備上 LLM 部署的關(guān)鍵步驟。然而，由于模型大小和計算資源開銷，LLM 的訓(xùn)練時間優(yōu)化非常昂貴。權(quán)重聚類 SOTA 算法之一 DKM，由于需要分析所有權(quán)重和所有可能的聚類選項之間的相互作用，其訓(xùn)練時間可變權(quán)重聚類對計算資源的需求過高。
因此，許多現(xiàn)有的 LLM 壓縮技術(shù)，如 GTPQ 和 AWQ，都依賴于訓(xùn)練后的優(yōu)化。在本文中，研究者提出了內(nèi)存優(yōu)化技術(shù)，以實現(xiàn)訓(xùn)練時間權(quán)重聚類及其在 DKM 中的應(yīng)用，也就是 eDKM。
本文使用的技術(shù)包括跨設(shè)備張量編排和權(quán)重矩陣唯一化及分片。在使用 eDKM 對 LLaMA 7B 模型進(jìn)行微調(diào)并將其壓縮為每個權(quán)重因子占位 3bit 時，研究者實現(xiàn)了解碼器堆棧約 130 倍的內(nèi)存占用減少，優(yōu)于現(xiàn)有的 3bit 壓縮技術(shù)。
提高 DKM 的內(nèi)存效率
如圖 1 所示，剪枝、量化和歸一化都是較為流行的權(quán)重優(yōu)化技術(shù)，這些方法將原始權(quán)重 W，優(yōu)化后得到權(quán)重，以優(yōu)化推理延遲、精度或模型大校在這些技術(shù)中，本文研究者主要關(guān)注的是權(quán)重聚類，特別權(quán)重聚類算法 DKM 。
權(quán)重聚類是一種非線性權(quán)重離散化，權(quán)重矩陣被壓縮成一個查找表和查找表的低精度索引列表，現(xiàn)代推理加速器可以處理這些索引。DKM 通過分析權(quán)重（以 W 表示）和中心點（以 C 表示）之間的相互作用來執(zhí)行可微權(quán)重聚類，并在壓縮比和準(zhǔn)確性之間做出權(quán)衡。
因此，使用 DKM 進(jìn)行 LLM 壓縮會產(chǎn)生高質(zhì)量的結(jié)果。然而，DKM 計算過程中產(chǎn)生的注意力圖較大，前向 / 后向傳遞的內(nèi)存復(fù)雜度為 O (|W||C|)（即圖 1 中的矩陣），這對 LLM 壓縮來說尤其困難。舉例來說，一個 LLaMA 7B 模型僅計算 4 bit 權(quán)重聚類的注意力圖就需要至少 224GB 的內(nèi)存。
圖 1：權(quán)重優(yōu)化系統(tǒng)概覽。DKM 中，系統(tǒng)內(nèi)部創(chuàng)建了一個可微分權(quán)重聚類的注意力圖譜。
因此，研究者需要利用 CPU 內(nèi)存來處理如此大的內(nèi)存需求，也就是先將信息存儲至到 CPU 內(nèi)存，然后在需要時再復(fù)制回 GPU。然而，這將在 GPU 和 CPU 之間產(chǎn)生大量的流量（會因此減慢訓(xùn)練速度），并需要巨大的 CPU 內(nèi)存容量。這意味著減少 CPU 和 GPU 之間的事務(wù)數(shù)量并最大限度地降低每次事務(wù)的流量至關(guān)重要。為了應(yīng)對這些難題，研究者在 PyTorch 中引入了兩種新型內(nèi)存優(yōu)化技術(shù)。
跨設(shè)備的張量編排：跟蹤跨設(shè)備復(fù)制的張量，避免冗余復(fù)制，從而減少內(nèi)存占用，加快訓(xùn)練速度。
權(quán)重唯一化及分片處理：利用 16 bit 權(quán)重僅有 216 個唯一值這一事實來減少注意力圖（如圖 1 所示）的表示，并進(jìn)一步將其分割給多個學(xué)習(xí)模型。
跨設(shè)備張量編排
PyTorch 用數(shù)據(jù)存儲來表示張量，數(shù)據(jù)存儲鏈接到實際的數(shù)據(jù)布局和元數(shù)據(jù)，元數(shù)據(jù)用于保存張量的形狀、類型等。這種張量架構(gòu)讓 PyTorch 可以盡可能地重復(fù)使用數(shù)據(jù)存儲，并有效減少內(nèi)存占用。然而，當(dāng)一個張量移動到另一個設(shè)備上時（如從 GPU 到 CPU），數(shù)據(jù)存儲就不能重復(fù)使用，需要創(chuàng)建一個新的張量。
表 1 舉例說明了張量在 PyTorch 設(shè)備間移動時的內(nèi)存占用情況。在第 0 行分配的張量 x0 在 GPU 上消耗了 4MB。當(dāng)其視圖在第 1 行中改變時，由于底層數(shù)據(jù)存儲可以重復(fù)使用（即 x0 和 x1 實際上是相同的），因此不需要額外的 GPU 內(nèi)存。然而，當(dāng) x0 和 x1 如第 2 行和第 3 行那樣移動到 CPU 時，盡管 y0 和 y1 可以在 CPU 上共享相同的數(shù)據(jù)存儲，但 CPU 內(nèi)存消耗卻變成了 8MB，這導(dǎo)致 CPU 內(nèi)存冗余，并增加了 GPU 到 CPU 的流量。
表 1：LLM 微調(diào)可能需要使用 CPU 內(nèi)存來卸載 GPU 上的內(nèi)存占用。缺乏跨設(shè)備的張量管理會導(dǎo)致跨設(shè)備的冗余拷貝（尤其是當(dāng)計算圖很復(fù)雜時），這對于 LLM 的訓(xùn)練時間優(yōu)化尤為不利。例如，雖然 x0 和 x1 是相同的張量，只是視圖不同，但當(dāng)復(fù)制到 CPU 時，生成的張量 y0 和 y1 并不共享數(shù)據(jù)存儲，而在 GPU 上 x0 和 x1 共享數(shù)據(jù)存儲。
為了解決這種低效問題，研究者在圖 2 (b) 中放置了一個編排層，其中黑色代表實際數(shù)據(jù)存儲和元數(shù)據(jù)，灰色僅表示元數(shù)據(jù)。圖 2 (a) 展示了表 1 中的示例，其中 x1 與 x0 共享數(shù)據(jù)布局，但 y0 和 y1 在 CPU 上擁有重復(fù)的數(shù)據(jù)存儲。如圖 2 (b) 所示，通過插入編排層，研究者避免了這種冗余，并減少了 GPU 傳至 CPU 的流量。研究者使用 PyTorch 中的 save-tensor-hook 來實現(xiàn)這樣的交換方案，檢查相同的數(shù)據(jù)存儲是否已經(jīng)被復(fù)制。
然而，使用這樣的方案來檢查目標(biāo)設(shè)備上是否存在相同的張量是很昂貴的。在圖 2 (b) 的示例中，研究者并沒有將 x1 復(fù)制到 CPU，而是簡單地返回了 y0 的引用以及 x1 和 y0 之間的視圖操作。
圖 2：將跨設(shè)備張量編排應(yīng)用于表 1 中的情況時，可以避免 CPU 端的重復(fù)，從而節(jié)省內(nèi)存及流量。
瀏覽計算圖會增加額外的計算周期，節(jié)省不必要的復(fù)制可以彌補此類開銷。研究者發(fā)現(xiàn)，4 hop 內(nèi)的搜索足以檢測原始 DKM 實現(xiàn)中計算圖中的所有合格的案例。
權(quán)重唯一化及分片處理
在大多數(shù) LLM 的訓(xùn)練中，權(quán)重普遍使用 16 bit 存儲（如 BF16 或 FP16），這意味著雖然 LLM 中有數(shù)十億個參數(shù)，但由于位寬的原因，只有 216 個唯一系數(shù)。這就為大幅壓縮權(quán)重和中心點之間的注意力圖提供了機會，如圖 3 所示。
圖 3：權(quán)重唯一化及分片
實驗結(jié)果
LLM 準(zhǔn)確率
本文將 eDKM 與其他基于量化的壓縮方案進(jìn)行了比較，包括：RTN、SmoothQuant、GPTQ 、AWQ 和 LLM-QAT 。對于 eDKM，研究者還對嵌入層進(jìn)行了 8 bit 壓縮。最終得出如下結(jié)論：
eDKM 使 3 bit 壓縮 LLaMA 7B 模型優(yōu)于所有其他 3 bit 壓縮方案。
eDKM 在 3 bit 和 4 bit 配置的 ARC-e 基準(zhǔn)測試中具有最佳精度。
在使用 4 bit 壓縮模型的 PIQA 和 MMLU 基準(zhǔn)測試中，eDKM 的性能極具競爭力。
消融實驗
在消融實驗中，研究者以 LLaMA 7B 解碼器棧中的一個注意層為例，測量了內(nèi)存占用與 3 bit 壓縮的前向后向速度之間的權(quán)衡。單是跨設(shè)備張量編排就減少了 2.9 倍的內(nèi)存占用，運行時開銷很小，而分片和唯一化模塊則分別節(jié)省了 23.5 倍和 16.4 倍。當(dāng)所有技術(shù)相結(jié)合時，eDKM 可節(jié)省約 130 倍。雖然這些步驟需要額外的計算和通信開銷，但由于 GPU 和 CPU 之間的流量大幅減少，因此運行時的開銷微不足道。
