老闆要求把公司自研的120B參數模型砍到一半大小、再壓成4-bit,好塞進公司只買得起的那幾張顯卡裡跑
這篇文章講的是一套幫AI模型『瘦身後療傷』的技術。模型瘦身通常分兩刀:第一刀是結構壓縮直接砍掉模型裡一部分的神經元、層數或注意力頭,讓整個模型變得比較瘦小,不是數字變粗略,是骨架本身變小了,第二刀是量化把模型裡的數字,從精細的小數四捨五入成只有少數幾種選擇的粗略數字,檔案變小、跑得更快,但代價是誤差。文中的例子是把一個120B(1200億參數)的模型砍成60B,再壓成4-bit每個數字只用4個0與1組成的密碼表示,總共只能表示16種不同的數值,是一種很粗糙的精度格式,叫做MXFP4。
問題是,這兩刀砍下去,模型通常會變笨、答非所問,所以砍完一定要『療傷』。過去主流做法有兩種。第一種是量化感知訓練(QAT)在訓練過程中刻意先模擬量化一次,讓模型習慣這種粗糙表示法後再繼續訓練,用真實任務重新調教,缺點是要重跑一次昂貴的完整訓練流程,而且練太久反而會變不穩定。第二種是知識蒸餾讓一個小模型(學生)去模仿一個已經很厲害的大模型(老師)的答案機率分佈,把老師的知識教給學生,做法是用KL散度一種數學方法,用來衡量兩個機率分佈有多不一樣,數值越小代表兩者的答案越像當作損失函數,這叫量化感知蒸餾(QAD)。
QAD在『只量化、沒有結構壓縮』時效果不錯,因為此時還存在一個貨真價實、沒被動過的全精度(bfloat16)模型完全沒被壓縮、沒被量化過的原始版本,數字保留得最細緻版本可以當老師。但一旦模型連骨架都被砍過,就找不到『同架構的全精度版』了,只能拿『砍完又修復過』的近似版本當老師,學生的天花板直接被這個不完美的老師鎖死。這篇論文提出的量化感知修復(QAH),做法只改一件事:不用近似版當老師,直接讓學生對齊砍之前、最原始的那個模型。結果在9項基準測試(Benchmark)一組事先設計好的標準考題,用來客觀比較不同模型的能力好壞中,有7項贏過完全沒被動過的原始全精度版本。
| 方法 | 老師是誰 | 遇到結構壓縮時會怎樣 |
|---|---|---|
| QAT(量化感知訓練) | 沒有老師,直接拿真實任務重新訓練 | 要重跑一次昂貴的完整訓練流程,而且練太久效果反而會變不穩定 |
| QAD(量化感知蒸餾) | 砍完、修復過的近似版本 | 老師自己就是打了折的近似值,學生的天花板被老師鎖死 |
| QAH(量化感知修復) | 砍之前、最原始的全精度模型 | 老師是最好的版本,學生沒有天花板限制,甚至可能比原始還準 |
這段示意程式碼濃縮了QAD跟QAH唯一的差異:兩行loss計算,對比的『老師』到底是誰,這個選擇決定了學生模型最終能有多準。
teacher_qad = load_model('recovered_bf16_checkpoint')teacher_qah = load_model('original_120b_full_precision')student = load_model('60b_mxfp4_quantized')for batch in calibration_data: student_logits = student(batch) loss_qad = kl_divergence(student_logits, teacher_qad(batch).detach()) loss_qah = kl_divergence(student_logits, teacher_qah(batch).detach()) loss_qah.backward() optimizer.step()