Miles 在 Blackwell GPU 上实现低精度强化学习，降低训练成本

为什么需要 Blackwell 原生方案？

方案 1：端到端 MXFP8 强化学习

方案 2：逐 Token NVFP4 强化学习

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

逐 Token 激活值缩放

高精度和反量化反向传播

反向传播开销和内存

位级精确的量化器契约

细粒度精度控制

层精度选择

案例研究：MXFP8 MLA

结果：在 8× B200 上进行 Qwen3-30B-A3B 方案消融实验

训练与推理不匹配

移除额外的 BF16 权重副本

偶发的梯度尖峰

重构权重更新接口

在 Miles 中试用 NVFP4 方案

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

在低精度强化学习中，rollout、训练、检查点转换和实时权重更新必须遵循同一套精度契约，否则采样器策略与训练器策略将产生偏离。Miles 和 SGLang 强化学习生态系统已经集成了低精度方案：LMSYS 的 FP8 强化学习文章表明，相比使用 BF16 训练、FP8 rollout，在训练和采样中都使用 FP8 可以减少不匹配；INT4 QAT 文章则表明，在训练期间使用伪量化并采用 W4A16 rollout，可以让 INT4 变得切实可用。我们将这项工作扩展到了 Blackwell 原生格式：在 Miles 中实现 MXFP8 和 NVFP4 方案，并为 SGLang、TransformerEngine、FlashInfer、Megatron 和 cuDNN frontend 贡献配套组件。公开路线图记录在 miles#615 中。

主要贡献包括：

端到端 MXFP8 强化学习，其中 rollout、前向传播、权重梯度 GEMM 和数据梯度 GEMM 全部使用 MXFP8。

面向 MoE 专家权重的逐 Token NVFP4 强化学习，使用在线逐 Token 激活值缩放。

细粒度精度控制，使最终层等选定张量能够始终保持 BF16。

TransformerEngine 与 FlashInfer 之间采用位级精确的量化器契约，使权重更新不会引入本可避免的训练与推理不匹配。

为什么需要 Blackwell 原生方案？

以往的低精度方法并非围绕 MXFP8 或 NVFP4 设计。Miles 现有路径采用类似 DeepSeek-V3 的分块缩放 FP8 方案：权重使用 128×128 分块缩放，激活值使用 1×128 分块缩放，并为每个分块在线计算缩放因子。这是一种出色的 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 表格。

对于强化学习系统，精度契约涵盖：

Megatron 和 TransformerEngine 训练。

Hugging Face 检查点转换。

Megatron 到 Hugging Face 的实时权重导出。

细粒度高精度例外配置。

MXFP8 是一种微缩放 FP8 格式。TransformerEngine 的 MXFP8 文档将其描述为一种 Blackwell 原生的分块缩放方案：每连续 32 个 E4M3 值共享一个局部 E8M0 缩放因子，并且分块是一维的。

由于 E8M0 缩放因子表示 2 的幂，因此解码后的缩放因子通常会向上取整，以避免裁剪分块中的最大值。

NVFP4 是 Blackwell 的原生 FP4 格式。正如 NVIDIA 对 NVFP4 的介绍所述，它存储 FP4 E2M1 值，每个包含 16 个值的分块共享一个 FP8 E4M3 缩放因子。由于 E4M3 的分辨率比 UE8M0 更精细，其缩放因子通常会舍入到最接近的可表示值。标准 NVFP4 方案还会为更大范围的张量添加一个 FP32 缩放因子，从而形成两级层次结构：

一个粗粒度 FP32 缩放因子，用于将张量或 Token 映射到 NVFP4 的可表示范围内。

一个细粒度 E4M3 缩放因子，用于适配每个 1×16 分块。

FP32 缩放因子可以基于不同的张量范围进行选择。这种选择属于方案决策，而不是格式本身的属性；对于强化学习而言，这一点尤为重要。

方案 1：端到端 MXFP8 强化学习

MXFP8 方案是对早期端到端 FP8 工作最直接的 Blackwell 原生扩展。rollout、前向传播、权重梯度 GEMM 和数据梯度 GEMM 全部使用 MXFP8，而选定张量则通过下文介绍的精度控制规则保持 BF16。

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

它与 DeepSeek-V3 FP8 方案之间的一个区别，在于反向传播激活值的表示方式。DeepSeek-V3 将前向激活值存储在 1×128 FP8 分块中，并在执行反向 GEMM 之前将其转换为反向传播所需的方向。这种方式存储的 FP8 数据更少，但会在反向 GEMM 之前引入反量化和重新量化。TransformerEngine 的 MXFP8 文档指出，按行排列的 1×32 分块与按列排列的 32×1 分块在数值上并不相同，必须分别基于全精度数据独立量化。因此，TransformerEngine 会在量化期间同时物化按行和按列量化的两个副本。这会占用更多内存，但能够避免额外的重新量化步骤，并减少反向传播路径中的额外量化误差。

这是强化学习系统中典型的取舍。我们采用 TransformerEngine 路径来维持统一的端到端 MXFP8 契约，避免引入另一个不匹配来源。

在 rollout 侧，SGLang 使用来自 FlashInfer 和 Triton 的 Blackwell MXFP8 内核。我们在 FlashInfer 和 SGLang 中实现并向上游提交了 rollout 路径（flashinfer#2581、sglang#17449、sglang#19537、sglang#21576 和 sglang#28459）。

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

方案 2：逐 Token NVFP4 强化学习

NVFP4 比 MXFP8 更激进，因此我们会选择性地应用它。我们对 MoE 专家网络进行量化，因为它们在模型大小和 rollout 内存流量中占据主导地位；除非显式进行其他配置，否则模型其余部分仍保持 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 的输入离群值。

这是一个很好的起点，但强化学习面临不同的失效场景：

NVFP4 强化学习方案并未纳入预训练方案的全部组成部分。我们的目标是实现 MoE 专家权重量化、逐 Token 激活值缩放、一致的精度控制，以及使用可选择原始操作数或反量化操作数的 BF16 反向 GEMM。

逐 Token 激活值缩放

两级 NVFP4 层级很强大，但 FP32 激活缩放的范围必须小心选择。如 Cursor Composer 2 技术报告所述，per-tensor NVFP4 缩放会使训练依赖批次，而跨令牌缩放共享会导致未来令牌信息泄漏到过去令牌的表示中。如果一个令牌与其他令牌共享其缩放，其量化表示依赖于批次组成。这对于 RL 特别有问题，因为在 RL 中，采样调度和序列长度会变化。

因此，我们的方案在线计算每个令牌一个 FP32 激活缩放。这将激活离群值限制在一个令牌，移除了静态激活校准的伪影，并让 SGLang 采样和 Megatron 训练使用相同的激活缩放范围。

在采样端，per-token FP32 缩放计算被融合到 FlashInfer 的激活量化内核路径中：发出打包 FP4 激活和 E4M3 块缩放的同一个调用也返回 per-token FP32 缩放。因此，per-token 激活缩放不需要单独的校准缩放通道。

训练-推理一致性也需要匹配的并行性。如果 FP32 缩放在专家-张量并行分区内按令牌计算，SGLang 和 Megatron 应该使用相同的 ETP 大小。否则，每一端可能看到张量的不同分区并计算不同的缩放。

SwiGLU MoE 层添加了另一个关键合约。SGLang 和 Megatron 通常将门和 up 投影融合成一个 GEMM，所以两个张量必须在转换和实时权重更新期间共享相同的 FP32 缩放，即使 Hugging Face 检查点分别存储它们。Miles 通过在 NVFP4 导出路径中将 gate/up 对一起量化来强制执行此操作。

我们在整个堆栈中实现并上游了 per-token NVFP4 方案：

TransformerEngine 训练方案：TransformerEngine#2931

cuDNN frontend 训练内核：cudnn-frontend#251

FlashInfer 采样内核：flashinfer#3027

SGLang 集成和权重更新修复：sglang#22918、sglang#22204

高精度反向和去量化反向

在高精度反向 NVFP4 变体中，前向通道和采样对 MoE 专家使用 NVFP4，而反向 GEMM 使用原始 BF16 操作数。

去量化反向是第二种可选模式。反向 GEMM 仍以 BF16 运行，但消耗前向生成的精确低精度操作数的 BF16 去量化，而不是原始 BF16 值。

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

相同的反向模式选择也适用于 MXFP8：

我们实现并上游了 NVTE_BACKWARD_OVERRIDE 作为可重用的 TransformerEngine 接口，用于选择高精度或去量化反向操作数 (TransformerEngine#2644)，然后通过 Miles 方案配置暴露两种模式。相应的 humans& 帖子涵盖了算法动机和额外的反向模式消融。

去量化反向添加了一个训练端的去量化步骤。我们与 NVIDIA 合作在 TransformerEngine#2865 中减少其开销。

高精度和去量化反向也可以相对于 TransformerEngine 的默认低精度反向路径减少峰值内存。两种模式都不需要生成和保留低精度反向 GEMM 使用的第二个按列量化的副本，该副本在上面的 MXFP8 训练部分中描述。

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

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

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

比特精确量化器合约

在 RL 中，量化不匹配可能在权重更新中累积。如果训练和采样端对张量量化方式不同，用于采样和学习的策略将不再是相同的低精度模型。因此，我们需要在 RL 堆栈的两端之间有一个明确的合约。

我们将 FlashInfer 和 TransformerEngine 量化器对齐到相同的 MXFP8 和 NVFP4 比特级合约。我们的 FlashInfer 单元测试检查与 TransformerEngine 风格参考在随机数据、量化边界数据、全零张量和最大值张量上的精确字节级一致性，对应于 init_data = ["random", "boundary", "zeros", "maxes"]。我们在 flashinfer#3387 中实现并上游了此量化器对齐。

服务和 RL 之间有一个实际的区别。对于仅服务工作负载，FlashInfer 可能在 FP4 量化路径的某些部分使用快速数学以获得性能。这是一个合理的服务默认值，但 RL 权重更新受益于与训练端量化器的精确一致。对于此方案，我们设置：

FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1

每个触及采样权重的后端应该要么精确实现此量化合约，要么使近似行为可选。

细粒度精度控制

实际上，单一的全局精度开关对于低精度 RL 是不足够的。某些张量应该保持为 BF16，但选择它们只是问题的一部分：相同的决定必须在 Hugging Face 检查点转换、Megatron 训练、SGLang 采样和实时权重导出中强制执行。

我们在 Miles 中通过基于计数和基于名称的 BF16 例外实现了此张量级精度控制，横跨检查点转换、训练、采样和实时导出 (miles#614、miles#1054 和 miles#1261)。我们还为生成的混合精度检查点实现了 SGLang 支持 (sglang#18742 和 sglang#20214)。具体来说，转换使用 --num-layers-at-start-in-bf16 和 --num-layers-at-end-in-bf16；Megatron 训练将这些计数与 --first-last-layers-bf16 结合；SGLang 服务生成的混合精度检查点。

层精度选择

如 NVIDIA NVFP4 预训练论文所推荐，我们保留最后几层的一小部分为更高精度。在我们的实验中，将最后 15% 的层保持为 BF16 有意义地减少了训练-推理不匹配并稳定了梯度。

在我们的实验中，为早期层打开 BF16 不会导致有意义的训练-推理不匹配减少：

将共享专家保持为高精度也减少了训练-推理不匹配，性能或内存开销很少。路由专家被稀疏选择，其输出通过高精度加权归约。共享专家始终处于活动状态，所以其精度误差影响通过该块的每个令牌。

对于 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 例外占用的空间很小。一个形状为 32768 x 512 的 kv_b_proj 张量，以 BF16 格式存储并分布在 61 层中，大约占用 1.90625 GB。

结果：在 8x B200 上对 Qwen3-30B-A3B 配方进行消融实验

为确保比较一致，所有实验均在 8x B200 上采用默认的 Miles 配置，进行同步 Qwen3-30B-A3B 强化学习。固定工作负载在 dapo-math-17k 上使用 GRPO 风格训练，每个提示生成 8 个 rollout 样本，最大响应长度为 8192 个 token。这只是一个配方消融实验配置，并非经过完整调优的训练或服务基准测试。实验启用了 KL 路径用于诊断，但其系数为 0.0，因此在本次消融实验中，KL 不作为优化惩罚项。

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

对于所有低精度配方：

启用 MoE rollout 路由重放。

最后 15% 的层保持使用 BF16。

为保证稳定性，低精度权重的权重衰减设为 0。

在低精度 MoE 路径中，SGLang rollout 使用 BF16 KV cache，并采用 FlashInfer TRTLLM routed MoE 后端。

我们比较了六种配置：

BF16 训练 + BF16 rollout。

端到端 MXFP8 训练 + MXFP8 rollout。

MXFP8 rollout 和前向传播，搭配高精度反向传播。

MXFP8 rollout 和前向传播，搭配反量化反向传播。

逐 token NVFP4 MoE rollout 和前向传播，搭配高精度反向传播。

逐 token NVFP4 MoE rollout 和前向传播，搭配反量化反向传播。

训练与推理之间的不匹配

正如预期，与 BF16 相比，两种低精度格式都表现出更高的训练—推理不匹配；而在每种格式内部，两种反向传播方式的表现相近。在本次消融实验中，这些数值仍处于合理范围内。

NVFP4 的初始参考 KL 高于 BF16 或 MXFP8。Miles 使用 Megatron BF16 参考模型计算这一诊断指标，因此除了强化学习期间策略自身的演化之外，该指标还包含各低精度策略与 BF16 参考模型之间的差异。不应将其视为独立的优化惩罚项。

尽管诊断得到的不匹配程度更高，但五种低精度配置的奖励曲线都与 BF16 奖励曲线高度一致。

这是本次配方消融实验的关键结果：在这一 Qwen3-30B-A3B B200 配置中，Blackwell 原生低精度在提高 rollout 效率的同时，保留了观测到的学习曲线。

与 BF16 相比，MXFP8 和 NVFP4 均能缩短 rollout 时间：

对于 NVFP4 rollout，FlashInfer 直接在激活量化内核路径中计算在线逐 token FP32 scale，而不是将其作为独立的预处理步骤。因此，报告的 rollout 性能已包含在线 scale 计算的开销。

在训练侧，MXFP8 变体比 BF16 更快，而在此处测量的实现中，使用反向传播覆盖的 NVFP4 变体更慢：

训练侧的性能差距来自本次消融实验采用的实现，而不是 FP4 Tensor Core 的固有限制。我们的 TransformerEngine 路径将逐 token FP32 scaling 作为独立的 PyTorch 操作执行（TransformerEngine#2931），而非使用原生的逐 token NVFP4 GEMM 路径，将 scaling 融合到内核 epilogue 中。我们已经实现并向上游提交了融合的 cuDNN frontend 内核（cudnn-frontend#251）；TransformerEngine 集成工作仍在等待完成。反量化反向传播还会增加前文所述的反量化步骤。由于这一强化学习工作负载以 rollout 为主，即使训练路径尚未得到全面加速，rollout 的速度提升仍然很有意义。

除了本次消融实验之外，humans& 还在大规模、长上下文、多智能体异步强化学习研究的生产环境中使用同一配方系列及其组件。

移除额外的 BF16 权重副本

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

Megatron 提供了 --fp8-param-gather 和 --fp4-param-gather，但 Blackwell 原生路径仍在逐步成熟。Megatron-Bridge 跟踪 issue Megatron-Bridge#3801 反映了实现稳健低精度参数 gather 仍需推进的工作。NVFP4 的 --fp4-param-gather 尚不支持此配方使用的一维 1x16 权重布局。

偶发的梯度尖峰

采用高精度反向传播的 NVFP4 变体仍可能偶尔出现梯度尖峰：

在本次消融实验中，反量化反向传播降低了最大尖峰的幅度，但未能将其完全消除。配套的 humans& 文章讨论了更高级的技术，包括 4/6 和符合链式法则的反向传播选择。

重构权重更新接口

低延迟 FlashInfer 后端通常需要进行填充、swizzling、重排以及使用后端特定的权重布局。这些转换对于服务场景而言很自然，但会增加强化学习实时权重更新和 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 以选择反量化反向传播。