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

基爾西 發自 凹非寺

量子位 | 公眾號 QbitAI

首個FP4精度的大模型訓練框架來了,來自微軟研究院!

在相同超參數的設置下,可以達到與FP8以及BF16相當的訓練效果

這意味著所需的存儲和計算資源可以更少

用這種方法訓練的模型規模最高可達130億參數規模,訓練Tokens數量也達到千億級別。

而且用的還不是真·FP4,而是通過FP8來模擬,如果採用真的FP4,效果還能進一步提升。

(註:研究開展時,尚未有原生支持FP4的硬件,故作者通過在FP8的TensorCore上模擬實現)

網民評論說,效率更高質量卻沒什麼損失,FP4真的是個game changer。

還有人說,如果這一發現廣為人知,恐怕老黃的股價又要跌了。

當然,因低訓練成本而成為當紅明星的DeepSeek也被網民cue了一下:

在FP8 TensorCore上模擬FP4

如開頭所述,在相同超參數的設置下,作者的方法可以達到與BF16的訓練效果。

具體來說,在1.3B、7B和13B的LLaMA模型上,從0到1千萬Tokens的訓練過程中,作者的FP4訓練與BF16的損失曲線基本一致。

在下遊任務上,表現同樣也和BF16相當。

為了實現FP4精度訓練,研究團隊採用了定製化的FP4矩陣乘法(GeMM)CUDA內核

在內核中,作者先用FP16將FP4的A和B矩陣讀入共享內存並做相應變換,然後用FP4完成分塊矩陣乘法,最後再用FP16對中間結果進行歸約,得到FP16格式的輸出矩陣。

首先需要確定量化的數據格式,該框架採用了E2M1的FP4格式,即用2位來表示指數,1位表示尾數,外加1位符號位,總共4位。

選擇這個格式是為了契合當前主流ML加速芯片的量化計算單元設計。

並且,這個框架對權重矩陣W和激活矩陣A採取了不同粒度的量化策略

對W做的是列方向(channel-wise)的量化,而對A做的是行方向(token-wise)的量化。

這種量化粒度是與GeMM在硬件上的並行實現方式相契合的,可以在不引入額外矩陣轉置操作的前提下,最大化發揮FP4在矩陣乘法上的加速效果。

在模型前向傳播開始時,框架對每一個線性層的權重矩陣W和輸入激活矩陣A同時進行FP4量化

量化時,先對矩陣中的數值進行縮放和偏移,將其映射到FP4所能表示的範圍內,然後通過查表的方式將其四捨五入到最近的FP4離散值。

由於不同層的數值範圍差異很大,所以需要對每一層的權重矩陣和激活矩陣分別確定一個獨立的量化範圍,即進行逐層的量化參數校準

這個框架採用的是scale+shift的校準方法,即先用一個縮放因子將數值從原始範圍映射到[-1,1],再用一個偏移因子把[-1,1]平移到FP4所能表示的範圍。

在反向傳播過程中,如果直接對量化後的矩陣求導,則權重矩陣的梯度幾乎處處為0,從而無法進行參數更新。

為此,作者提出了一種新穎的可微分梯度估計方法

它在前向計算時仍然使用硬量化,以保證計算效率,但在反向傳播時,用一個連續可微的函數來重新擬合這個量化函數,並求導得到一個對梯度的修正項。

另外在訓練過程中,模型的隱層激活分佈通常呈現出明顯的長尾特徵,少數維度上的數值明顯偏大,導致出現「離群點」(outlier)

針對這個問題,作者提出了一種「離群點削峰和補償」的策略。

具體來說,先在激活矩陣A中,通過分位數檢索的方法找出那些幅值最大的離群點,將它們限幅到某一個預設的閾值範圍內,得到削峰後的矩陣A_clamped。

然後,再基於原矩陣A和削峰後的A_clamped,構造出一個稀疏補償矩陣∆A,其中只有那些被削峰的位置是非零的。

此外在部分環節當中,作者還採用了混合精度設計。

比如在梯度通信時採用了FP8,在優化器狀態(如動量)的存儲時選擇了FP16。在系統的其他部分,如非矩陣乘操作的計算、Loss Scaling等,也都採用了FP16。

通過這些混合精度的設計,在保證訓練數值穩定性的前提下,儘可能地降低了計算和存儲開銷。

中科大博士生一作

這個框架由微軟亞洲研究院和SIGMA團隊打造,所有研究人員都是華人。

第一作者Ruizhe Wang是中科大在讀博士生,目前在MSRA實習,研究方向就包括低精度量化。

中科大科研部部長、類腦智能國家工程實驗室執行主任、博士生導師查正軍教授也參與了這一項目。

通訊作者為MSRA高級首席研究經理(Senior Principal Research Manager)程鵬和首席研究經理(Principal Research Manager)Yeyun Gong

程鵬曾先後就讀於北航和清華,在清華讀博期間還到UCLA訪問學習;Yeyun Gong則是複旦博士,畢業後即加入微軟。

MSRA傑出科學家、常務副院長郭百寧也參與了本項目,他先後就讀於北大和康奈爾大學,1999年就已加入微軟。

此外還有其他作者,完整名單如下:

論文地址:

https://arxiv.org/abs/2501.17116

參考鏈接:

https://x.com/arankomatsuzaki/status/1884446877837582598