Towards Blackwell-Native 8-bit and 4-bit RL: End-to-End MXFP8 and NVFP4 RL in Miles

发表时间: 2026-07 · Blog post by Ziang Li, humans& and Miles Team (lmsys.org)

原文: https://lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl

作者/机构:Ziang Li, humans& and Miles Team

速读

一句话结论 本文在 Miles 框架中实现了面向 Blackwell 架构的原生端到端 MXFP8 和针对 MoE 专家的逐 token NVFP4 强化学习训练方案,在保证奖励曲线与 BF16 基线高度重合的同时,显著降低了经验采样的时间开销。

要解决什么问题 在低精度强化学习中,经验采样(rollout)、模型训练、权重转换和在线权重更新必须遵循完全一致的精度契约,否则采样策略和训练策略会产生分歧,导致性能崩溃。此前的低精度方案并非为 Blackwell 架构的原生微缩放格式设计,要么依赖软件层面的 FP32 缩放,要么在计算路径上仍需反量化回 BF16,无法充分利用新硬件的稠密 Tensor Core 算力。此外,直接套用现有的 NVFP4 预训练方案在强化学习中会失效,因为强化学习的梯度充满噪声且奖励方差大,微弱的有效更新信号极易被量化噪声淹没;同时,预训练中常用的张量级或跨 token 共享激活缩放会使量化结果受到 batch 组成的影响,甚至导致未来 token 的信息泄漏到过去 token 的表示中,这对于序列长度多变且高度依赖调度的强化学习采样来说是致命的机制卡点。

怎么做的 核心思路是构建贯穿整个强化学习栈(从 Megatron 训练到 SGLang 采样)的比特级精确量化契约,并针对强化学习的脆弱性对 Blackwell 原生格式进行定制。方法由四个关键部件构成。第一是端到端 MXFP8 方案,前向传播、经验采样以及权重和数据梯度的矩阵乘法全部采用 MXFP8 格式。MXFP8 是一种一维微缩放格式,每 32 个连续的 E4M3 值共享一个 E8M0 缩放因子。为避免反向传播时额外的反量化和重量化误差,训练端在量化时会同时物化按行和按列的量化副本。第二是针对 MoE(混合专家)权重的逐 token NVFP4 方案。由于 MoE 专家占据了极大的参数量(如 DeepSeek-V3 中占比约 97.8%),仅量化专家层即可获得大部分显存收益。NVFP4 采用双层缩放结构,其量化表示可抽象为 $X \approx s_{\text{token}} \times s_{\text{block}} \times X_{\text{FP4}}$,其中 $X_{\text{FP4}}$ 是 16 个 E2M1 格式的值,$s_{\text{block}}$ 是对应的 E4M3 块缩放因子。为绕开跨 token 缩放带来的信息泄漏卡点,该方案在算子内部在线计算逐 token 的 FP32 粗粒度缩放因子 $s_{\text{token}}$,将激活异常值隔离在单个 token 内部。第三是高精度反向传播机制。该方案摒弃了预训练中使用的随机舍入和随机哈达玛变换,转而提供两种反向模式:高精度模式(反向矩阵乘法直接使用原始 BF16 操作数)和反量化模式(反向矩阵乘法使用前向低精度操作数的 BF16 反量化值),以此牺牲部分反向吞吐量来换取计算精度。第四是细粒度精度控制与比特级对齐。强制 FlashInfer 和 TransformerEngine 采用完全一致的量化器逻辑,并允许将最后 15% 的网络层、共享专家以及多头潜在注意力的特定投影矩阵(如 kv_b_proj,以避免一维块在不同收缩轴下的量化错位)保持在 BF16 精度,从而在底层阻断训练与推理的精度漂移。

效果如何 实验在 8 张 B200 GPU 上搭建(4 张用于采样,4 张用于训练),使用 Qwen3-30B-A3B 模型在 dapo-math-17k 数据集上进行 GRPO 风格的同步强化学习训练,每个提示生成 8 个采样,最大响应长度为 8192 个 token。对比基线共六组,包括代表性能上限的纯 BF16 训练与采样基线,以及代表本文路线的五组低精度变体:端到端 MXFP8、带高精度反向的 MXFP8、带反量化反向的 MXFP8、带高精度反向的逐 token NVFP4 MoE、带反量化反向的逐 token NVFP4 MoE。量化结果表明,在保持最后 15% 层为 BF16 的设置下,所有五种低精度配置的原始奖励曲线均与 BF16 基线高度重合,成功维持了学习轨迹。在性能上,MXFP8 和 NVFP4 均显著降低了采样时间,且 MXFP8 的训练速度快于 BF16。但该方案也存在代价与局限性:由于 TransformerEngine 尚未完全集成融合算子,当前 NVFP4 的逐 token 缩放仍依赖独立的 PyTorch 操作,导致其训练端耗时反而高于 BF16;反量化反向模式会引入额外的反量化步骤开销;高精度反向 NVFP4 偶尔仍会出现梯度尖峰,反量化模式能缓解但无法彻底消除;此外,由于底层对 1D 1x16 权重的低精度参数收集支持尚不成熟,Megatron 目前仍需保存一份额外的 BF16 权重副本,这限制了该方案在实际应用中的显存收益上限。

主要贡献

在低精度强化学习(RL)中,rollout(经验回放/采样)、训练、检查点转换和实时权重更新必须遵循统一的精度契约,否则采样器和训练器策略将会产生分歧。先前的研究已经证明了在Hopper架构下使用FP8或INT4进行RL的可行性,但这些方法并非针对最新的Blackwell架构原生设计。本文旨在将低精度RL扩展到Blackwell原生格式,通过在Miles框架及SGLang、TransformerEngine、FlashInfer、Megatron和cuDNN前端等生态组件中实现MXFP8和NVFP4配方,解决低精度RL中的训练-推理不匹配问题。

本文的主要贡献包括:

  1. 实现端到端的MXFP8 RL,其中rollout、前向传播、权重梯度GEMM和数据梯度GEMM全部使用MXFP8格式。
  2. 针对MoE专家权重实现基于单token(per-token)激活缩放的NVFP4 RL。
  3. 引入细粒度的精度控制机制,确保选定的关键张量(如最终层)能一致地保持在BF16精度。
  4. 在TransformerEngine和FlashInfer之间建立了位级精确(bit-exact)的量化器契约,从而避免权重更新引入不必要的训练-推理不匹配。

背景知识与设计原则

为何需要Blackwell原生方案:之前的低精度方法并非围绕MXFP8或NVFP4设计。现有的Miles路径遵循DeepSeek-V3风格的块缩放FP8配方:权重使用128x128块缩放,激活使用1x128切片缩放,并且在线计算每个切片或块的缩放因子。这是一个强大的Hopper架构时代的配方,但在Blackwell架构上,其FP32缩放因子仍然是在Tensor Core路径周围通过软件应用的,而不是通过原生的微缩放(microscaling)硬件。

INT4 QAT的局限性:INT4 QAT(量化感知训练)解决的是另一个问题。训练时使用伪量化使模型适应INT4权重,而rollout使用W4A16。尽管这种方式内存效率高,但计算路径实际上仍然使用BF16激活和反量化后的INT4权重。下表将NVIDIA HGX平台的密集Tensor Core规格标准化为单GPU吞吐量(B200和B300来自8-GPU HGX系统,Rubin来自HGX Rubin NVL8表格),展示了不同精度的算力差异。

GPU BF16 密集 Tensor Core FP8 密集 Tensor Core FP4 密集 Tensor Core
B200 2.25 PFLOPS 4.5 PFLOPS 9 PFLOPS
B300 2.25 PFLOPS 4.5 PFLOPS 13.5 PFLOPS
Rubin GPU (NVL8) 4 PFLOPS 17.5 PFLOPS 35 PFLOPS

对于RL系统,精度契约必须跨越SGLang rollout、Megatron与TransformerEngine训练、Hugging Face检查点转换、Megatron到Hugging Face的实时权重导出,以及细粒度的高精度异常处理。

MXFP8格式背景:MXFP8是一种微缩放FP8格式。TransformerEngine的MXFP8文档将其描述为Blackwell原生的逐块缩放配方:每32个连续的E4M3值共享一个局部的E8M0缩放因子,并且该块是一维的。因为E8M0缩放因子表示2的幂,所以解码后的缩放因子通常向上取整,以确保块中的最大值不会被截断。

NVFP4格式背景:NVFP4是Blackwell原生的FP4格式。正如NVIDIA的NVFP4介绍所述,它将FP4 E2M1值与每16个值块的一个FP8 E4M3缩放因子一起存储。因为E4M3的分辨率比UE8M0更细,其缩放因子通常被舍入到最接近的可表示值。标准的NVFP4配方还为更大的张量范围添加了一个FP32缩放因子,从而创建了一个两级层次结构:首先是一个粗粒度的FP32缩放因子,将张量或token映射到NVFP4的可表示范围内;其次是一个细粒度的E4M3缩放因子,用于适配每个1x16的块。

带有FP32张量缩放和E4M3块缩放的NVFP4两级缩放
带有FP32张量缩放和E4M3块缩放的NVFP4两级缩放

FP32缩放范围的选择:FP32缩放因子可以在不同的张量范围上进行选择。这种选择是一个配方层面的决定,而不是格式本身的属性,这对于强化学习(RL)来说尤为重要。

方法细节

配方1:端到端MXFP8 RL的设计:MXFP8配方是早期端到端FP8工作最直接的Blackwell原生扩展。在这个配方中,rollout、前向传播、权重梯度GEMM和数据梯度GEMM全部使用MXFP8,而选定的张量通过下文描述的精度控制规则保持为BF16。

端到端MXFP8 RL配方
端到端MXFP8 RL配方

训练端的MXFP8实现:TransformerEngine和Megatron将MXFP8实现为性能优化的、一等的Blackwell训练路径,这包含了DeepSeek-V3 GB200优化文档中描述的优化。在Miles的集成中,研究人员使用这条路径作为端到端MXFP8 RL的训练端基础。

反向激活表示的权衡:该方法与DeepSeek-V3 FP8配方的一个区别在于反向激活的表示方式。DeepSeek-V3将前向激活存储为1x128的FP8切片,并在反向GEMM之前将其转换为反向方向。这种方法存储的FP8数据较少,但会在反向GEMM之前引入反量化加上重新量化的过程。TransformerEngine的MXFP8文档指出,按行划分的1x32块和按列划分的32x1块在数值上是不同的,必须从全精度数据中独立量化。因此,TransformerEngine在量化期间同时生成并保留按行和按列的量化副本。这种方式使用了更多的内存,但避免了额外的重新量化步骤,并减少了反向路径中的额外量化误差。这在RL中是一个典型的系统权衡,研究人员选择使用TransformerEngine路径以保持单一的端到端MXFP8契约,避免引入另一个不匹配的来源。

Rollout端的MXFP8实现:在rollout端,SGLang使用了来自FlashInfer和Triton的Blackwell MXFP8内核。研究人员在FlashInfer和SGLang中实现并向上游提交了rollout路径。几乎所有主要的GEMM都可以量化为MXFP8,包括注意力机制的投影层和MoE专家层。主要的例外是显式控制的高精度层,例如下文描述的BF16 MLA投影层。

配方2:单Token NVFP4 RL的设计:NVFP4比MXFP8更为激进,因此研究人员选择性地应用它。研究人员仅对MoE专家层进行量化,因为它们主导了模型大小和rollout内存流量,而模型的其余部分保持为BF16(除非显式配置为其他格式)。例如,DeepSeek-V3总共有约671B参数。其MoE专家占据了 $(61 - 3) \times (256 + 1) \times 3 \times 7168 \times 2048 / 10^9 = 656.5B$ 个参数。这大约占模型的97.8%。因此,以MoE专家为目标可以捕获大部分内存收益,而无需强制将每一层都转换为最激进的精度格式。

为何不直接使用NVFP4预训练配方:原始的NVFP4预训练配方专为大规模预训练设计,其目标是在使用FP4 GEMM的同时,在许多token上保留粗略的优化方向。它将FP4线性层GEMM与几个稳定器结合在一起:选定的层保持较高精度,前向和反向的权重缩放保持一致,并且训练路径使用随机舍入(Stochastic Rounding, SR)和随机哈达玛变换(Random Hadamard Transforms, RHT)。在原论文中,SR应用于梯度以减少量化偏差并产生无偏的量化梯度,而RHT用于分散大幅值的块级异常值,特别是对于权重梯度GEMM的输入。

原始NVFP4预训练配方
原始NVFP4预训练配方

预训练与RL的数值环境对比:这虽然是一个很好的起点,但RL具有不同的失效面。在预训练中,梯度信号稳定,权重更新幅度大,模型权重具有适应性,且激活和梯度对精度敏感且具有高动态范围,底线是保留粗略的优化方向和收敛性。而在RL中,梯度充满噪声,奖励方差大,有用的更新微小且脆弱,底线是必须将量化噪声保持在真实的更新信号之下;否则它可能会覆盖脆弱的能力并导致性能崩溃。因此,NVFP4 RL配方并未包含预训练配方的所有部分。研究人员的目标是MoE专家权重量化、单token(per-token)激活缩放、一致的精度控制,以及具有可选原始或反量化操作数的BF16反向GEMM。

单Token激活缩放的必要性:两级NVFP4层次结构非常强大,但必须仔细选择FP32激活缩放的作用范围。正如Cursor Composer 2技术报告中所讨论的,单张量(per-tensor)的NVFP4缩放会使训练变得随批次变化(batch-variant),而跨token的缩放共享可能会将未来token的信息泄漏到过去token的表示中。如果一个token与其他token共享其缩放因子,其量化表示将取决于批次的组成。这对于RL来说尤其成问题,因为rollout调度和序列长度是变化的。

单Token缩放的在线计算实现:因此,本配方在线为每个token计算一个FP32激活缩放因子。这将激活异常值局部化到一个token,消除了静态激活校准的伪影,并让SGLang rollout和Megatron训练使用相同的激活缩放范围。在rollout端,单token的FP32缩放计算被融合到FlashInfer的激活量化内核路径中:发出打包的FP4激活和E4M3块缩放因子的同一个调用,也会返回单token的FP32缩放因子。因此,单token激活缩放不需要单独的校准缩放传递。

并行度与张量转换的一致性契约:训练与推理的一致性还需要匹配的并行度。如果FP32缩放因子是在专家张量并行(ETP)分区内按token计算的,SGLang和Megatron应该使用相同的ETP大小。否则,每一端可能会看到张量的不同分区并计算出不同的缩放因子。SwiGLU MoE层增加了另一个关键契约。SGLang和Megatron通常将门控(gate)和上投影(up projections)融合到一个GEMM中,因此这两个张量在转换和实时权重更新期间必须共享相同的FP32缩放因子,即使Hugging Face检查点将它们分开存储。Miles通过在NVFP4导出路径中将gate/up对一起量化来强制执行这一点。研究人员在整个软件栈中实现并向上游提交了单token NVFP4配方,包括TransformerEngine训练配方、cuDNN前端训练内核、FlashInfer rollout内核,以及SGLang集成和权重更新修复。

高精度与反量化反向模式:在高精度反向NVFP4变体中,前向传递和rollout对MoE专家使用NVFP4,而反向GEMM使用原始的BF16操作数。

具有高精度反向的NVFP4
具有高精度反向的NVFP4

反量化反向模式:反量化反向(Dequantized backward)是第二种可选模式。反向GEMM仍然在BF16中运行,但不使用原始的BF16值,而是消耗前向产生的精确低精度操作数的BF16反量化值。

具有反量化反向的NVFP4
具有反量化反向的NVFP4

反向模式的适用性与扩展:这两种模式都避免了低精度反向GEMM,因此这些配置不使用原始NVFP4预训练配方中的RHT或随机舍入。它们以牺牲反向吞吐量为代价换取更高精度的计算,但RL通常受限于rollout,长上下文注意力和通信进一步降低了端到端的影响。相同的反向模式选择也适用于MXFP8:

具有高精度反向的MXFP8
具有高精度反向的MXFP8
具有反量化反向的MXFP8
具有反量化反向的MXFP8

反向接口的实现:研究人员实现并向上游提交了 NVTE_BACKWARD_OVERRIDE 作为可重用的TransformerEngine接口,用于选择高精度或反量化反向操作数,然后通过Miles配方配置公开了这两种模式。伴随的humans&博客文章涵盖了算法动机和额外的反向模式消融实验。

反向成本与内存分析:反量化反向在训练端增加了一个反量化步骤。研究人员与NVIDIA合作降低了其在TransformerEngine中的开销。高精度和反量化反向还可以降低相对于TransformerEngine默认低精度反向路径的峰值内存。这两种模式都不需要生成和保留MXFP8训练部分中描述的低精度反向GEMM所使用的第二个按列量化的副本。研究人员在验证TransformerEngine反向模式实现时测量了以下内存数据。alloc列报告分配的内存,resrv列报告保留的内存,所有值单位均为MB。

MXFP8线性层内存,dtype=torch.bfloat16, input_shape=(2048, 2048), out_features=8192:

模式 fwd_alloc bwd_alloc e2e_alloc fwd_resrv bwd_resrv e2e_resrv delta_fwd delta_bwd delta_e2e
默认 73.75 73.02 94.14 474.00 474.00 474.00 +0.00 (+0.00%) +0.00 (+0.00%) +0.00 (+0.00%)
高精度 53.12 40.02 53.12 474.00 474.00 474.00 -20.62 (-27.97%) -33.00 (-45.20%) -41.02 (-43.57%)
反量化 53.25 80.02 84.64 474.00 474.00 474.00 -20.50 (-27.80%) +7.00 (+9.59%) -9.50 (-10.09%)

NVFP4线性层内存,dtype=torch.bfloat16, input_shape=(2048, 2048), out_features=8192:

模式 fwd_alloc bwd_alloc e2e_alloc fwd_resrv bwd_resrv e2e_resrv delta_fwd delta_bwd delta_e2e
默认 55.75 146.02 150.27 478.00 478.00 478.00 +0.00 (+0.00%) +0.00 (+0.00%) +0.00 (+0.00%)
高精度 44.50 40.02 44.50 478.00 478.00 478.00 -11.25 (-20.18%) -106.00 (-72.60%) -105.77 (-70.39%)
反量化 44.50 80.02 75.27 478.00 478.00 478.00 -11.25 (-20.18%) -66.00 (-45.20%) -75.00 (-49.91%)

位级精确的量化器契约:在RL中,量化不匹配会在权重更新中累积。如果训练端和rollout端对张量的量化方式不同,用于采样和学习的策略就不再是同一个低精度模型。因此,RL栈的两端之间需要一个明确的契约。研究人员将FlashInfer和TransformerEngine量化器对齐到相同的MXFP8和NVFP4位级契约。FlashInfer单元测试针对随机数据、量化边界数据、全零张量和最大值张量(对应于 init_data = ["random", "boundary", "zeros", "maxes"]),检查与TransformerEngine风格参考的精确字节级一致性。研究人员在FlashInfer中实现并向上游提交了这种量化器对齐。

服务与RL在量化上的实际区别:服务和RL之间存在一个实际区别。对于仅服务的工作负载,FlashInfer可能会在FP4量化路径的部分区域使用快速数学(fast math)以提高性能。这是一个合理的服务默认设置,但RL权重更新受益于与训练端量化器的精确一致。对于此配方,研究人员设置了 FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1。任何涉及rollout权重的后端都应该要么精确实现这个量化契约,要么将近似行为作为可选加入(opt-in)。

细粒度精度控制的实现:在实践中,单一的全局精度开关不足以满足低精度RL的需求。一些张量应保持在BF16,但选择它们只是问题的一部分:必须在Hugging Face检查点转换、Megatron训练、SGLang rollout和实时权重导出中强制执行相同的决定。研究人员在Miles中通过基于计数和基于名称的BF16异常处理跨越检查点转换、训练、rollout和实时导出实现了这种张量级的精度控制。研究人员还实现了SGLang对生成的混合精度检查点的支持。具体来说,转换使用 --num-layers-at-start-in-bf16 和 --num-layers-at-end-in-bf16;Megatron训练将这些计数与 --first-last-layers-bf16 结合;而SGLang服务于生成的混合精度检查点。

层精度选择的策略:正如NVIDIA NVFP4预训练论文中所推荐的,研究人员将一小部分最终层保持在较高精度。在实验中,将最后15%的层保持在BF16有意义地减少了训练-推理的不匹配并稳定了梯度。

将最终层保持在BF16的效果
将最终层保持在BF16的效果

早期层精度的影响:在实验中,为早期层开启BF16并没有带来有意义的训练-推理不匹配的减少:

首层BF16与末层BF16的对比
首层BF16与末层BF16的对比

共享专家的高精度保留:将共享专家保持在高精度也能减少训练-推理不匹配,且几乎没有性能或内存开销。路由专家被稀疏地选择,它们的输出通过高精度加权归约传递。共享专家始终处于活动状态,因此它们的精度误差会影响通过该块的每个token。

共享专家高精度消融实验
共享专家高精度消融实验

案例研究:MXFP8 MLA投影张量:对于MLA模型,kv_b_proj 是一个重要的MXFP8案例。吸收(Absorbed)和非吸收的MLA模式可能使用不同的收缩轴,而MXFP8使用一维的微缩放块。改变收缩轴可能会改变哪些元素共享一个缩放因子。同样的担忧也适用于其他一维格式,包括NVFP4。原始DeepSeek-V3 FP8配方在权重端没有这个确切的问题,因为它使用128x128的权重缩放块而不是一维块。研究人员将这些投影张量保持在BF16,以避免隐藏的重新量化并保持训练-推理的一致性。

configs:
  bf16:
    transformer_engine_config_type: "TEQuantizationParams"
    training_recipe: {}
matchers:
  mla_kv_up_proj_bf16:
    type: "glob"
    enabled: true
    pattern: "*.self_attention.linear_kv_up_proj"
    config: "bf16"
  absorbed_k_up_proj_bf16:
    type: "glob"
    enabled: true
    pattern: "*.self_attention.linear_k_up_proj"
    config: "bf16"
  absorbed_v_up_proj_bf16:
    type: "glob"
    enabled: true
    pattern: "*.self_attention.linear_v_up_proj"
    config: "bf16"

MLA匹配名称的参数配置:匹配的基于名称的转换时和rollout时参数为:

--extra-high-precision-layers-hf .kv_b_proj.
--extra-high-precision-layers-megatron .linear_kv_up_proj .linear_k_up_proj .linear_v_up_proj

MLA异常的内存开销:对于DeepSeek-V3风格的MLA,这种BF16异常很小。一个形状为 32768 x 512 的 kv_b_proj 张量,分布在61层中以BF16存储,大约占用1.90625 GB。

实验环境

  • 硬件配置:8张 B200 GPU。硬件划分为4张GPU用于rollout,4张GPU用于训练。
  • 软件配置:Miles框架,SGLang(rollout使用BF16 KV Cache,低精度MoE路径使用FlashInfer TRTLLM路由MoE后端),Megatron,TransformerEngine。
  • 数据集与工作负载:在 dapo-math-17k 数据集上进行GRPO风格的同步强化学习训练。每个prompt采样8个rollout样本,最大响应长度为8192个token。启用了KL散度路径用于诊断,但其系数设为0.0,因此KL不作为优化惩罚项。
  • 模型与参数:Qwen3-30B-A3B模型。启用MoE rollout路由重放,最后15%的层保持为BF16,低精度权重使用0权重衰减以维持稳定性。

实验结果

本次实验为配方消融测试(Recipe Ablation),比较了六种配置:

  1. BF16训练 + BF16 rollout(基线)。
  2. 端到端MXFP8训练 + MXFP8 rollout。
  3. MXFP8 rollout和前向 + 高精度反向。
  4. MXFP8 rollout和前向 + 反量化反向。
  5. 单token NVFP4 MoE rollout和前向 + 高精度反向。
  6. 单token NVFP4 MoE rollout和前向 + 反量化反向。

  7. 训练-推理不匹配(Train-inference mismatch):正如预期,两种低精度格式均表现出比BF16更高的训练-推理不匹配度。在每种格式内,两种反向模式表现相似。NVFP4的初始参考KL高于BF16或MXFP8(该诊断指标是针对Megatron BF16参考模型计算的,包含了低精度策略与参考模型之间的差异以及RL期间策略的演变)。
    跨反向模式的训练-rollout对数概率差异
    跨反向模式的KL损失对比

  8. 奖励(Reward):尽管诊断指标显示不匹配度较高,但所有五种低精度配置的奖励曲线都紧密跟踪BF16的基线奖励曲线。这表明Blackwell原生的低精度方案能够保留观察到的学习曲线。
    跨反向模式的原始奖励对比

  9. Rollout性能(Performance):与BF16相比,MXFP8和NVFP4都显著减少了rollout时间。对于NVFP4 rollout,FlashInfer直接在激活量化内核路径内部在线计算单token的FP32缩放因子,因此报告的rollout性能已包含了在线缩放计算的成本。
    跨反向模式的Rollout时间对比

  10. 训练性能(Performance):在训练端,MXFP8变体比BF16更快。但在本次测量的实现中,NVFP4反向覆盖变体较慢。这是由于当前TransformerEngine路径将单token FP32缩放作为单独的PyTorch操作应用,而不是使用将缩放融合到内核epilogue中的原生单token NVFP4 GEMM路径(融合的cuDNN前端内核已实现并提交,但TE集成待定)。反量化反向也增加了额外的反量化步骤。由于此RL工作负载重度依赖rollout,因此即使在训练路径完全加速之前,端到端的rollout加速仍然具有重要意义。
    跨反向模式的训练时间对比

补充细节

在线NVFP4模型服务:相同的单token NVFP4配方可用于服务FP8/BF16模型检查点,通过在线训练后量化(PTQ),在不产生校准伪影的情况下提供良好的模型质量,同时保持计算性能与传统NVFP4服务相当。在线单token激活缩放将激活异常值局部化到每个token,消除了从校准传递中携带静态激活FP32缩放因子的需要。在FlashInfer中,该FP32缩放计算被融合到激活量化内核路径中,因此在线缩放不需要单独的预处理内核。通过结合4/6(four-over-six)技术和零暴露解码延迟,可以进一步提高训练后量化(PTQ)的模型质量。研究人员在SGLang中实现并开源了这条在线NVFP4服务路径。在SGLang中使用此功能非常简单,只需在任何现有的BF16/FP8 MoE模型服务命令中添加 --quantization nvfp4_online 即可。

在Miles中尝试NVFP4配方:以下环境变量设置重现了单token NVFP4高精度反向的设置:

NVTE_NVFP4_ROW_SCALED_ACTIVATION=1
NVTE_BACKWARD_OVERRIDE=high_precision
NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1
NVTE_NVFP4_DISABLE_RHT=1
NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1
TRTLLM_DISABLE_FP4_QUANT_FAST_MATH=1
FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1
SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1

选择反量化反向变体:设置 NVTE_BACKWARD_OVERRIDE=dequantized 可以选择反量化反向变体,而无需更改配方的其余部分。

启动脚本配置:对于Miles启动脚本,该配方将这些环境变量与 --fp4-format e2m1、--fp4-recipe nvfp4、在检查点转换期间使用的相同的BF16首/尾层控制,以及以下TransformerEngine精度配置配对:

configs:
    nvfp4:
        transformer_engine_config_type: "TEQuantizationParams"
        training_recipe:
            fp4_quantization_recipe: "nvfp4"
    bf16:
        transformer_engine_config_type: "TEQuantizationParams"
        training_recipe: {}
matchers:
    routed_experts_fc1_nvfp4:
        type: "glob"
        enabled: true
        pattern: "*.mlp.experts.linear_fc1"
        config: "nvfp4"
    routed_experts_fc2_nvfp4:
        type: "glob"
        enabled: true
        pattern: "*.mlp.experts.linear_fc2"
        config: "nvfp4"
    default_bf16:
        type: "glob"
        enabled: true
        pattern: "*"
        config: "bf16"

结论

本文成功在Miles框架中实现了面向Blackwell架构的端到端MXFP8和基于MoE专家的单token NVFP4强化学习配方。通过精细的精度控制、一致的单token激活缩放机制以及位级精确的量化器契约,该方案有效控制了低精度RL中容易出现的训练-推理不匹配问题。实验表明,在保证模型学习能力(奖励曲线与BF16高度重合)的前提下,MXFP8和NVFP4均能显著提升rollout效率。

未来工作展望:

  1. 移除额外的BF16权重副本:尽管rollout和训练执行相同的低精度配方,Megatron目前仍保存一个额外的BF16权重副本,这增加了内存消耗。Megatron-Bridge正在推进稳健的低精度参数收集(parameter gather)工作,目前NVFP4的 --fp4-param-gather 尚未支持本配方使用的1D 1x16权重布局。
  2. 解决偶尔的梯度尖峰问题:高精度反向NVFP4变体偶尔会出现梯度尖峰。
    偶尔的NVFP4梯度尖峰
    反量化反向模式能减少最大的尖峰,但无法完全消除。未来需要探索更高级的技术,如4/6技术和符合链式法则的反向选择。

  3. 重构权重更新接口:低延迟的FlashInfer后端通常需要填充(padding)、重排(swizzling)、乱序(shuffling)和特定于后端的权重布局。这些转换对服务很有利,但由于训练端通常拥有不同的规范张量布局,这使得实时RL权重更新和RDMA变得复杂。Miles和SGLang未来的工作旨在保留高性能服务布局的同时,使每个权重转换变得明确、可验证,并减少对后端私有细节的依赖。