首個 FP4 精度的大模型訓(xùn)練框架來了，來自微軟研究院！在相同超參數(shù)的設(shè)置下，可以達到與 FP8 以及 BF16 相當?shù)挠?xùn)練效果。
這意味著所需的存儲和計算資源可以更少。
用這種方法訓(xùn)練的模型規(guī)模最高可達 130 億參數(shù)規(guī)模，訓(xùn)練 Tokens 數(shù)量也達到千億級別。
而且用的還不是真?FP4，而是通過 FP8 來模擬，如果采用真的 FP4，效果還能進一步提升。
（注：研究開展時，尚未有原生支持 FP4 的硬件，故作者通過在 FP8 的 TensorCore 上模擬實現(xiàn)）
網(wǎng)友評論說，效率更高質(zhì)量卻沒什么損失，F(xiàn)P4 真的是個 game changer。
還有人說，如果這一發(fā)現(xiàn)廣為人知，恐怕老黃的股價又要跌了。
當然，因低訓(xùn)練成本而成為當紅明星的 DeepSeek 也被網(wǎng)友 cue 了一下：
在 FP8 TensorCore 上模擬 FP4
如開頭所述，在相同超參數(shù)的設(shè)置下，作者的方法可以達到與 BF16 的訓(xùn)練效果。
具體來說，在 1.3B、7B 和 13B 的 LLaMA 模型上，從 0 到 1 千萬 Tokens 的訓(xùn)練過程中，作者的 FP4 訓(xùn)練與 BF16 的損失曲線基本一致。
在下游任務(wù)上，表現(xiàn)同樣也和 BF16 相當。
為了實現(xiàn) FP4 精度訓(xùn)練，研究團隊采用了定制化的 FP4 矩陣乘法（GeMM）CUDA 內(nèi)核。
在內(nèi)核中，作者先用 FP16 將 FP4 的 A 和 B 矩陣讀入共享內(nèi)存并做相應(yīng)變換，然后用 FP4 完成分塊矩陣乘法，最后再用 FP16 對中間結(jié)果進行歸約，得到 FP16 格式的輸出矩陣。
首先需要確定量化的數(shù)據(jù)格式，該框架采用了 E2M1 的 FP4 格式，即用 2 位來表示指數(shù)，1 位表示尾數(shù)，外加 1 位符號位，總共 4 位。
選擇這個格式是為了契合當前主流 ML 加速芯片的量化計算單元設(shè)計。
并且，這個框架對權(quán)重矩陣 W 和激活矩陣 A 采取了不同粒度的量化策略。
對 W 做的是列方向（channel-wise）的量化，而對 A 做的是行方向（token-wise）的量化。
這種量化粒度是與 GeMM 在硬件上的并行實現(xiàn)方式相契合的，可以在不引入額外矩陣轉(zhuǎn)置操作的前提下，最大化發(fā)揮 FP4 在矩陣乘法上的加速效果。
在模型前向傳播開始時，框架對每一個線性層的權(quán)重矩陣 W 和輸入激活矩陣 A 同時進行 FP4 量化。
量化時，先對矩陣中的數(shù)值進行縮放和偏移，將其映射到 FP4 所能表示的范圍內(nèi)，然后通過查表的方式將其四舍五入到最近的 FP4 離散值。
由于不同層的數(shù)值范圍差異很大，所以需要對每一層的權(quán)重矩陣和激活矩陣分別確定一個獨立的量化范圍，即進行逐層的量化參數(shù)校準。
這個框架采用的是 scale+shift 的校準方法，即先用一個縮放因子將數(shù)值從原始范圍映射到 [-1,1]，再用一個偏移因子把 [-1,1] 平移到 FP4 所能表示的范圍。
在反向傳播過程中，如果直接對量化后的矩陣求導(dǎo)，則權(quán)重矩陣的梯度幾乎處處為 0，從而無法進行參數(shù)更新。
為此，作者提出了一種新穎的可微分梯度估計方法。
它在前向計算時仍然使用硬量化，以保證計算效率，但在反向傳播時，用一個連續(xù)可微的函數(shù)來重新擬合這個量化函數(shù)，并求導(dǎo)得到一個對梯度的修正項。
另外在訓(xùn)練過程中，模型的隱層激活分布通常呈現(xiàn)出明顯的長尾特征，少數(shù)維度上的數(shù)值明顯偏大，導(dǎo)致出現(xiàn)“離群點”（outlier）。
針對這個問題，作者提出了一種“離群點削峰和補償”的策略。
具體來說，先在激活矩陣 A 中，通過分位數(shù)檢索的方法找出那些幅值最大的離群點，將它們限幅到某一個預(yù)設(shè)的閾值范圍內(nèi)，得到削峰后的矩陣 A_clamped。
然后，再基于原矩陣 A 和削峰后的 A_clamped，構(gòu)造出一個稀疏補償矩陣?A，其中只有那些被削峰的位置是非零的。
此外在部分環(huán)節(jié)當中，作者還采用了混合精度設(shè)計。
比如在梯度通信時采用了 FP8，在優(yōu)化器狀態(tài)（如動量）的存儲時選擇了 FP16。在系統(tǒng)的其他部分，如非矩陣乘操作的計算、Loss Scaling 等，也都采用了 FP16。
通過這些混合精度的設(shè)計，在保證訓(xùn)練數(shù)值穩(wěn)定性的前提下，盡可能地降低了計算和存儲開銷。
中科大博士生一作
這個框架由微軟亞洲研究院和 SIGMA 團隊打造，所有研究人員都是華人。
第一作者 Ruizhe Wang 是中科大在讀博士生，目前在 MSRA 實習(xí)，研究方向就包括低精度量化。
中科大科研部部長、類腦智能國家工程實驗室執(zhí)行主任、博士生導(dǎo)師查正軍教授也參與了這一項目。
通訊作者為 MSRA 高級首席研究經(jīng)理（Senior Principal Research Manager）程鵬和首席研究經(jīng)理（Principal Research Manager）Yeyun Gong。
程鵬曾先后就讀于北航和清華，在清華讀博期間還到 UCLA 訪問學(xué)習(xí)；Yeyun Gong 則是復(fù)旦博士，畢業(yè)后即加入微軟。
MSRA 杰出科學(xué)家、常務(wù)副院長郭百寧也參與了本項目，他先后就讀于北大和康奈爾大學(xué)，1999 年就已加入微軟。
此外還有其他作者，完整名單如下：
論文地址：
https://arxiv.org/abs/2501.17116
參考鏈接：
https://x.com/arankomatsuzaki/status/1884446877837582598
本文來自微信公眾號：量子位（ID：QbitAI），作者：克雷西
廣告聲明：文內(nèi)含有的對外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。