微軟研究院推出首個 FP4 精度大模型訓練框架，效率提升顯著

1月29日消息，微軟研究院發布了首個 FP4 精度的大模型訓練框架。在相同超參數設置下，該框架可以達到與 FP8 和 BF16 相當的訓練效果。這意味著所需的儲存和計算資源可以更少。
模型規模與訓練效果
使用該方法訓練的模型規模最高可達 130 億參數，訓練 Tokens 數量也達到千億級別。雖然目前使用的是通過 FP8 模擬的 FP4，但如果採用真正的 FP4，效果將進一步提升。
FP4 模擬與訓練方法
由於研究開展時尚未有原生支持 FP4 的硬體，研究團隊通過在 FP8 的 TensorCore 上模擬實現 FP4。具體來說，在 1.3B、7B 和 13B 的 LLaMA 模型上，從 0 到 1 千萬 Tokens 的訓練過程中，FP4 訓練與 BF16 的損失曲線基本一致。
為了實現 FP4 精度訓練，研究團隊採用了定制化的 FP4 矩陣乘法（GeMM）CUDA 內核。在內核中，先用 FP16 將 FP4 的 A 和 B 矩陣讀入共享內存並做相應變換，然後用 FP4 完成分塊矩陣乘法，最後再用 FP16 對中間結果進行歸約，得到 FP16 格式的輸出矩陣。
量化策略與硬體適配
該框架採用了 E2M1 的 FP4 格式，即用 2 位表示指數，1 位表示尾數，外加 1 位符號位，共 4 位。選擇這個格式是為了契合當前主流 ML 加速芯片的量化計算單元設計。對權重矩陣 W 和激活矩陣 A 採用了不同粒度的量化策略，最大化發揮 FP4 在矩陣乘法上的加速效果。
前向與反向傳播的創新
在模型前向傳播開始時，框架對每一個線性層的權重矩陣 W 和輸入激活矩陣 A 同時進行 FP4 量化。量化時，先對矩陣中的數值進行縮放和偏移，將其映射到 FP4 所能表示的範圍內，然後通過查表的方式將其四捨五入到最近的 FP4 離散值。
在反向傳播過程中，研究團隊提出了一種新穎的可微分梯度估計方法，以保證計算效率和梯度的準確性。此外，針對激活矩陣中的 “離群點” 問題，研究團隊提出了 “離群點削峰和補償” 的策略。
混合精度設計與應用
在部分環節中，研究團隊採用了混合精度設計，如在梯度通信時採用 FP8，在優化器狀態存儲時選擇 FP16。在系統的其他部分，如非矩陣乘操作的計算、Loss Scaling 等，也都採用了 FP16。這些設計在保證訓練數值穩定性的前提下，盡可能地降低了計算和存儲開銷。
研究團隊與背景
這個框架由微軟亞洲研究院和 SIGMA 團隊打造，所有研究人員都是華人。第一作者 Ruizhe Wang 是中科大在讀博士生，目前在 MSRA 實習，研究方向包括低精度量化。中科大科研部部長查正軍教授也參與了這一項目。通訊作者為 MSRA 高級首席研究經理程鵬和首席研究經理 Yeyun Gong。MSRA 傑出科學家郭百寧也參與了本項目。
總結
微軟研究院推出的 FP4 精度大模型訓練框架，通過創新的量化策略和混合精度設計，實現了更高效的模型訓練。該框架不僅提升了訓練效率，還顯著降低了計算和存儲資源的需求，為大模型訓練開闢了新的路徑。
論文地址：arXiv
參考鏈接: Twitter