微软首创FP4训练框架，训练效果媲美BF16，助力大模型训练更高效
微软革新训练框架，超越BF16，引领大模型训练新时代
微软研究院推出了首个支持FP4精度的大规模模型训练框架，在相同的超参数配置下，其训练效果能够媲美FP8和BF16。
这意味着所需的存储和计算资源可以更少。
采用此方法训练的模型最大可达到130亿参数规模，训练的Token数量也能达到千亿级别。
采用的技术并不是真正的FP4，而是通过FP8来模拟实现的。若能使用真实的FP4技术，其效果还会有所提升。 从目前的技术应用来看，虽然通过FP8模拟FP4在一定程度上实现了预期的效果，但显然还有很大的提升空间。这表明，尽管现有的技术已经能够满足基本的需求，但在追求更高质量的图像处理过程中，我们依然需要不断地探索和引入新的技术手段。未来，随着更高级别精度技术的普及，我们有理由相信，图像处理领域将会迎来更为显著的进步。
（注：在研究进行之际，尚无任何硬件原生支持FP4，因此作者通过在FP8的TensorCore上模拟实现了这一功能。
网友评论说，效率更高质量却没什么损失，FP4 真的是个 game changer。
还有人说，若此消息广泛传播，恐怕老黄公司的股票价格又要下滑了。
当然，因低成本训练而走红的DeepSeek近日也引起了网友们的关注。
如开头所述，在相同的超参数设置下，作者提出的方法能够达到与BF16相当的训练效果。这一成果不仅展示了算法设计上的创新性，也进一步证明了混合精度训练在提升深度学习模型效率方面的巨大潜力。这无疑为研究者们提供了一种新的视角来优化现有的训练流程，同时也可能为未来的模型开发带来新的启发。这种进步对于推动整个机器学习领域的发展具有重要意义，特别是在计算资源有限的情况下，该方法能够以更高效的方式实现高质量的训练目标。
具体来说，在1.3B、7B和13B参数量的LLaMA模型上，从0到1千万Tokens的训练过程中，作者发现采用FP4量化技术训练得到的损失曲线与使用BF16精度时的基本一致。 这一发现表明，FP4量化技术不仅能够显著降低模型存储和计算的需求，而且还能在一定程度上保持训练效果的稳定性。这对于推动大语言模型的大规模应用具有重要意义，特别是在资源有限的环境下，如边缘设备或低功耗服务器。此外，这也为未来探索更高效的训练方法提供了新的视角，或许可以进一步优化模型性能，同时减少硬件需求。
在下游任务上，表现同样也和 BF16 相当。
为实现FP4精度训练，研究团队开发了专门的FP4矩阵乘法（GeMM）CUDA内核。
在内核中，作者先用 FP16 将 FP4 的 A 和 B 矩阵读入共享内存并做相应变换，然后用 FP4 完成分块矩阵乘法，最后再用 FP16 对中间结果进行归约，得到 FP16 格式的输出矩阵。
首先需要确定量化的数据格式，该框架采用了 E2M1 的 FP4 格式，即用 2 位来表示指数，1 位表示尾数，外加 1 位符号位，总共 4 位。
选择这个格式是为了契合当前主流 ML 加速芯片的量化计算单元设计。
并且，这个框架在处理权重矩阵W和激活矩阵A时采用了差异化的量化策略，这种设计思路非常新颖。我认为这不仅能够提高模型的效率，还能在一定程度上减少计算资源的消耗，对于推动深度学习技术的实际应用具有重要意义。这样的创新方法可能会引领未来深度学习领域的一个新趋势，值得我们持续关注。
对 W 做的是列方向（channel-wise）的量化，而对 A 做的是行方向（token-wise）的量化。
这种量化的精细程度与GeMM在硬件上的并行执行方式高度匹配，能够在不增加额外矩阵转置操作的情况下，最大程度地利用FP4格式在矩阵乘法中的加速优势。
在模型前向传播开始时，框架对每一个线性层的权重矩阵 W 和输入激活矩阵 A 同时进行 FP4 量化。
量化时，首先对矩阵中的数值进行缩放和偏移，使其适应FP4的表示范围，随后采用查表的方法将其四舍五入至最接近的FP4离散值。
由于不同层的数值范围差异很大，因此需要对每一层的权重矩阵和激活矩阵分别确定一个独立的量化范围，以实现更精准的逐层量化参数校准。这种做法不仅能够有效减少模型在压缩过程中可能出现的信息损失，还能保证各层间数据处理的一致性和准确性。在实际应用中，这种方法已被证明能显著提升模型性能，特别是在深度学习领域，其优势更加明显。通过精细调整每层的量化参数，可以更好地平衡计算效率与模型精度之间的关系，为未来人工智能技术的发展提供了新的思路。
该框架采用了ScaleShift的校准方法，即首先使用一个缩放因子将数值从原始范围映射到[-1,1]区间内，然后再通过一个偏移因子将[-1,1]的范围调整至FP4所能表示的范围内。 这种处理方式在技术层面上展示了一种巧妙的数据转换策略，它不仅能够有效地管理数据范围，还能确保在转换过程中最大限度地保留原始数据的信息。这种方法对于需要在有限数值精度下进行高效计算的应用场景来说，无疑是一个值得借鉴的优秀案例。它体现了现代技术在优化算法效率与精度之间寻求平衡的努力，也预示着未来可能有更多类似的创新方法被应用于不同领域。
在反向传播过程中，若直接对量化后的矩阵求导，则权重矩阵的梯度几乎处处为零，导致无法进行有效的参数更新。
为此，作者提出了一种新颖的可微分梯度估计方法。
在前向计算过程中，模型依然采用硬量化方法以确保高效的运算速度，但在反向传播阶段，则引入了一个连续且可微的函数来替代原有的量化函数，以此来重新构建梯度的计算方式。这样的处理不仅保持了模型运行时的高效性，同时也使模型在训练过程中能够更好地调整参数，从而提升整体性能。这种方法体现了技术领域内对计算效率与模型精度之间平衡点的不断探索，也展示了研究人员如何通过巧妙的设计来解决复杂的技术挑战。
在训练过程中，我们观察到模型的隐层激活分布通常具有显著的长尾特征。这意味着，在少数维度上，数值会明显偏高，从而产生“离群点”。这种情况表明，某些特征可能对模型的学习过程产生了不成比例的影响，这不仅增加了模型训练的复杂性，还可能导致过拟合的风险增加。因此，如何有效地管理和调整这些离群值，以确保模型能够更加稳健地学习，成为了研究者们需要重点关注的问题。
针对这个问题，作者提出了一种“离群点削峰和补偿”的策略。
具体来说，首先在激活矩阵A中，利用分位数检索方法识别出那些幅值最高的离群点，然后将它们限制在一个预先设定的阈值范围内，从而得到经过削峰处理的矩阵A_clamped。
然后，根据原始矩阵A和削峰处理后的矩阵A_clamped，构建一个稀疏补偿矩阵∆A，该矩阵只有在被削峰的位置上才具有非零值。
此外在部分环节当中，作者还采用了混合精度设计。
例如，在梯度通信过程中使用了FP8格式，在优化器状态（如动量）的存储时选择了FP16格式。系统中的其他部分，如非矩阵运算和Loss Scaling等，也均采用FP16格式。
通过这些混合精度的设计，在保证训练数值稳定性的前提下，尽可能地降低了计算和存储开销。
这个框架由微软亚洲研究院和 SIGMA 团队打造，所有研究人员都是华人。
第一作者Ruizhe Wang是中科大的在读博士生，目前在微软亚洲研究院（MSRA）实习，他的研究领域涵盖了低精度量化技术。
中科大科研部部长兼类脑智能国家工程实验室执行主任查正军教授也加入了这一重要项目。查教授在人工智能领域有着丰富的研究经验，他的加入无疑为项目的成功增添了更多的保障。我们期待他在项目中的贡献能够推动相关技术的发展，并为中国在这一领域的领先地位做出更大的贡献。 这种高水平专家的参与不仅体现了项目的重要性和影响力，也展示了中国在科技创新方面的强大实力和决心。通过这样的合作与努力，相信中国在未来的人工智能领域将取得更多突破性的成果。
通讯作者为微软亚洲研究院的高级首席研究经理程鹏和首席研究经理龚艳君。
程鹏先后在北航和清华深造，他在清华大学攻读博士学位时还赴加州大学洛杉矶分校进行了学术交流。而Yeyun Gong则是在复旦大学获得博士学位后，直接加入了微软公司。 这样的教育背景和职业路径展示了他们各自领域的深厚积累与广泛视野。程鹏通过在国内外顶尖学府的学习和交流，无疑积累了丰富的专业知识和技术经验。这种跨文化的学术经历对他的科研工作和未来的职业发展都将产生积极的影响。另一方面，Yeyun Gong选择在完成学业后立即投身于业界领先的科技公司，这表明了他对实际应用和技术创新的高度关注。这样的选择也使得他能够在实践中进一步深化理论知识，并为推动技术进步贡献力量。
MSRA杰出科学家及常务副院长郭百宁亦有参与此项目。他曾先后在北大和康奈尔大学深造，早在1999年便已加入微软。
此外还有其他作者，完整名单如下：
论文地址：
https://arxiv.org/abs/2501.17116
参考链接：
https://x.com/arankomatsuzaki/status/1884446877837582598
本文来自微信公众号：量子位（ID：QbitAI），作者：克雷西