預調節器空間捨入:重構 4-bit AdamW 優化器狀態量化
Rounding in Preconditioner Space: Redesigning 4-bit AdamW Quantization
將 AdamW 優化器狀態量化為 4-bit 能顯著節省記憶體,但量化誤差會在動量遞迴中累積並破壞更新。本研究揭示狀態空間的最小誤差並不代表預調節器空間的低誤差。研究團隊為此重構量化捨入機制,提出 ZIP-SR(包含零值的預調節器空間隨機捨入)與 ZE-EDEN(不含零值的校準區塊重縮放)兩種方案,並在訓練最後 10% 階段對 LM-head 的第一動量進行特定隨機捨入。在 130M 至 2.7B 參數的模型預訓練與全參數 SFT 實驗中,顯著超越 TorchAO 4-bit AdamW,彌合高達 70% 與 32-bit AdamW 的驗證損失差距。
核心重點
預調節器空間捨入理論
證明二階動量在狀態空間的低量化誤差不等於預調節器空間的低誤差,改以預調節器座標選擇量化階梯。
ZIP-SR 與 ZE-EDEN 新配方
ZIP-SR 保留零值並在預調節器空間計算隨機捨入;ZE-EDEN 則排除零值並重縮放區塊以減緩量化下限產生的扭曲。
LM-Head 第一動量針對性處理
一階動量採用 NF4 量化,並於訓練最後 10% 階段對 LM-head 的一階動量實施特定隨機捨入以穩定收斂。
縮小 70% 效能差距
在 GPT 與 Llama 架構(130M 至 2.7B 參數)預訓練中,驗證損失大幅接近 FP32 AdamW,差距最多減少 70%。
技術圖解
為什麼重要
大語言模型訓練常受限於優化器狀態所佔用的巨大 GPU 記憶體。過去 4-bit 量化優化器(如 TorchAO)雖然能省下記憶體,但因為誤差累積導致模型效果明顯下降。這項研究透過重構捨入數學空間,在幾乎不增加記憶體開銷的前提下大幅修復了精度損失,讓大模型預訓練與全參數微調(SFT)能夠同時享有低記憶體需求與高品質收斂。
對誰有影響
- AI 開發者
- AI 研究人員
- 企業決策者
可以怎麼使用
- 1大語言模型(130M 至 2.7B+ 參數)的高記憶體效率預訓練
- 2大幅降低 VRAM 消耗的大模型全參數監督式微調(SFT)
限制與注意事項
- 實驗規模目前驗證至 2.7B 參數,在超大規模模型(如 70B+)上的效果仍需更多實證。
- 與標準 32-bit AdamW 相比,4-bit 量化仍存在極微小的收斂品質差距。
延伸閱讀
重新審視 METR 時間跨度指標:以統計模型量化 AI 的真實任務能力
Re-Evaluating AI Time Horizons: A Statistical Assessment of the METR Benchmark
本研究利用樣條函數與項目反應理論重新分析 METR 資料集,發現 AI 處理任務的難度與人類時間對數非線性相關,突破 2 至 30 分鐘任務的難度遠低於從 30 分鐘跨越至 5 小時。
Bi-FORK:高維分歧物理系統的生成式建模框架
Bi-FORK: Generative Modeling for High-Dimensional Bifurcating Physical Systems
Bi-FORK 結合隱空間流匹配與推斥引導取樣,解決物理分歧現象中的一對多相干軌跡生成難題,可擴展至高達 26 萬離散點的高維系統。
單一區塊實現多重深度:具備深度編程專家庫的循環 Vision Transformer
One Block, Multiple Depths: Recurrent Vision Transformers with Depth-Programmed Experts
論文提出 reViT 架構,僅利用單一重複使用的 Transformer 區塊,透過動態組合專家權重來模擬不同深度的特徵轉換,在大幅減少 70% 參數量的同時達到與標準全深度模型相當的精準度。