微软研究院推出首个 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 精度大模型训练框架，通过创新的量化策略和混合精度设计，实现了更高效的模型训练。该框架不仅提升了训练效率，还显著降低了计算和存储资源的需求，为大模型训练开辟了新的路径。