Unified FP8: Moving Beyond Mixed Precision for Stable and Accelerated MoE RL
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$:
其中 $V_{\max}$ 为 FP8 能表示的最大值。随后利用该缩放因子计算量化后的值 $Q(x)$:
效果如何 实验在 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 架构不断推进这一理念:
- 扩大规模:允许 Tensor Core 每次操作处理更大的矩阵,从而提高计算访存比。
- 降低精度:不断增加对 FP/BF16、FP8 甚至更低精度数据格式的支持。
| 架构 | 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 |
低精度浮点格式的核心优势:在这种硬件趋势下,使用更低精度进行存储和计算变得极具吸引力。具体而言,低精度浮点格式提供了以下潜在优势:
- 显著减少显存占用:与主流的 BF16 相比,FP8 理论上可以将模型权重和激活值消耗的显存减半,直接缓解日益增长的 VRAM 压力。
- 理论上 2 倍的计算吞吐量:在主流 GPU(如 H100 SXM)上,FP8 Tensor Core 提供高达 1979 TFLOPS 的理论性能,是 BF16 单元(989 TFLOPS)的两倍。这一巨大的性能提升是推动 FP8 训练的关键因素。
- 缓解显存带宽瓶颈:由于数据表示更紧凑,需要从 GPU HBM 传输到计算核心的数据量更少。这意味着数据移动所花费的时间更少,有效地降低了显存带宽压力。
FP8 数据格式
FP8 格式的定义与分类:FP8 是一种使用 8 个比特位来表示数值的浮点格式。与 FP32(32 位)和 FP16/BF16(16 位)相比,FP8 可以将相同数据量的存储和传输成本降低到 1/4 或 1/2,极大地缓解了 VRAM 和带宽瓶颈,并提高了训练和推理性能。目前,主要有两种 FP8 格式:
- E4M3:4 位指数 + 3 位尾数。动态范围较小,但精度较高。
- E5M2:5 位指数 + 2 位尾数。动态范围较大,但精度较低。
图片来源:OCP 白皮书【5,OCP whitepaper/URL】
这种设计允许 FP8 在最大化硬件吞吐量的同时,保持足够的数值范围和精度。
FP8 缩放因子(Scale)的选择
FP32 与 E8M0 缩放因子的对比:
- 格式定义:FP32 为 IEEE 754 单精度浮点数;E8M0 为仅包含 8 位指数、0 位尾数的格式。
- 数值特性:FP32 可以表示任意精度的实数;E8M0 仅支持 2 的幂(如 1, 2, 0.5),无法表示 1.5 这样的数值。
- 核心思想:FP32 旨在通过高精度管理缩放因子,确保训练期间的数值稳定性;E8M0 则将缩放因子带入低精度领域,利用位运算提升效率。
- 主要优势:FP32 精度高、训练稳定(准确捕捉动态范围,减少量化误差,防止发散),且生态支持广泛(是 NVIDIA Transformer Engine 等主流库的默认选择);E8M0 对硬件极其友好(缩放可实现为简单的位移操作,速度快且节能),并且能实现统一流水线(整个流水线包括 scale 都在 8 位下运行,简化硬件设计)。
- 主要劣势:FP32 存在存储开销(每个量化张量需额外存储一个 FP32 的 scale 值)和计算开销(Scale 的计算和转换必须在 FP32 下进行);E8M0 存在精度损失风险(强制舍入到 2 的幂会引入量化噪声,在反向传播中可能累积并导致发散),且动态范围分辨率有限(难以精细适应复杂的张量分布)。
选择 FP32 作为训练缩放精度的原因:经过全面评估,本文最终选择 FP32 作为训练时的缩放精度。原因包括:
- 精度对齐与训练稳定性:FP32 缩放因子提供了细粒度的数值缩放,能够捕捉张量的动态范围,并使 FP8 训练的损失曲线尽可能接近 BF16 基线。
- 与推理生态系统的致性:主流推理模型也使用 FP32 作为量化缩放格式。
- 现实的硬件收益考量:
- 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(词元级)。无论粒度如何,量化通常遵循两个简单的步骤:
图片来源:InfiR2 论文【7,InfiR2: A Comprehensive FP8 Training Recipe for Reasoning-Enhanced Language Models/2025】
-
步骤 1:计算缩放因子 $S$。取给定张量(或块)中的最大绝对值 $\max|X|$,并将其除以 FP8 的最大可表示值 $V_{\max}$:
$S = \frac{\max|X|}{V_{\max}}$ -
步骤 2:计算量化值 $Q$。使用缩放因子 $S$,将原始张量 $X$ 中的每个元素 $x$ 除以 $S$ 并四舍五入到最近的整数,以获得量化值:
$Q(x) = \mathrm{round}\left(\frac{x}{S}\right)$
混合粒度量化策略的应用:由于 FP8 的精度低于 FP16/BF16,在实践中必须在训练稳定性和效率之间做出权衡,因此前向和反向传播通常使用不同的量化策略和粒度:
- 激活值(Activations):通常使用 per-token 量化。激活值通常包含显著的异常值;更细的量化粒度可以局部化异常值的影响,并更好地保留整体精度。
- 权重(Weights):通常使用 per-block 量化。收敛后,权重分布通常是平滑的(接近高斯分布),异常值很少,但对量化误差高度敏感。块级量化(例如 block_size × block_size)在保持精度的同时,能很好地配合硬件优化,平衡了计算效率和内存节省。
- 梯度(Gradients):通常使用 per-token 量化。梯度具有较大的动态范围变化,但绝对精度要求相对较低。历史上,大多数方案使用 per-tensor E5M2 来确保动态范围,但 DeepSeek-V3 表明细粒度的 E4M3 也能平衡精度和范围。
图片来源: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 带来的内存节省和加速往往不如理论上那么显著,主要原因在于:
-
有限的内存优化:
- 冗余的权重副本:为了加速反向传播,TransformerEngine 会保留一份额外的转置量化权重副本。这阻碍了权重内存使用量按预期的 2 倍比例减少。
- 高精度的激活值副本:在注意力层和激活层的前向传播中,框架通常会保留一份高精度的激活值副本,以便后续进行准确的梯度计算。FP8 并不能减少这部分内存占用。
-
计算效率瓶颈:
- 小批量(small batch sizes)下的性能退化:当 batch_size 较小时,FP8 训练可能无法充分利用 GPU 计算单元,甚至表现不如 BF16。根本原因是 FP8 引入了额外的量化和反量化操作,增加了 CPU 开销。在通常使用小批量(例如 batch_size=4)的 Agentic RL 场景中,这个问题尤为突出——频繁的 CPU 开销会导致 FP8 训练比 BF16 更慢。(如下图所示,GPU 内核调度不密集;通常 GPU 已经完成了前一个工作,但下一个内核启动被延迟,因为系统处于 CPU 瓶颈状态。)
图:FP8 训练中受限于 CPU 的行为
挑战二:精度对齐与累积误差的影响
低精度导致的数值差异及其放大效应:FP8 的低精度特性不可避免地会引入相对于 BF16 的数值差异,这些差异在深层模型中会被放大并导致训练不稳定性问题:
- 固有的量化误差:即使在 FP32 下进行累加,对单个 GEMM 操作的 FP8 输入进行量化也会引入误差。实验表明,与 BF16 GEMM 相比,典型的误差约为 0.0007。
- 逐层累积效应:在深层 Transformer 模型中,这些微小的误差会在前向和反向传播期间逐层累积:
- 在预训练和微调(SFT)中:梯度主要由真实标签(ground-truth)的对数概率主导。通过细粒度的块级量化,误差通常可以控制在可接受的范围内,模型不太可能崩溃。
- 在强化学习(RL)中:梯度通常由两次前向传播的两个对数概率之间的差异决定。在这种情况下,累积的 FP8 误差会被放大,导致梯度偏离其理想方向并影响收敛效率——甚至将优化推向“错误的方向”。
挑战三:框架适配与 TransformerEngine 版本兼容性
Megatron-Core 与 TE 集成的工程挑战:除了算法挑战外,Megatron-Core 与 Transformer Engine (TE) 的集成方式还有待改进,特别是考虑到 TE 的快速迭代:
- 版本依赖和迁移开销:TE 的快速迭代带来了新功能,但也带来了严格的版本依赖。在实践中发现,即使是相同的训练脚本,在不同的 TE 版本中也可能产生不同的数值行为,有时需要调整代码以避免 NaN 等问题。
- 特定架构的成熟度:对所有主流模型架构的全面 FP8 支持是一个持续的过程。对于一些非标准或较新的组件(如 MLA),观察到 FP8 训练支持仍在成熟中。即使在较后的版本(如 2.4.0 → 2.8.0)中,某些错误和限制仍有待解决。
- 与内存优化策略的冲突:在 RL 训练中,启用优化器 CPU 卸载(Optimizer CPU Offload)可以显著减少内存使用,但当前的 TE 不支持将其与
--fp8-param-gather一起使用。由于此限制,端到端 FP8 训练最终可能会消耗比使用 FP8 rollout 的 BF16 训练更多的内存,这需要社区和维护者的进一步优化。
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 散度损失为何更高,分析了量化过程中两个潜在的误差来源:
- 量化计算内核的误差:来自特定 FP8 GEMM 实现的数值误差。
- 固有的量化误差:量化和反量化本身带来的精度损失。
量化计算内核的误差分析:
最初怀疑 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 散度损失。
固有的量化误差分析:
对于第二个潜在来源,设计了一组对比实验来隔离并验证量化的固有误差:
- 基线环境:单台 H800 上的 Qwen3-4B。
-
实验模式设计:
- Baseline(基线):权重和输入使用 BF16,使用 BF16 GEMM。
- FP8 Real Quant(真实量化):权重和输入使用 FP8,使用 FP8 GEMM(例如 DeepGEMM/cuBLAS GEMM;主要测试 cuBLAS 以避免对 TE 进行大量修改)。
- FP8 Fake Quant(伪量化):权重和输入保持在 BF16,但模拟量化过程(量化到 FP8 然后反量化回 BF16),最后使用 BF16 GEMM。
-
对比维度:
- FP8 Real Quant vs. FP8 Fake Quant:用于验证 FP8 GEMM 内核(cuBLAS)的精度,隔离实现带来的任何额外误差。
- Baseline vs. FP8 Fake Quant:用于忽略 GEMM 内核的影响,专注于量化/反量化本身引入的固有误差。
-
评估指标:收集 RL 训练开始时(Step 0 和 Step 1)所有 GEMM 操作的输出差异(Diff)。
- 结果分析:下图可视化了在一个完整的前向+反向传播过程中,按执行顺序排列的所有 GEMM 输出的误差分布:
该图展示了 GEMM 输出误差在一个完整迭代中的演变。
- 灰色/高点 (Baseline vs. FP8 Fake Quant):代表量化本身的误差。可以看到 BF16 基线和伪量化之间存在显著差异。
- 绿色/低点 (FP8 Real Quant vs. FP8 Fake Quant):代表内核实现产生的误差。这些差异极小,接近于零。
由此得出结论:
- 误差主要来自量化原理,而非内核实现:Fake Quant 和 Real Quant 都与基线存在显著差异(相差两个数量级),这强烈表明主要的误差来源是有损的量化/反量化本身,而不是计算过程。
- FP8 GEMM 内核高度可靠:Real Quant 和 Fake Quant 输出之间的微小差异表明,TE 中使用的 cuBLAS FP8 GEMM 极其准确,并且非常接近理想的数学模拟,可以安全地用于生产环境。
量化误差如何导致训练异常的机制验证:
基于上述实验,提出以下假设:
- 训练中的主要误差在量化步骤中就已经引入,并且数量可观。
- FP8 训练中较高的初始 KL 散度损失很可能来自于这种量化误差。
- 在混合 BF16 训练 + FP8 推理(rollout)中,相同的量化误差也会导致训练-推理不一致。
为了验证这些假设,修改了 Transformer Engine (TE) 并设计了以下实验:
- 基线环境:H800 集群上的 Qwen3-4B。
- 测试案例(Cases):
- Case 1:BF16 训练,FP8 rollout(推理)。
- Case 2:BF16 训练,FP8 rollout;在训练的前向传播中,将 BF16 权重和激活值量化为 FP8,然后在运行 BF16 GEMM 之前反量化回 BF16。
- Case 3:BF16 训练,FP8 rollout;在训练的前向和反向传播中,将输入矩阵 A 和 B 都量化为 FP8,然后在 BF16 GEMM 之前反量化回 BF16。
- Case 4 (FP8-TI):FP8 训练,FP8 rollout。
验证假设 2 —— KL 散度损失分析:
下图显示了四种情况的 KL 散度损失曲线。可以看到,Case 2、Case 3 和 Case 4 (FP8-TI) 在第一步具有几乎相同的 KL 散度损失,且均显著高于 Case 1:
验证假设 3 —— TIS-clipfrac 分析:
引入截断重要性采样(Truncated Importance Sampling, TIS)中的 clipfrac 指标来验证假设 3。该指标反映了离策略(off-policy)训练的程度,即用于训练的模型与用于生成经验的模型之间的一致性。较高的 clipfrac 通常表示训练-推理不一致性更严重。
从图中可以看出,Case 2、Case 3 和 Case 4 (FP8-TI) 的 clipfrac 值大致处于同一数量级,且均显著低于 Case 1。这证实了:
- 初始 KL 散度损失升高的根本原因是量化误差。
- 与混合 BF16 训练 + FP8 rollout (Case 1) 相比,FP8-TI (Case 4) 可以显著缓解训练-推理不一致性。
- 对于训练偏差,前向传播中的量化误差比反向传播中的量化误差影响更大(如 Case 2 和 Case 3 的相似性所示)。同样,对于训练-推理一致性,前向传播量化误差是主要因素。
实验环境
- 模型架构:包含 Qwen3-4B(密集模型)、Qwen3-30B-A3B(MoE 模型,30B 规模)、Qwen3-235B-A22B(MoE 模型,235B 规模)。
-
硬件配置:
- Qwen3-4B 实验:H800 集群(单机及集群)。
- Qwen3-30B-A3B 实验:2 台 H20 服务器。
- Qwen3-235B-A22B 实验:16 台 H20 服务器。
-
软件配置:基于 Miles 框架,底层依赖 Megatron-Core 和 NVIDIA TransformerEngine (TE) 库,并使用 cuBLAS 和 DeepGEMM 进行 FP8 算子计算。
实验结果
密集模型实验已经证明 FP8-TI 有效地抑制了训练-推理不一致性。在此基础上,Ant Group AQ 团队将研究扩展到 RL 中的 MoE 模型,以评估 FP8-TI 是否适用于更复杂的架构。实验发现 FP8-TI 具有以下优势:
- 降低 TIS clip fraction:其 TIS-clipfrac 显著低于 BF16 Train / FP8 Rollout,这意味着被截断的更新更少,训练稳定性更高。
- 缩小训练-rollout对数概率差距:与 BF16 Train / FP8 Rollout 相比,这种 FP8 方案在训练和 rollout 的对数概率之间产生了更小且更稳定的差异。
MoE 实验设计
为了隔离变量进行清晰的比较,设置了两种实验方案:
- Case 1(混合精度):BF16 训练,FP8 rollout。
- Case 2(统一精度):FP8 训练,FP8 rollout。
关键评估指标:
- TIS clip fraction (TIS-clipfrac):衡量离策略训练的稳定性;越低越好。
- 训练和 rollout 对数概率的绝对差值 (train_rollout_logprob_abs_diff):衡量模型在训练和 rollout 行为之间的一致性;越小且越稳定越好。
MoE 结果与分析
实验一:Qwen3-30B-A3B 结果
在 30B 规模的 MoE 模型上,结果清楚地显示了 FP8-TI 的优势:
- 更低的 TIS-clipfrac:FP8-TI 实现了比 BF16 Train / FP8 Rollout 基线显著更低的 TIS-clipfrac,表明被截断的更新更少,训练更稳定。
- 更小的训练-rollout对数概率差距:FP8-TI 为 Train_rollout_logprob_abs_diff 产生了更窄、更稳定的范围,表明训练和推理之间的行为更加一致。
实验二:Qwen3-235B-A22B 结果
为了评估可扩展性,在 235B 规模的模型上复制了实验,并得出了高度一致的结论:
- TIS-clipfrac 和训练-推理差异的持续改善:如下图所示,即使在 235B 规模下,与 BF16 Train / FP8 Rollout 相比,FP8-TI 继续降低了 TIS-clipfrac 和 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)的必要性和重要性。
结论与未来工作
本文提出并验证了在强化学习中应用端到端全 FP8 流水线(FP8-TI)的有效性。通过实验证实,量化误差是导致混合精度策略下初始 KL 散度异常及训练-推理不一致的核心原因。采用 Unified FP8 策略,能够显著抑制这种不一致性,提升 MoE 模型的 RL 训练稳定性和速度,且该优势在 30B 至 235B 的模型中具备良好的可扩展性。
未来工作将集中在以下几个方向:
- 更深入地研究训练-推理不一致性,分析其根本原因,并探索更好的解决方案。
- 更彻底地研究量化策略,理解量化误差是如何产生的,并设计具有更低误差的方案。
- 通过更好的算法、框架以及软硬件协同设计来提高低精度训练效率,隐藏内核启动和量化的延迟,真正实现训练和推理的双重加速。
致谢:感谢 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) 的技术支持。
参考文献汇总
- 【1】miles low_precision 支持代码库:https://github.com/radixark/miles/tree/main/examples/low_precision (引用于主要贡献,提及 Qwen3-4B 和 Qwen3-30B-A3B 的 FP8 RL 训练已开箱即用)
- 【2】miles true_on_policy 代码库:https://github.com/radixark/miles/tree/main/examples/true_on_policy (引用于主要贡献,提及 SGLang 和 FSDP 后端对齐以实现严格零 KL 散度)
- 【3】Speculative Decoding 教程:https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/slime/spec/readme-en.md (引用于主要贡献,提及为草稿模型提供在线 SFT 的投机解码技术)
- 【4】Hopper FP22 累加精度说明:[arXiv:2505.09343v1/2025] (引用于 Tensor Core 演进表格,说明 Hopper 架构的累加精度限制)
- 【5】OCP microscaling formats 白皮书:https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf (引用于 FP8 格式图表,展示 E4M3 与 E5M2 的差异)
- 【6】MXFP8 块级缩放硬件加速:[arXiv:2506.08027/2025] (引用于 FP8 缩放因子选择,说明 Blackwell 架构引入的 MXFP8 支持)
- 【7】InfiR2: A Comprehensive FP8 Training Recipe for Reasoning-Enhanced Language Models:https://arxiv.org/html/2509.22536v4 / 2025 (引用于 FP8 量化流程图、混合粒度量化策略图,以及证明 FP8 在预训练和微调任务上的可行性)
💬 评论讨论
欢迎在这里分享您的想法和见解!