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

中科大博士生一作

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

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

微軟亞洲研究院和SIGMA團隊近日發(fā)布了首個FP4精度的大模型訓練框架，該框架在相同超參數(shù)設(shè)置下，能達到與FP8和BF16相當?shù)挠柧毿Ч@著降低了存儲和計算資源需求。此框架最高可訓練130億參數(shù)規(guī)模、千億級別Tokens的模型。

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

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

2. 技術(shù)細節(jié)：定制化FP4矩陣乘法內(nèi)核

框架采用E2M1的FP4格式（2位指數(shù)，1位尾數(shù)，1位符號位），契合主流ML加速芯片設(shè)計。權(quán)重矩陣和激活矩陣分別采用列方向和行方向的量化策略，最大化FP4在矩陣乘法中的加速效果，避免額外矩陣轉(zhuǎn)置操作?？蚣苁褂胹cale+shift方法進行逐層量化參數(shù)校準。

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

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

4. 混合精度設(shè)計：提升效率

框架采用混合精度設(shè)計，在梯度通信時使用FP8，優(yōu)化器狀態(tài)存儲使用FP16，其他部分使用FP16，在保證數(shù)值穩(wěn)定性的前提下降低計算和存儲開銷。

5. 團隊介紹

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

6. 總結(jié)

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

聯(lián)系作者

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