第 6 章 — 剪枝與知識蒸餾

《LLM Primer VI:擴展 AI 系統》章節走讀的第六篇。這一章直接攻擊權重的數量 — 先把不重要的權重歸零,再把大模型的行為轉移到一個更小的模型裡。

這一章為什麼存在

量化縮小每個權重的位元組數。剪枝縮小根本會被相乘的權重數量。知識蒸餾把大老師的行為轉移到一個從頭到尾都跑得比較便宜的更小學生身上。剪枝是一種同架構的編輯,削減開銷;蒸餾是一種架構變更,拿訓練期成本換一次永久的推論期節省。第 6 章走過兩者 — 各自仰賴的機制、實務工作者會踩到的失敗模式、以及在 Hopper 與 Blackwell 上讓結構化剪枝重新變得經濟上有意思的硬體加速樣式(2:4 稀疏)。

6.1 2:4 稀疏是硬體會獎勵的剪枝樣式

非結構化的量級剪枝 — 把最小的 50% 權重歸零然後短暫微調 — 數值上有用,但不會加速。GPU 執行的是密集 GEMM;散亂的零還是會被乘。結構化剪枝拿掉的是硬體能跳過的整塊單位:注意力頭、MLP 神經元、或(重要的情況)權重矩陣裡的 N:M 樣式。Hopper 與 Blackwell 原生支援 2:4 稀疏 — 沿著輸入維度,每連續四個權重裡剛好有兩個是零。晶片把每組四個權重存成兩個非零值加上兩位元的 mask,稀疏 tensor core 跳過為零的乘法,頻寬與 FLOPs 大約都減半。實測 decoding 吞吐量會有 1.7 到 1.9 倍的加速。訣竅是 2:4 必須在剪枝的時候就強加上去:演算法一組四個地走每一個矩陣,強制其中兩個最小的變成零,然後不是短暫微調、就是套用 SparseGPT 補償。mask 是 metadata;執行期節省是自動的。

6.2 蒸餾轉移的是老師的分布,而不是它的 argmax

訓練過的老師 LLM 在每個位置產生一個關於字彙的分布 — 一個 12.8 萬維的 softmax。argmax 是它在 greedy decoding 時會吐出的 token。但那個分布編碼了老師的不確定性、它在近義詞之間的偏好結構、以及它的校準。訓練學生去符合那個分布,會把這些東西全部轉移過去:損失是老師與學生輸出經過軟化後的 KL 散度,搭配一個能把梯度訊號拓寬到低機率尾巴的溫度。長出來的學生行為上像一個更小、更快版本的老師,而不是一個用同樣資料從頭訓練起來的更小模型。Patient Knowledge Distillation 把這個配方延伸到中間層,在對應的學生-老師層對之間比對隱藏狀態。MiniLLM 處理一個更微妙的失敗 — 模式平均問題,forward-KL 學生會跨越好幾個看起來都合理的老師延續做折衷,產生平淡的輸出 — 它改用 reverse-KL,傾向對老師某一個模式做出銳利的匹配,而不是對所有模式做柔軟的覆蓋。

6.3 三種壓縮會疊在一起,而且有特定順序

先蒸餾、再剪枝、最後量化。蒸餾是最貴、最觸及架構的操作;道理上應該在其他選擇鎖定之前,先在一顆密集、高精度的模型上做一次。用 SparseGPT 把學生剪成 2:4,再短暫恢復微調。用 llm-compressor 把稀疏化的學生量化到 FP8。一個具體的堆疊:70B BF16 老師 → 13B BF16 蒸餾學生 → 2:4 稀疏化 → FP8 量化。每次前向傳遞的權重位元組數從 140 GB 掉到大約 6.5 GB — 22 倍的縮減 — 而 Hopper 的 FP8 與 2:4 路徑會把它變成每 GPU 8 到 10 倍的吞吐量收益。這個堆疊不會減少 KV 快取,因為它取決於隱藏維度、頭數、上下文長度,而不是權重數量。那是第 8 章的主題。而且模型層級的壓縮是「每次推論」的;它讓 GPU 有餘裕,但把餘裕填成同時在線使用者的執行期槓桿,是批次處理。

第 6 章鋪陳出來的東西

第 5、6 章把模型端的壓縮走完了 — 量化每個權重的位元組數、剪枝權重的數量、蒸餾整個模型往下縮。GPU 現在每一次前向傳遞都有餘裕了。第 7 章接著轉到那個把餘裕填成有用工作的系統級槓桿:批次處理。靜態批次差一點就成功,卻在最快完成者問題上塌陷。連續批次搭配迭代級排程,是生產引擎在跑的東西,也是把受頻寬所困的 decode 階段從閒置晶片變成同時服務許多使用者的晶片的關鍵。

下一篇 — 第 7 章:進階批次處理策略。靜態、動態、與連續批次 — 以及連續批次為 KV 快取製造出來的債務。