大模型訓練開銷還能更小！微軟推出首個FP4訓練框架，訓練效果與BF16相當

中科大博士生一作

原標題：大模型訓練開銷還能更小！微軟推出首個FP4訓練框架，訓練效果與BF16相當
文章來源：量子位
內容字數：3919字

微軟研究院發布FP4精度大模型訓練框架

微軟亞洲研究院和SIGMA團隊近日發布了首個FP4精度的大模型訓練框架，該框架在相同超參數設置下，能達到與FP8和BF16相當的訓練效果，顯著降低了存儲和計算資源需求。此框架最高可訓練130億參數規模、千億級別Tokens的模型。

1. 突破性進展：FP4精度訓練

該框架通過在FP8的TensorCore上模擬FP4精度實現，而非使用真正的FP4硬件（目前尚未有原生支持FP4的硬件）。即便如此，該框架在1.3B、7B和13B的LLaMA模型訓練中，損失曲線與BF16基本一致，下游任務表現也相當。網友評論認為，這將是游戲規則的改變者，并預測此發現可能影響英偉達股價。

2. 技術細節：定制化FP4矩陣乘法內核

框架采用E2M1的FP4格式（2位指數，1位尾數，1位符號位），契合主流ML加速芯片設計。權重矩陣和激活矩陣分別采用列方向和行方向的量化策略，最大化FP4在矩陣乘法中的加速效果，避免額外矩陣轉置操作。框架使用scale+shift方法進行逐層量化參數校準。

3. 創新性解決方法：可微分梯度估計和離群點處理

為了解決直接對量化矩陣求導導致梯度幾乎為零的問題，該框架提出了一種新穎的可微分梯度估計方法。在前向計算中使用硬量化保證效率，在反向傳播中使用連續可微函數擬合量化函數，并計算梯度修正項。針對模型激活分布的長尾特征和離群點問題，框架采用“離群點削峰和補償”策略，先限幅離群點，再構造稀疏補償矩陣。

4. 混合精度設計：提升效率

框架采用混合精度設計，在梯度通信時使用FP8，優化器狀態存儲使用FP16，其他部分使用FP16，在保證數值穩定性的前提下降低計算和存儲開銷。

5. 團隊介紹

該框架由微軟亞洲研究院和SIGMA團隊的華人研究人員打造，第一作者是中科大在讀博士生王瑞哲（Ruizhe Wang），通訊作者為微軟亞洲研究院的程鵬和龔業耘。中科大查正軍教授、微軟亞洲研究院郭百寧也參與了該項目。

6. 總結

微軟研究院的這項研究成果，通過巧妙的模擬和創新性算法，實現了FP4精度的大模型訓練，為降低大模型訓練成本提供了新的方向，并有望推動大模型訓練技術進一步發展。 論文已發表在arXiv上。

聯系作者

文章來源：量子位
作者微信：
作者簡介：追蹤人工智能新趨勢，關注科技行業新突破