FP8训练新范式：减少40%显存占用，训练速度提高1.4倍
FP8训练新范式：减少40%显存占用，训练速度提高1.4倍
近期，来自伯克利、英伟达、MIT和清华的研究者们提出了显存高效的FP8训练方法：COAT（CompressingOptimizer states andActivation for Memory-Efficient FP8Training）。该方法通过FP8量化来压缩优化器状态和激活值，实现了端到端内存占用减少1.54倍，端到端训练速度提高1.43倍，同时保持模型精度。
近期DeepSeek V3引爆国内外的社交媒体，他们在训练中成功应用了FP8精度，显著降低了GPU内存使用和计算开销。这表明，FP8量化技术在优化大型模型训练方面正发挥着越来越重要的作用。
FP8优化器状态
FP8量化优化器状态的难点在于，当前的量化方法无法充分利用FP8的表示范围。对于FP8的E4M3格式，我们希望量化组X的动态范围覆盖E4M3的最小可表示值（0.00195）和最大可表示值（448）之间的整个跨度，以充分利用其表示能力。然而，E4M3的动态范围通常未被充分利用：E4M3的动态范围约为200000，但一阶动量的每个量化组的最大值最小值之比通常为1000，二阶动量的该比值则通常为10，远小于E4M3的动态范围。这使得用FP8来量化优化器状态的误差非常大。
解决方案：动态范围扩展
研究者发现，在量化之前引入一个扩展函数f(・)，能够扩大量化组的动态范围，并使其与E4M3对齐。使用的扩展函数为：
其中，k是即时计算的参数，每个量化组共享一个k。当k>1时，动态范围将被扩大，并更接近E4M3的动态范围。在每一步训练中，都可以即时的计算出最优的k，从而可以充分利用E4M3的表示范围，而原始的量化方法只能利用其中的一小部分。
动态范围扩展方法可以大大减少量化误差，并充分利用E4M3的动态范围。除此之外，还发现，E4M3比E5M2更适合一阶动量。而对于二阶动量，虽然在原始设置中E4M3优于E5M2，但在应用扩展函数后，它们的量化误差几乎相同。因此，建议在量化优化器状态时使用E4M3+E4M3量化策略或E4M3+E5M2量化策略。
FP8激活
动机：非线性层占用大量内存
在语言模型的前向传播中，必须保留激活值以用于反向传播计算梯度。在Llama模型系列中，非线性层通常占内存占用的约50%。相比之下，线性层的贡献不到25%。因此，优化线性和非线性层以减少激活内存占用至关重要。
解决方案：混合粒度FP8精度流
FP8精度流要求所有线性和非线性层的输入和输出采用FP8格式。通过直接以FP8格式保存输入张量用于反向传播，这消除了额外的量化操作需求，从而减少了相关开销。FP8精度流自然地将非线性和线性层的内存占用减少了50%，因为它们只需要保存FP8激活值，而不是BF16。为了进一步提高该方法的准确性，作者提出在不同层中变化量化粒度，以混合粒度的方式平衡精度和效率。
实验结果
COAT在多个任务中展示了其在内存占用和训练速度方面的优势，同时保持了模型性能。
训练加速1.43倍，显存降低1.54倍
在使用4张H100训练Llama-2-13B模型时，COAT将每个GPU的内存占用从BF16的76.1GB减少到49.1GB，实现了1.54倍的内存缩减。同时，COAT将训练速度从BF16的每秒2345个token提升至每秒5295个token，达到1.43倍的加速。在几乎所有的训练场景下，COAT都能够使Batch Size翻倍，或是让训练所需的卡数减小。
训练完全不掉点，FP8训练表现和BF16吻合
COAT在各种应用场景下，均展现出了出色的精度，完全不会导致模型性能下降。例如，在大语言模型预训练任务中，COAT可以保持近乎无损的模型性能，训练中的loss曲线也和BF16完全吻合。
COAT在视觉语言模型微调中同样实现了和BF16训练完全一致的表现。无论是loss曲线，还是下游任务上的表现，COAT均和BF16基准相持平。
在一些实际的下游任务例子中，经过COAT训练过的模型也有着相当优秀的生成和总结能力。
总结
COAT的核心价值在于使用FP8进行训练的同时做到了显存优化。动态范围扩展减少量化误差，混合粒度量化优化激活存储，两者协同作用使得端到端内存占用降低1.54倍。这种优化不仅适用于单机训练，更在分布式训练中发挥关键作用——通过批量大小翻倍，可在相同硬件条件下处理更多数据，显著提升训练效率。而对于显存资源紧张的研究者，COAT也提供了全参数训练的可行路径，降低了大模型训练的门槛。
论文标题：COAT: Compressing Optimizer States and Activation for memory efficient FP8 Training