Unweight：如何在保證質量的前提下將大型語言模型壓縮22%

Cloudflare 致力於在全球95%的網際網路使用者附近提供低延遲推理，這意味著必須高效利用 GPU 記憶體。繼去年推出基於 Rust 的推理引擎 Infire 和模型排程平臺 Omni 後，現在他們著手解決下一個瓶頸：模型權重。

大型語言模型生成單個令牌需要從 GPU 記憶體中讀取每一個權重。在 NVIDIA H100 GPU 上，張量核心處理資料的速度比記憶體提供資料的速度快近600倍，因此瓶頸在於記憶體頻寬，而非計算。Unweight 正是為此而生：一種無失真壓縮系統，可將模型權重縮小15-22%，同時保持位精確輸出，且不依賴特殊硬體。其核心突破在於在快速片上記憶體中解壓權重並直接饋送給張量核心，避免了透過慢速主記憶體的額外往返。

壓縮為什麼比聽起來更難

常見的量化技術是有失真壓縮，會不可預測地影響模型質量。Unweight 追求無損，保留精確模型行為。現有系統如 ZipNN、Huff-LLM 和 ZipServ 各有側重，但都不適合 Cloudflare 的需求：在 Hopper GPU 上進行推理時無損解壓，並整合到他們的 Rust 推理引擎中。核心挑戰不在於壓縮本身——BF16 權重中的指數字節高度冗餘，熵編碼效果良好——而在於解壓速度要足夠快，不拖慢推理。

如何有效壓縮模型權重

每個 AI 模型中的數字以16位“腦浮點”（BF16）儲存，包含符號位、指數位和尾數位。符號和尾數變化不可預測，但指數具有規律性：在訓練好的 LLM 中，前16個最常見指數值覆蓋了超過99%的權重，理論上只需約2.6位表示。Unweight 利用這一冗餘，保留符號和尾數，僅對指數字節應用霍夫曼編碼，實現約30%的指數壓縮。該技術選擇性應用於 MLP 權重矩陣（門、上、下投影），這些矩陣約佔模型引數的三分之二，在令牌生成期間主導記憶體流量。注意力權重、嵌入和層歸一化則不壓縮。

GPU 記憶體瓶頸與四種執行流水線

NVIDIA H100 GPU 有兩種相關記憶體：高頻寬記憶體（HBM）和共享記憶體（SMEM）。推理時，每個令牌需要從 HBM 讀取完整權重矩陣，記憶體匯流排是效能瓶頸。Unweight 提供了四種壓縮執行流水線：

1. 完全解碼：完全重構原始 BF16 權重，交給 NVIDIA cuBLAS 庫進行標準矩陣乘法。簡單但寫回主記憶體的資料最多，適合小批次。
2. 僅指數解碼：僅解碼指數字節，將預處理流量減半，使用重構矩陣乘法。
3. 調色盤轉碼：執行時將權重轉碼為緊湊的4位索引格式，預處理流量減至四分之一。
4. 直接調色盤：跳過預處理，權重在載入時預轉碼為4位格式，矩陣乘法核即時重構 BF16 值。

沒有單一流水線始終最優。Unweight 透過自動調優器測量端到端吞吐量，為每個權重矩陣和批次大小選擇最佳流水線。

重構矩陣乘法的工作原理

三種流水線使用自定義矩陣乘法核，融合解壓與計算。該核從 HBM 載入壓縮資料，在共享記憶體中重構 BF16 值，然後直接饋入張量核心，重構權重從不駐留主記憶體。GPU 執行緒組分為生產者（使用TMA載入資料）和消費者（重構並計算），透過迴圈緩衝區實現流水線。

自動調優與結果

自動調優器遍歷候選配置（流水線、矩陣乘法變體、SM分配），為每個投影和批次大小確定最優方案。在 Llama 3.1 8B 上，Unweight 實現推理時模型尺寸減少約13%（僅壓縮門/上投影），分發時減少約22%（壓縮所有 MLP 投影），均為100%位精確無損。當前吞吐量開銷約30-40%，主要來自小批次固定成本、冗餘權重瓦片重構和未壓縮的下投影，這些正在最佳化中。

未來方向

Cloudflare 計劃壓縮下投影（約佔可壓縮權重的三分之一），最佳化核心，並將 Unweight 擴充套件到更大模型及混合專家模型。Unweight 已開源，旨在為 GPU 效率和壓縮研究做出貢獻。