大模型训练开销还能更小！微软推出首个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上。

联系作者

文章来源：量子位
作者微信：
作者简介：追踪人工智能新趋势，关注科技行业新突破