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