通过创新的FP4训练框架和低精度量化技术，微软实现大模型训练成本大幅降低，性能媲美BF16，混合精度设计助力高效训练

微软推出首个FP4精度大模型训练框架，显著降低计算和存储成本

近日，微软研究院推出了首个支持FP4（4位浮点）精度的大模型训练框架。这一创新使得在相同超参数设置下，训练效果与FP8和BF16相当，大幅减少了所需的存储和计算资源。这意味着模型规模最高可达130亿参数，训练Tokens数量也达到了千亿级别。

尽管目前尚无原生支持FP4的硬件，研究人员通过在FP8 TensorCore上模拟实现了FP4训练。未来，如果采用真正的FP4硬件，效果有望进一步提升。不少网友表示，这种方法不仅提高了效率，还几乎不损失训练质量，堪称“游戏规则改变者”。

FP4训练框架的技术细节

为了实现FP4精度训练，研究团队开发了定制化的FP4矩阵乘法（GeMM）CUDA内核。该内核首先用FP16将FP4的A和B矩阵读入共享内存并做相应变换，然后用FP4完成分块矩阵乘法，最后再用FP16对中间结果进行归约，得到FP16格式的输出矩阵。整个过程中采用了E2M1的FP4格式，即用2位表示指数，1位表示尾数，外加1位符号位，总共4位。

对于权重矩阵W和激活矩阵A，框架采取了不同粒度的量化策略：权重矩阵W进行列方向（channel-wise）量化，激活矩阵A进行行方向（token-wise）量化。这种量化方式可以最大化发挥FP4在矩阵乘法上的加速效果，同时避免引入额外的矩阵转置操作。

在前向传播开始时，框架对每个线性层的权重矩阵W和输入激活矩阵A同时进行FP4量化。量化时，先对数值进行缩放和偏移，将其映射到FP4所能表示的范围内，再通过查表四舍五入到最近的离散值。为确保每一层的数值范围差异不影响整体性能，框架采用了逐层的量化参数校准方法。

解决反向传播中的梯度问题

在反向传播过程中，直接对量化后的矩阵求导会导致权重矩阵的梯度几乎处处为0，无法进行参数更新。为此，研究团队提出了一种新颖的可微分梯度估计方法。该方法在前向计算时使用硬量化以保证效率，而在反向传播时用一个连续可微的函数重新拟合量化函数，并求导得到梯度修正项。

此外，针对模型隐层激活分布中出现的“离群点”问题，团队提出了一种“离群点削峰和补偿”的策略。具体来说，通过分位数检索找出幅值最大的离群点，将其限幅到预设阈值范围内，构造出稀疏补偿矩阵以保持数值稳定性。

混合精度设计优化性能

为了进一步优化性能，研究团队还在部分环节采用了混合精度设计。例如，在梯度通信时采用FP8，在优化器状态存储时选择FP16。在非矩阵乘操作的计算、Loss Scaling等环节也采用了FP16。这些设计在保证训练数值稳定性的前提下，尽可能地降低了计算和存储开销。

研究团队背景

这个框架由微软亚洲研究院和SIGMA团队共同打造，所有研究人员均为华人。第一作者王瑞哲是中科大在读博士生，目前在MSRA实习，专注于低精度量化研究。其他参与者包括中科大科研部部长查正军教授、MSRA高级首席研究经理程鹏和首席研究经理龚耶云，以及MSRA杰出科学家郭百宁。

论文地址：[https：//arxiv.org/abs/2501.17116](https：//arxiv.org/abs/2501.17116)

参考链接：[https：//x.com/arankomatsuzaki/status/1884446877837582598](https：//x.com/arankomatsuzaki/status/1884446877837582598)

本文来源：