Miles 在 Blackwell 架构上实现端到端 MXFP8 与逐 token NVFP4 强化学习方案

LMSYS:Blog(Chatbot Arena 团队)·2026-07-29 00:00·48天前·Ziang Li, humans& and Miles Team
AI 导读

Miles 团队在 Blackwell 架构上实现了两种原生低精度强化学习方案:端到端 MXFP8 和 MoE 专家权重的逐 token NVFP4。在 8x B200 上对 Qwen3-30B-A3B 的消融实验中,BF16 与所有五种低精度配置的原始奖励曲线高度重合,且 MXFP8 和 NVFP4 减少了推理时间。

LMSYS:Blog(Chatbot Arena 团队)
精选
69AI 编辑部评分,满分 100

Miles 在 Blackwell 架构上实现端到端 MXFP8 与逐 token NVFP4 强化学习方案

2026-07-29 00:00· 48天前· Ziang Li, humans& and Miles Team
AI 导读

Miles 团队在 Blackwell 架构上实现了两种原生低精度强化学习方案:端到端 MXFP8 和 MoE 专家权重的逐 token NVFP4。在 8x B200 上对 Qwen3-30B-A3B 的消融实验中,BF16 与所有五种低精度配置的原始奖励曲线高度重合,且 MXFP8 和 NVFP4 减少了推理时间。

推荐理由

虽然只针对Blackwell,但这份报告把MXFP8和NVFP4在RL里的实现细节讲清楚了,从量化合约到每token缩放都给了可复现配置,做大规模RL训练的能直接上手。

正文 · AI 翻译

简而言之:我们在 Miles 中实现了两种 Blackwell 原生 RL 方案:端到端 MXFP8,以及针对 MoE 专家的逐 token NVFP4。两者均通过细粒度精度控制,在检查点转换、Megatron 训练、SGLang rollout 以及实时权重更新中得到支持。MXFP8 覆盖 rollout、前向传播、权重梯度 GEMM 和数据梯度 GEMM。NVFP4 在其 MoE 专家路径中使用在线逐 token 激活缩放,且两种格式都支持高精度或反量化反向模式。在 8x B200 上进行的 Qwen3-30B-A3B 方案消融实验中,BF16 与全部五种低精度配置的原始奖励曲线高度重叠,而 MXFP8 和 NVFP4 缩短了 rollout 时间。

引言

在低精度 RL 中,rollout、训练、检查点转换和实时权重更新必须遵循同一套精度约定,否则采样器与训练器的策略将会发散。Miles 和 SGLang RL 生态已经纳入了低精度方案:LMSYS 的 FP8 RL 文章表明,相对于 BF16 训练搭配 FP8 rollout,在训练和采样中统一使用 FP8 可以减少不匹配;INT4 QAT 文章表明,训练期间的伪量化加上 W4A16 rollout 可以让 INT4 变得实用。我们将这项工作扩展到 Blackwell 原生格式,在 Miles 中实现了 MXFP8 和 NVFP4 方案,并在 SGLang、TransformerEngine、FlashInfer、Megatron 和 cuDNN frontend 中贡献了配套组件。公开路线图记录在 miles#615

主要贡献如下:

  1. 端到端 MXFP8 RL,其中 rollout、前向传播、权重梯度 GEMM 和数据梯度 GEMM 全部使用 MXFP8。
  2. 针对 MoE 专家权重,采用逐 token 的 NVFP4 强化学习,并使用在线逐 token 激活缩放。
  3. 细粒度精度控制,因此诸如最终层等选定张量可以始终保持 BF16。
  4. TransformerEngine 与 FlashInfer 之间采用位精确的量化器契约,因此权重更新不会引入本可避免的训练-推理不匹配。

为什么需要 Blackwell 原生方案?

此前的低精度方案并非围绕 MXFP8 或 NVFP4 设计。现有的 Miles 路径遵循 DeepSeek-V3 风格的块缩放 FP8 方案:权重使用 128x128 块缩放,激活使用 1x128 tile 缩放,并且每个 tile 或块的缩放因子都是在线计算的。这是一套强大的 Hopper 时代方案,但在 Blackwell 上,其 FP32 缩放因子仍然是在 Tensor Core 路径外围以软件方式应用的,而非通过原生微缩放硬件实现。

INT4 QAT 解决的是另一个问题。训练使用伪量化来让模型适应 INT4 权重,而 rollout 使用 W4A16。尽管内存效率高,但计算路径实际上仍然使用 BF16 激活和反量化后的 INT4 权重。下表将 NVIDIA 的 HGX 平台 稠密 Tensor Core 规格归一化为每 GPU 吞吐量:B200 和 B300 来自 8-GPU HGX 系统,Rubin 来自 HGX Rubin NVL8 表格。

GPUBF16 稠密 Tensor CoreFP8 稠密 Tensor CoreFP4 稠密 Tensor Core
B2002.25 PFLOPS4.5 PFLOPS9 PFLOPS
B3002.25 PFLOPS4.5 PFLOPS13.5 PFLOPS
Rubin GPU(NVL8)4 PFLOPS17.5 PFLOPS35 PFLOPS

对于强化学习系统,精度契约涵盖:

  • SGLang rollout。
  • Megatron 与 TransformerEngine 训练。
  • Hugging Face checkpoint 转换。
  • 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 块。

NVFP4 two-level scaling with FP32 tensor scale and E4M3 block scales

FP32 缩放因子可以在不同的张量范围上进行选择。这一选择属于方案层面的决策,而非格式本身的属性,并且它对 RL 尤为重要。

方案一:端到端 MXFP8 强化学习

MXFP8 方案是此前端到端 FP8 工作最直接的 Blackwell 原生扩展。Rollout、前向传播、权重梯度 GEMM 以及数据梯度 GEMM 全部使用 MXFP8,而部分选定的张量则通过下文所述的精度控制规则保持为 BF16。

End-to-end MXFP8 RL recipe

训练

TransformerEngine 和 Megatron 将 MXFP8 实现为一条经过性能优化的一等 Blackwell 训练路径,其中包括 deepseek-v3-gb200-optimization.md 中所述的 GB200 DeepSeek-V3 优化。在我们的 Miles 集成中,我们将这条路径作为端到端 MXFP8 强化学习训练侧的基础。

DeepSeek-V3 FP8 方案 的一个区别在于反向传播激活值的表示方式。DeepSeek-V3 将前向激活值存储为 1x128 FP8 分块,并在反向 GEMM 之前将其转换为反向方向。这种做法存储的 FP8 数据更少,但在反向 GEMM 之前引入了反量化加重新量化。TransformerEngine 的 MXFP8 文档指出,行向 1x32 分块和列向 32x1 分块在数值上是不同的,必须从全精度数据独立量化。因此,TransformerEngine 在量化过程中会同时物化行向和列向的量化副本。这会占用更多内存,但避免了额外的重新量化步骤,并减少了反向路径中额外的量化误差。

这是强化学习中典型的系统权衡。我们采用 TransformerEngine 这条路径,以保持单一的端到端 MXFP8 契约,而不引入另一个不匹配来源。

推出

在推出侧,SGLang 使用了来自 FlashInfer 和 Triton 的 Blackwell MXFP8 kernel。我们在 FlashInfer 和 SGLang 中实现并向上游贡献了推出路径(flashinfer#2581sglang#17449sglang#19537sglang#21576 以及 sglang#28459)。

几乎所有主要的 GEMM 都可以量化为 MXFP8,包括注意力投影和 MoE 专家。主要例外是显式控制的高精度层,例如下文描述的 BF16 MLA 投影。

方案 2:Per-Token NVFP4 RL

NVFP4 比 MXFP8 更为激进,因此我们有选择性地应用它。我们对 MoE 专家进行量化,因为它们主导了模型规模和推出内存流量,而模型的其余部分除非显式配置,否则保持 BF16。

例如,DeepSeek-V3 总参数量约为 671B。其 MoE 专家占比为:

(61 - 3) * (256 + 1) * 3 * 7168 * 2048 / 1e9 = 656.5B parameters

这约占模型的 97.8%。因此,针对 MoE 专家进行优化可以捕获大部分内存收益,而无需强制每一层都采用最激进的精度格式。

为什么不直接使用 NVFP4 预训练方案?

最初的 NVFP4 预训练配方是为大规模预训练设计的,其目标是在使用 FP4 GEMM 的同时,在大量 token 上保持一个粗略的优化方向。它将 FP4 线性层 GEMM 与若干稳定机制相结合:选定层保持较高精度,权重缩放在前向与反向传播中保持一致,训练路径使用随机舍入(SR)和随机 Hadamard 变换(RHT)。在论文中,SR 被应用于梯度,以减少量化偏差并产生无偏的量化梯度,而 RHT 则用于分散块级别的大幅值离群点,尤其是针对权重梯度 GEMM 的输入。

Original NVFP4 pretraining recipe

这是一个不错的起点,但 RL 有着不同的失效面:

设置数值情况结论
预训练梯度信号稳定,权重更新幅度可观,模型权重具有自适应性,而激活值和梯度对精度敏感且具有高动态范围。保持粗略的优化方向和收敛性。
RL梯度噪声大,奖励方差高,而有效的更新幅度小且精细。让量化噪声低于真实的更新信号;否则它可能覆盖脆弱的能力,导致性能崩溃。

NVFP4 强化学习方案并未纳入预训练方案的每一个部分。我们针对的是 MoE 专家权重量化、逐 token 激活缩放、一致的精度控制,以及可选原始或反量化操作数的 BF16 反向 GEMM。

逐 token 激活缩放

两级 NVFP4 层级结构很强大,但 FP32 激活缩放的作用范围必须谨慎选择。正如 Cursor Composer 2 技术报告中所讨论的,逐张量 NVFP4 缩放会使训练依赖于批次,而 token 间共享缩放则可能将未来 token 的信息泄漏到过去 token 的表示中。如果一个 token 与其他 token 共享其缩放,那么它的量化表示就取决于批次构成。这对强化学习尤其成问题,因为其 rollout 调度和序列长度都会变化。

因此,我们的方案在线为每个 token 计算一个 FP32 激活缩放。这将激活离群值局部化到单个 token,消除了静态激活校准的伪影,并让 SGLang rollout 和 Megatron 训练使用相同的激活缩放作用范围。

在 rollout 一侧,逐 token 的 FP32 缩放计算被融合进 FlashInfer 的激活量化内核路径:同一个既输出打包 FP4 激活和 E4M3 块缩放的调用,也会返回逐 token 的 FP32 缩放。因此,逐 token 激活缩放不需要单独的校准缩放流程。

训练与推理一致性还要求并行方式匹配。如果 FP32 scale 是在专家张量并行划分内按 token 计算的,那么 SGLang 和 Megatron 应使用相同的 ETP 大小。否则,双方可能看到张量的不同划分,从而计算出不同的 scale。

SwiGLU MoE 层还引入了另一个关键约定。SGLang 和 Megatron 通常将 gate 和 up 投影融合为一个 GEMM,因此即使在 Hugging Face checkpoint 中它们分开存储,两个张量在转换和实时权重更新期间也必须共享相同的 FP32 scale。Miles 通过在 NVFP4 导出路径中将 gate/up 成对量化来强制实现这一点。

我们在整个技术栈中实现并向上游贡献了 per-token NVFP4 方案:

高精度与反量化反向传播

在高精度反向传播的 NVFP4 变体中,前向传播和 rollout 对 MoE 专家使用 NVFP4,而反向 GEMM 使用原始 BF16 操作数。

NVFP4 with high-precision backward

反量化反向传播是第二种可选模式。反向 GEMM 仍以 BF16 运行,但使用的是前向过程中生成的精确低精度操作数的 BF16 反量化结果,而非原始的 BF16 值。

NVFP4 with dequantized backward

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

相同的反向模式选择同样适用于 MXFP8:

MXFP8 with high-precision backward

MXFP8 with dequantized backward

我们实现并向上游贡献了 NVTE_BACKWARD_OVERRIDE,作为可复用的 TransformerEngine 接口,用于选择高精度或反量化的反向操作数(TransformerEngine#2644),随后通过 Miles 配方配置暴露了这两种模式。配套的 humans& 文章涵盖了算法动机以及额外的反向模式消融实验。

反向成本与内存

反量化反向传播增加了一个训练侧的反量化步骤。我们与 NVIDIA 合作,在 TransformerEngine#2865 中降低了其开销。

相对于 TransformerEngine 默认的低精度反向路径,高精度和反量化反向传播还可以降低峰值内存。这两种模式都不需要生成并保留上文 MXFP8 训练部分所述低精度反向 GEMM 所使用的第二份按列量化的副本。

我们在验证 TransformerEngine#2644 中的 TransformerEngine 反向模式实现时,测量了下面的内存数据。alloc 列报告已分配内存,resrv 列报告已预留内存,所有数值均以 MB 为单位。

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

模式前向分配反向分配端到端分配前向预留反向预留端到端预留前向差值反向差值端到端差值
默认73.7573.0294.14474.00474.00474.00+0.00 (+0.00%)+0.00 (+0.00%)+0.00 (+0.00%)
高精度53.1240.0253.12474.00474.00474.00-20.62 (-27.97%)-33.00 (-45.20%)-41.02 (-43.57%)
反量化53.2580.0284.64474.00474.00474.00-20.50 (-27.80%)+7.00 (+9.59%)-9.50 (-10.09%)

NVFP4 线性内存,dtype=torch.bfloat16input_shape=(2048, 2048)out_features=8192

模式前向分配反向分配端到端分配前向预留反向预留端到端预留前向增量反向增量端到端增量
默认55.75146.02150.27478.00478.00478.00+0.00 (+0.00%)+0.00 (+0.00%)+0.00 (+0.00%)
高精度44.5040.0244.50478.00478.00478.00-11.25 (-20.18%)-106.00 (-72.60%)-105.77 (-70.39%)
反量化44.5080.0275.27478.00478.00478.00-11.25 (-20.18%)-66.00 (-45.20%)-75.00 (-49.91%)

位精确量化器契约

在强化学习中,量化不一致会在权重更新过程中不断累积。如果训练侧和 rollout 侧对同一个张量采用不同的量化方式,那么用于采样和学习的策略就不再是同一个低精度模型了。因此,我们需要在 RL 栈的两侧之间建立一份明确的契约。

我们将 FlashInfer 和 TransformerEngine 的量化器对齐到同一套 MXFP8 和 NVFP4 位级契约。我们的 FlashInfer 单元测试会在随机数据、量化边界数据、全零张量和最大值张量上,检查与 TransformerEngine 风格参考实现是否达到精确的字节级一致,对应 init_data = ["random", "boundary", "zeros", "maxes"]。我们在 flashinfer#3387 中实现并向上游提交了这次量化器对齐。

在推理服务和强化学习之间有一个实际的区别。对于纯推理服务工作负载,FlashInfer 可能会在 FP4 量化路径的某些部分使用快速数学运算以提升性能。作为推理服务的默认设置,这是合理的,但 RL 的权重更新则受益于与训练侧量化器的精确一致。针对这套方案,我们设置如下:

FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1

每一个涉及 rollout 权重的后端,要么精确实现这一量化契约,要么将近似行为设为可选启用。

细粒度精度控制

在实践中,对于低精度 RL,单一的全局精度开关是不够的。有些张量应保持 BF16,但选择哪些张量只是问题的一部分:同样的决策必须在 Hugging Face checkpoint 转换、Megatron 训练、SGLang rollout 和实时权重导出中一致地强制执行。

我们在 Miles 中通过基于计数和基于名称的 BF16 例外,在 checkpoint 转换、训练、rollout 和实时导出中实现了这种张量级精度控制(miles#614miles#1054miles#1261)。我们还为由此产生的混合精度 checkpoint 实现了 SGLang 支持(sglang#18742sglang#20214)。具体而言,转换使用 --num-layers-at-start-in-bf16--num-layers-at-end-in-bf16;Megatron 训练将这些计数与 --first-last-layers-bf16 结合;SGLang 则服务由此产生的混合精度 checkpoint。

层级精度选择

按照 NVIDIA NVFP4 预训练论文中的建议,我们将一小部分最后几层保持在较高精度。在我们的实验中,将最后 15% 的层保持为 BF16 能显著减少训练-推理不匹配并稳定梯度。

Effect of keeping final layers in BF16

在我们的实验中,对早期层启用 BF16 并不会带来有意义的训练-推理不匹配减少:

First-layer BF16 versus last-layer BF16

将共享专家保持在高精度同样能减少训练-推理不匹配,且性能或内存开销很小。路由专家是稀疏选择的,其输出会经过高精度加权归约。共享专家始终处于激活状态,因此其精度误差会影响通过该模块的每一个 token。

Shared expert high-precision ablation

案例研究:MXFP8 MLA

对于 MLA 模型,kv_b_proj 是一个重要的 MXFP8 案例。吸收式和非吸收式 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"

基于名称匹配的转换时和发布时参数为:

--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

对于DeepSeek-V3风格的 MLA,这个 BF16 例外很小。一个kv_b_proj形状为32768 x 512的张量,以 BF16 存储于 61 层中,占用约 1.90625 GB。

结果:Qwen3-30B-A3B 配方消融实验(8x B200)

为进行一致的比较,所有实验均使用同步 Qwen3-30B-A3B RL,采用默认 Miles 配置,运行于 8x B200 上。固定工作负载在 dapo-math-17k 上使用 GRPO 风格训练,每个提示词 8 个 rollout 样本,最大响应长度为 8192 tokens。这仅是配方消融实验设置,并非完全调优的训练或服务基准测试。KL 路径为诊断目的而启用,但其系数为 0.0,因此在本消融实验中 KL 并非优化惩罚项。

硬件分配为 4 块 GPU 用于 rollout,4 块 GPU 用于训练。

对于所有低精度配方:

  • 启用 MoE rollout 路由重放。
  • 最后 15% 的层保持为 BF16。
  • 低精度权重使用 0 权重衰减以保证稳定性。
  • SGLang rollout 使用 BF16 KV cache,低精度 MoE 路径使用 FlashInfer TRTLLM 路由 MoE 后端。

我们比较六种配置:

  1. BF16 训练 + BF16 rollout。
  2. 端到端 MXFP8 训练 + MXFP8 rollout。
  3. MXFP8 rollout 和前向传播,配合高精度反向传播。
  4. MXFP8 rollout 和前向传播,配合反量化反向传播。
  5. 逐 token NVFP4 MoE rollout 与前向传播,配合高精度反向传播。
  6. 逐 token NVFP4 MoE rollout 与前向传播,配合反量化反向传播。

训练-推理不匹配

正如预期,两种低精度格式都表现出比 BF16 更高的训练-推理不匹配,而在每种格式内部,两种反向传播选择的表现相似。在此消融实验中,这些数值仍处于合理范围内。

Train-rollout logprob difference across backward modes KL loss comparison across backward modes

NVFP4 起始时的参考 KL 高于 BF16 或 MXFP8。Miles 是针对 Megatron BF16 参考模型计算这一诊断指标的,因此该指标除了包含 RL 过程中的策略演化外,还包含了每种低精度策略与 BF16 参考模型之间的差异。不应将其单独解读为优化惩罚。

奖励

尽管诊断不匹配程度更高,但五条低精度奖励曲线都与 BF16 奖励曲线紧密吻合。

Raw reward comparison across backward modes

这是配方消融实验的关键结果:在这个 Qwen3-30B-A3B B200 配置中,Blackwell 原生低精度在保持所观测到的学习曲线的同时,提升了 rollout 效率。

性能

与 BF16 相比,MXFP8 和 NVFP4 都缩短了 rollout 时间:

Rollout time comparison across backward modes

对于 NVFP4 rollout,FlashInfer 直接在激活量化 kernel 路径内部计算在线 per-token FP32 scale,而不是将其作为单独的预处理步骤。因此,所报告的 rollout 性能包含了在线 scale 计算的开销。

在训练侧,MXFP8 变体比 BF16 更快,而 NVFP4 backward-override 变体在此处所测量的实现中更慢:

Training time comparison across backward modes

训练侧的差距来自此消融实验中所使用的实现,而非 FP4 Tensor Core 的固有限制。我们的 TransformerEngine 路径将 per-token FP32 scaling 作为单独的 PyTorch 操作来应用(TransformerEngine#2931),而不是采用将 scaling 融合进 kernel epilogue 的原生 per-token NVFP4 GEMM 路径。我们已经实现并向上游提交了融合的 cuDNN 前端 kernel(cudnn-frontend#251);TransformerEngine 集成仍待完成。反量化 backward 增加了上述反量化步骤。由于此 RL 工作负载以 rollout 为主,即使在训练路径完全加速之前,rollout 的加速仍然意义重大。

除本消融实验之外,humans& 在生产环境中使用相同的配方系列和组件,用于大规模、长上下文、多智能体异步 RL 研究。

在线 NVFP4 模型服务

相同的 per-token NVFP4 配方可以通过在线训练后量化(PTQ)来服务 FP8/BF16 模型 checkpoint,在无需校准产物的情况下提供良好的模型质量,同时保持计算性能与传统 NVFP4 服务相当。

在线逐 token 激活缩放将激活离群值定位到每个 token,并免除了从校准过程携带静态激活 FP32 缩放的需求。在 FlashInfer 中,该 FP32 缩放计算被融合进激活量化 kernel 路径,因此在线缩放不需要单独的预处理 kernel。

通过将 four-over-six 与零暴露解码延迟相结合,训练后量化(PTQ)的模型质量可以进一步提升,详见配套的 humans& 文章。

我们在 sglang#26083 中实现并开源了这条 SGLang 在线 NVFP4 服务路径。在 SGLang 中使用该功能,只需在任意现有的 BF16/FP8 MoE 模型服务命令中添加 --quantization nvfp4_online 即可。

未来工作

去掉额外的 BF16 权重副本

尽管 rollout 和训练执行相同的低精度配方,Megatron 仍会额外保存一份 BF16 权重副本。这增加了内存消耗,并限制了低精度路径在实际中的内存收益。

Megatron 有 --fp8-param-gather--fp4-param-gather,但 Blackwell 原生路径仍在成熟过程中。Megatron-Bridge 跟踪 issue Megatron-Bridge#3801 反映了为实现稳健的低精度参数收集仍需持续开展的工作。NVFP4 --fp4-param-gather 尚不支持该配方所使用的 1D 1x16 权重布局。

偶发的梯度尖峰

高精度反向传播的 NVFP4 变体仍可能出现偶发的梯度尖峰:

Occasional NVFP4 gradient spike

在此消融实验中,反量化反向传播减少了最大的尖峰,但并未将其消除。更先进的技术,包括 4/6 以及链式法则一致的反向传播选择,在配套的 humans& 文章中讨论。

重构权重更新接口

低延迟的 FlashInfer 后端通常需要填充、swizzling、shuffling 以及后端特定的权重布局。这些变换对于服务部署而言很自然,但它们会使实时 RL 权重更新和 RDMA 变得复杂,因为训练侧通常拥有不同的规范张量布局。

Miles 和 SGLang 中的工作旨在保留高性能服务布局,同时使每一次权重变换都显式、可验证,并且更少依赖后端私有细节。

在 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"

致谢

本博文所述的配方设计及大部分实现工作由 humans& 的 Ziang Li 完成。

我们感谢以下合作者提供的工程支持、集成帮助和审阅:

  • SGLang 团队与 Miles 团队。
  • NVIDIA DevTech Compute 团队(Siyuan Fu、Yigong Qin、Zhongbo Zhu)、TransformerEngine 团队和 FlashInfer 团队。

我们还要感谢 Cursor 团队提出的 per-token NVFP4 激活缩放思路。

来源:LMSYS:Blog(Chatbot Arena 团队)· lmsys.org