Unified FP8: Moving Beyond Mixed Precision for Stable and Accelerated MoE RL

发表时间: 2025-11 · Blog post by InfiXAI Team, Ant Group AQ Team, SGLang RL Team, Miles Team (lmsys.org)

原文: https://lmsys.org/blog/2025-11-25-fp8-rl

作者/机构：InfiXAI Team, Ant Group AQ Team, SGLang RL Team, Miles Team

速读

一句话结论 本文在强化学习阶段实现了端到端的统一 FP8 训练与采样（FP8-TI），通过消除混合精度方案中由量化误差导致的训练与推理不一致问题，显著提升了大规模 MoE 模型强化学习的稳定性。

要解决什么问题 现有的强化学习（RL）训练通常采用 BF16 训练结合 FP8 采样（Rollout）的混合精度方案，这种做法卡在严重的训练与推理不一致问题上。具体而言，FP8 格式本身固有的量化与反量化操作会引入数值误差。在预训练或微调阶段，这种误差通常可控；但在 RL 阶段，梯度的计算往往依赖于两次前向传播输出的对数概率（log probabilities）之差。混合精度方案下，采样时使用 FP8，而训练前向传播时使用 BF16，两者之间的固有量化误差会在计算概率差时被急剧放大，导致初始 KL 散度异常升高，甚至将优化推向错误的方向。此外，随着混合专家（MoE）模型参数规模的增大，这种由量化误差引起的训练与推理不一致现象会愈发严重，直接导致严重的离策略（off-policy）偏移，使得 RL 训练极不稳定。同时，在 Agentic RL 常用的小批量（small batch size）场景下，频繁的量化操作会引发 CPU 瓶颈，导致显存和计算效率的实际收益远不及理论预期。

怎么做的 核心思路是采用统一的 FP8 训练与推理架构（FP8-TI），即在 RL 的训练和采样阶段全部使用 FP8 精度。它能绕开上述卡点的根本原因在于：作者通过构建“伪量化（Fake Quant）”与“真量化（Real Quant）”的对比实验证明，导致 KL 散度异常和训练偏离的误差并非来自底层的 FP8 矩阵乘法算子（如 cuBLAS），而是来自量化与反量化过程本身的固有精度损失，且训练前向传播中的量化误差对一致性的影响起决定性作用。通过在训练和采样两端统一使用 FP8，模型在生成经验和更新策略时经历完全相同的量化过程，从而在机制上消除了两者之间的对数概率计算偏差。
在关键设计上，该方法由两个核心部件构成。首先是缩放因子（Scale）的精度选择，放弃了极度追求硬件效率的 E8M0 纯整数缩放，而是采用 FP32 作为缩放精度。这不仅能提供细粒度的数值缩放以捕捉张量的动态范围，使 FP8 训练的损失曲线最大程度贴合 BF16 基线，还能兼容当前 Hopper 架构缺乏 E8M0 专用计算单元的硬件现状。其次是混合粒度的 FP8 量化策略，为了在训练稳定性和计算效率之间取得平衡，对不同张量采用不同的量化粒度。量化的基本定义由以下两步构成，首先计算给定张量 $X$ 的缩放因子 $S$：

$$S = \frac{\max|X|}{V_{\max}}$$


其中 $V_{\max}$ 为 FP8 能表示的最大值。随后利用该缩放因子计算量化后的值 $Q(x)$：

$$Q(x) = \mathrm{round}\left(\frac{x}{S}\right)$$
在具体分配上，激活值（Activations）包含较多异常值，采用逐 Token（per-token）量化以局部化异常值的影响；权重（Weights）分布平滑但对误差敏感，采用逐块（per-block）量化以兼顾精度与硬件优化；梯度（Gradients）动态范围大但绝对精度要求低，同样采用逐 Token 量化。这套组合拳在保证动态范围的同时，最大程度抑制了深层 Transformer 累积的量化噪声。

效果如何 实验在 H800 和 H20 集群上搭建，评估了稠密模型 Qwen3-4B 以及不同规模的 MoE 模型（使用 2 台 H20 运行 Qwen3-30B-A3B，使用 16 台 H20 运行 Qwen3-235B-A22B）。对比的最核心基线是“BF16 训练 + FP8 采样”的混合精度路线（代表当前常规的加速方案）。量化结果表明，在 30B 和 235B 规模的 MoE 模型上，统一 FP8（FP8-TI）方案均显著优于混合精度基线。在衡量离策略训练稳定性的截断重要性采样裁剪比例（TIS-clipfrac）指标上，FP8-TI 大幅降低了该数值，意味着被裁剪的更新更少，训练更稳定；同时，训练与采样对数概率的绝对差值（train_rollout_logprob_abs_diff）也变得更小且更平稳。作者还专门测试了模型规模的影响，发现在混合精度基线下，从 30B 到 1T 规模，上述两个指标均显著恶化，这反向证明了在大模型上使用统一 FP8 的必要性。
不过，作者也坦诚了该方案在当前框架下的局限与代价。首先是显存收益打折，由于 Transformer Engine 为了加速反向传播会保留量化权重的转置副本，且前向传播会保留高精度激活值副本，加上目前不支持将优化器 CPU 卸载与 FP8 参数收集同时开启，导致端到端 FP8 训练的显存占用甚至可能高于混合精度方案。其次是计算效率存在失效场景，在 RL 常用的小批量（如 batch_size=4）设置下，FP8 引入的额外量化与反量化操作会带来严重的 CPU 开销，导致 GPU 调度不饱和，此时 FP8 训练速度反而可能慢于 BF16。

主要贡献

在强化学习（RL）训练中，混合精度策略（如 BF16 训练配合 FP8 推理/采样）在面对混合专家（MoE）模型时暴露出严重的问题：随着模型规模的增加，训练与推理之间的不一致性会变得愈发严重。为了解决这一核心问题，本文提出并实现了一种端到端的全 FP8 强化学习训练与采样流水线（Unified FP8）。

本研究的核心目标是通过在训练和采样（Rollout）阶段统一使用 FP8 精度，从根本上消除由量化误差引起的训练-推理不一致性。实验表明，该方案不仅显著提升了 RL 训练的稳定性，还加速了训练过程。目前，针对 Qwen3-4B 和 Qwen3-30B-A3B 模型的 FP8 RL 训练已在 Miles 框架中得到全面支持【1，miles low_precision/URL】。本工作建立在 SGLang RL 团队和 Miles 社区早期关于严格零 KL 散度的 SGLang 与 FSDP 后端对齐【2，miles true_on_policy/URL】以及结合在线 SFT 的投机解码【3，Speculative Decoding/URL】等探索的基础之上。

背景知识与设计原则

Tensor Core 与低精度支持

低精度计算的硬件基础：低精度计算是软硬件协同设计的结晶，其核心硬件基础是 Tensor Core。这是一种专门为深度学习中的核心计算——大规模矩阵乘加（Matrix Multiplication and Accumulation）设计的 GPU 硬件加速单元。与传统的 CUDA 核心相比，Tensor Core 能够以高得多的吞吐量处理低精度数据格式（如 FP16、BF16、FP8）。

Tensor Core 的演进历程：其演进始于基础的 FMA（融合乘加）指令和通过 DP4A 实现的早期向量化。真正的里程碑是 Volta 架构，它首次引入了 Tensor Core 作为大规模矩阵运算的专用单元。此后，Ampere、Hopper 以及最新的 Blackwell 架构不断推进这一理念：

架构 FP64 FP16 INT8 INT4 FP8 MXFP
Volta ❌ ✅ FP16 ❌ ❌ ❌ ❌
Turing ❌ ✅ FP16 ✅ ✅ ❌ ❌
Ampere ✅ ✅ FP16/BF16 ✅ ✅ ❌ ❌
Hopper ✅ ✅ FP16/BF16 ✅ ❌ ✅ (累加精度仅支持 FP22【4，arXiv:2505.09343v1/2025】) ❌
Blackwell ✅ ✅ FP16/BF16 ✅ ❌ ✅ ✅ MXFP(8/6/4) NVFP4
Blackwell Ultra ✅ (降低的 FLOPs) ✅ FP16/BF16 ✅ (降低的 FLOPs) ❌ ✅ ✅ MXFP(8/6/4) NVFP4

低精度浮点格式的核心优势：在这种硬件趋势下，使用更低精度进行存储和计算变得极具吸引力。具体而言，低精度浮点格式提供了以下潜在优势：

  1. 显著减少显存占用：与主流的 BF16 相比，FP8 理论上可以将模型权重和激活值消耗的显存减半，直接缓解日益增长的 VRAM 压力。
  2. 理论上 2 倍的计算吞吐量：在主流 GPU（如 H100 SXM）上，FP8 Tensor Core 提供高达 1979 TFLOPS 的理论性能，是 BF16 单元（989 TFLOPS）的两倍。这一巨大的性能提升是推动 FP8 训练的关键因素。
  3. 缓解显存带宽瓶颈：由于数据表示更紧凑，需要从 GPU HBM 传输到计算核心的数据量更少。这意味着数据移动所花费的时间更少，有效地降低了显存带宽压力。

FP8 数据格式

FP8 格式的定义与分类：FP8 是一种使用 8 个比特位来表示数值的浮点格式。与 FP32（32 位）和 FP16/BF16（16 位）相比，FP8 可以将相同数据量的存储和传输成本降低到 1/4 或 1/2，极大地缓解了 VRAM 和带宽瓶颈，并提高了训练和推理性能。目前，主要有两种 FP8 格式：

FP8 E4M3 与 E5M2
FP8 E4M3 与 E5M2

图片来源：OCP 白皮书【5，OCP whitepaper/URL】

这种设计允许 FP8 在最大化硬件吞吐量的同时，保持足够的数值范围和精度。

FP8 缩放因子（Scale）的选择

FP32 与 E8M0 缩放因子的对比：

选择 FP32 作为训练缩放精度的原因：经过全面评估，本文最终选择 FP32 作为训练时的缩放精度。原因包括：

  1. 精度对齐与训练稳定性：FP32 缩放因子提供了细粒度的数值缩放，能够捕捉张量的动态范围，并使 FP8 训练的损失曲线尽可能接近 BF16 基线。
  2. 与推理生态系统的致性：主流推理模型也使用 FP32 作为量化缩放格式。
  3. 现实的硬件收益考量：
    • Hopper (H100/H800)：虽然支持 FP8 Tensor Core，但没有专门用于 E8M0 缩放的计算单元。
    • Blackwell (B100/B200)：引入了对 MXFP8（微缩放）的支持，为类似 E8M0 的块级缩放提供了硬件加速【6，arXiv:2506.08027/2025】。
      因此，在当前的 H 系列集群下，强制使用 E8M0 不仅无法带来明显的加速，反而可能引入额外的软件模拟开销和精度风险。

FP8 量化策略

量化步骤与原理：常见的量化策略包括 per-tensor（张量级）、per-block（块级）和 per-token（词元级）。无论粒度如何，量化通常遵循两个简单的步骤：

FP8 量化流程
FP8 量化流程

图片来源：InfiR2 论文【7，InfiR2: A Comprehensive FP8 Training Recipe for Reasoning-Enhanced Language Models/2025】

混合粒度量化策略的应用：由于 FP8 的精度低于 FP16/BF16，在实践中必须在训练稳定性和效率之间做出权衡，因此前向和反向传播通常使用不同的量化策略和粒度：

Megatron 中的混合粒度 FP8 量化
Megatron 中的混合粒度 FP8 量化

图片来源：InfiR2 论文【7，InfiR2: A Comprehensive FP8 Training Recipe for Reasoning-Enhanced Language Models/2025】

该图展示了 Megatron 中使用的混合粒度 FP8 策略与标准 BF16 流水线的对比。在 FP8 流水线中，应用了不同的量化方法：权重使用 per-block 量化（蓝色），而激活值使用 per-token 量化（紫色）。图中展示了完整的训练过程，包括前向传播（FProp）、权重梯度计算（Wgrad）和输入梯度计算（Dgrad），并详细说明了 FProp 的工作流程。

方法与问题分析细节

尽管 FP8 展现出巨大潜力，但在实际工程实践中（尤其是结合 Megatron-Core 和 TransformerEngine (TE) 时），面临着三大主要挑战。本文将这种统一的 FP8 训练和推理设置称为 FP8-TI (FP8 Training & Inference)。

挑战一：内存与计算效率的理论与现实差距

内存和计算效率未达预期的原因分析：在实践中，FP8 带来的内存节省和加速往往不如理论上那么显著，主要原因在于：

FP8 训练受限于 CPU
FP8 训练受限于 CPU

图：FP8 训练中受限于 CPU 的行为

挑战二：精度对齐与累积误差的影响

低精度导致的数值差异及其放大效应：FP8 的低精度特性不可避免地会引入相对于 BF16 的数值差异，这些差异在深层模型中会被放大并导致训练不稳定性问题：

挑战三：框架适配与 TransformerEngine 版本兼容性

Megatron-Core 与 TE 集成的工程挑战：除了算法挑战外，Megatron-Core 与 Transformer Engine (TE) 的集成方式还有待改进，特别是考虑到 TE 的快速迭代：

FP8 + RL：异常 KL 散度损失的归因分析

InfiXAI 团队已经成功在预训练和微调任务上运行了完整的 FP8 训练【7，InfiR2: A Comprehensive FP8 Training Recipe for Reasoning-Enhanced Language Models/2025】。在此基础上，将其应用于 RL。得益于 Miles 对 Megatron FP8 训练的良好支持，团队顺利运行了一系列 FP8 RL 实验。

观察到异常的初始 KL 散度损失：当直接从 BF16 切换到 FP8 并开始训练时，观察到一个引人注目的现象：与 BF16 训练相比，FP8 训练在第一步具有显著更高的 KL 散度损失（KL loss）。如下图所示，FP8-TI 的初始 KL 散度损失明显高于使用 FP8 推理的 BF16 训练（T 代表训练，I 代表推理）：

初始 KL 散度损失对比
初始 KL 散度损失对比

定位误差来源的实验设计：为了理解初始 KL 散度损失为何更高，分析了量化过程中两个潜在的误差来源：

  1. 量化计算内核的误差：来自特定 FP8 GEMM 实现的数值误差。
  2. 固有的量化误差：量化和反量化本身带来的精度损失。

量化计算内核的误差分析：
最初怀疑 TransformerEngine 中使用的闭源 cuBLAS GEMM 实现可能不如广泛使用的开源 DeepGEMM 准确。因此设计了实验，比较这两种 FP8 GEMM 实现相对于 BF16 的精度。在各种形状（基于 TE 的测试用例）下评估了它们的误差，结果如下：

内核形状 (M, K, N) cuBLAS(TE)误差 DeepGEMM误差
128, 128, 128 0.00068 0.00036
256, 128, 256 0.00068 0.00037
320, 128, 336 0.000684 0.00037
320, 64, 336 0.00067 0.00024
320, 256, 336 0.00068 0.00048
1024, 4096, 1024 0.000681 0.00065
2048, 2048, 512 0.00068 0.00063
1024, 1024, 1024 0.000683 0.0006

结果表明，两种 GEMM 实现的误差处于同一数量级，没有显著差异。因此，替换 TE 的 FP8 GEMM 并不能降低初始 KL 散度损失。

固有的量化误差分析：
对于第二个潜在来源，设计了一组对比实验来隔离并验证量化的固有误差：

FP8 量化误差分布
FP8 量化误差分布

该图展示了 GEMM 输出误差在一个完整迭代中的演变。

由此得出结论：

量化误差如何导致训练异常的机制验证：
基于上述实验，提出以下假设：

  1. 训练中的主要误差在量化步骤中就已经引入，并且数量可观。
  2. FP8 训练中较高的初始 KL 散度损失很可能来自于这种量化误差。
  3. 在混合 BF16 训练 + FP8 推理（rollout）中，相同的量化误差也会导致训练-推理不一致。

为了验证这些假设，修改了 Transformer Engine (TE) 并设计了以下实验：

验证假设 2 —— KL 散度损失分析：
下图显示了四种情况的 KL 散度损失曲线。可以看到，Case 2、Case 3 和 Case 4 (FP8-TI) 在第一步具有几乎相同的 KL 散度损失，且均显著高于 Case 1：

不同情况下的 KL 散度损失对比
不同情况下的 KL 散度损失对比

验证假设 3 —— TIS-clipfrac 分析：
引入截断重要性采样（Truncated Importance Sampling, TIS）中的 clipfrac 指标来验证假设 3。该指标反映了离策略（off-policy）训练的程度，即用于训练的模型与用于生成经验的模型之间的一致性。较高的 clipfrac 通常表示训练-推理不一致性更严重。

不同情况下的 TIS-clipfrac 对比
不同情况下的 TIS-clipfrac 对比

从图中可以看出，Case 2、Case 3 和 Case 4 (FP8-TI) 的 clipfrac 值大致处于同一数量级，且均显著低于 Case 1。这证实了：

  1. 初始 KL 散度损失升高的根本原因是量化误差。
  2. 与混合 BF16 训练 + FP8 rollout (Case 1) 相比，FP8-TI (Case 4) 可以显著缓解训练-推理不一致性。
  3. 对于训练偏差，前向传播中的量化误差比反向传播中的量化误差影响更大（如 Case 2 和 Case 3 的相似性所示）。同样，对于训练-推理一致性，前向传播量化误差是主要因素。

实验环境

实验结果

密集模型实验已经证明 FP8-TI 有效地抑制了训练-推理不一致性。在此基础上，Ant Group AQ 团队将研究扩展到 RL 中的 MoE 模型，以评估 FP8-TI 是否适用于更复杂的架构。实验发现 FP8-TI 具有以下优势：

  1. 降低 TIS clip fraction：其 TIS-clipfrac 显著低于 BF16 Train / FP8 Rollout，这意味着被截断的更新更少，训练稳定性更高。
  2. 缩小训练-rollout对数概率差距：与 BF16 Train / FP8 Rollout 相比，这种 FP8 方案在训练和 rollout 的对数概率之间产生了更小且更稳定的差异。

MoE 实验设计

为了隔离变量进行清晰的比较，设置了两种实验方案：

关键评估指标：

MoE 结果与分析

实验一：Qwen3-30B-A3B 结果
在 30B 规模的 MoE 模型上，结果清楚地显示了 FP8-TI 的优势：

Qwen3-30B-A3B TIS-clipfrac
Qwen3-30B-A3B train_rollout_logprob_abs_diff

实验二：Qwen3-235B-A22B 结果
为了评估可扩展性，在 235B 规模的模型上复制了实验，并得出了高度一致的结论：

Qwen3-235B-A22B TIS-clipfrac
Qwen3-235B-A22B train_rollout_logprob_abs_diff

MoE 实验结论：对于 MoE RL 任务，与 BF16 Train / FP8 Rollout 相比，在训练和推理中统一使用 FP8 提高了训练稳定性，并有效地抑制了训练-推理不一致性。从 30B 到 235B 的 MoE 模型中，一致地观察到了这一优势。

MoE 模型规模对训练-推理不一致性的影响

进一步研究了在混合精度设置（BF16 训练 / FP8 Rollout）下，MoE 模型大小如何影响训练-推理不一致性。实验表明，随着 MoE 模型规模的增加，训练-推理不一致性变得更加严重。

如下图所示，从 30B 到 1T 规模，TIS-clipfrac 和 Train_rollout_logprob_abs_diff 均显著增加。这表明，对于 BF16 Train / FP8 Rollout 方案，越大的模型往往遭受越严重的训练-推理不一致性，这间接突显了采用统一精度方案（如 FP8-TI）的必要性和重要性。

不同模型规模下的 TIS-clipfrac
不同模型规模下的 train_rollout_logprob_abs_diff

结论与未来工作

本文提出并验证了在强化学习中应用端到端全 FP8 流水线（FP8-TI）的有效性。通过实验证实，量化误差是导致混合精度策略下初始 KL 散度异常及训练-推理不一致的核心原因。采用 Unified FP8 策略，能够显著抑制这种不一致性，提升 MoE 模型的 RL 训练稳定性和速度，且该优势在 30B 至 235B 的模型中具备良好的可扩展性。

未来工作将集中在以下几个方向：

  1. 更深入地研究训练-推理不一致性，分析其根本原因，并探索更好的解决方案。
  2. 更彻底地研究量化策略，理解量化误差是如何产生的，并设计具有更低误差的方案。
  3. 通过更好的算法、框架以及软硬件协同设计来提高低精度训练效率，隐藏内核启动和量化的延迟，真正实现训练和推理的双重加速。

致谢：感谢 InfiXAI Team (Congkai Xie, Mingfa Feng, Shuo Cai)、Ant Group AQ Team (Yanan Gao, Zhiling Ye, Hansong Xiao)、SGLang RL Team (JiLi, Yefei Chen, Xi Chen, Zilin Zhu)、Miles Team (Chenyang Zhao) 的共同努力，以及 NVIDIA (Juan Yu, NeMo-RL Team) 的技术支持。

参考文献汇总