Miles 在 Blackwell 架构上实现端到端 MXFP8 与逐 token NVFP4 强化学习方案
Blog Towards Blackwell-Native 8-bit and 4-bit RL: End-to-End MXFP8 and NVFP4 RL in Miles TL;DR: We implemented two Blackwell-native RL recipes in Miles: end-to-end MXFP8 and per-token NVFP4 for MoE experts. Both are supported by fine-grained precision control across checkpoint conversion,... Ziang Li, humans& and Miles Team
Miles 团队在 Blackwell 架构上实现了两种原生低精度强化学习方案:端到端 MXFP8 和 MoE 专家权重的逐 token NVFP4。在 8x B200 上对 Qwen3-30B-A3B 的消融实验中,BF16 与所有五种低精度配置的原始奖励曲线高度重合,且 MXFP8 和 NVFP4 减少了推理时间。
虽然只针对Blackwell,但这份报告把MXFP8和NVFP4在RL里的实现细节讲清楚了,从量化合约到每token缩放都给了可复现配置,做大规模RL训练的能直接上手。
简而言之:我们在 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。
主要贡献如下:
- 端到端 MXFP8 RL,其中 rollout、前向传播、权重梯度 GEMM 和数据梯度 GEMM 全部使用 MXFP8。
- 针对 MoE 专家权重,采用逐 token 的 NVFP4 强化学习,并使用在线逐 token 激活缩放。
- 细粒度精度控制,因此诸如最终层等选定张量可以始终保持 BF16。
- 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 表格。
| 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 |
对于强化学习系统,精度契约涵盖:
- 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 块。
FP32 缩放因子可以在不同的张量范围上进行选择。这一选择属于方案层面的决策,而非格式本身的属性,并且它对 RL 尤为重要。
方案一:端到端 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 将前向激活值存储为 1x128 FP8 分块,并在反向 GEMM 之前将其转换为反向方向。这种做法存储的 FP8 数据更少,但在反向 GEMM 之前引入了反量化加重新量化。TransformerEngine 的 MXFP8 文档指出,行向 1x32 分块和列向 32x1 分块在数值上是不同的,必须从全精度数据独立量化。因此,TransformerEngine 在量化过程中会同时物化行向和列向的量化副本。这会占用更多内存,但避免了额外的重新量化步骤,并减少了反向路径中额外的量化误差。
这是强化学习中典型的系统权衡。我们采用 TransformerEngine 这条路径,以保持单一的端到端 MXFP8 契约,而不引入另一个不匹配来源。
推出
在推出侧,SGLang 使用了来自 FlashInfer 和 Triton 的 Blackwell MXFP8 kernel。我们在 FlashInfer 和 SGLang 中实现并向上游贡献了推出路径(flashinfer#2581、sglang#17449、sglang#19537、sglang#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 的输入。
这是一个不错的起点,但 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 方案:
- TransformerEngine 训练方案:TransformerEngine#2931
- cuDNN 前端训练 kernel:cudnn-frontend#251
- FlashInfer rollout kernel:flashinfer#3027
- SGLang 集成与权重更新修复:sglang#22918、sglang#22204
高精度与反量化反向传播
在高精度反向传播的 NVFP4 变体中,前向传播和 rollout 对 MoE 专家使用 NVFP4,而反向 GEMM 使用原始 BF16 操作数。
反量化反向传播是第二种可选模式。反向 GEMM 仍以 BF16 运行,但使用的是前向过程中生成的精确低精度操作数的 BF16 反量化结果,而非原始的 BF16 值。
两种模式都避免了低精度反向 GEMM,因此这些配置不使用原始 NVFP4 预训练方案中的 RHT 或随机舍入。它们以反向吞吐量换取更高精度的计算,但 RL 通常受限于 rollout,而长上下文注意力和通信进一步降低了端到端的影响。
相同的反向模式选择同样适用于 MXFP8:
我们实现并向上游贡献了 NVTE_BACKWARD_OVERRIDE,作为可复用的 TransformerEngine 接口,用于选择高精度或反量化的反向操作数(TransformerEngine#2644),随后通过 Miles 配方配置暴露了这两种模式。配套的 humans& 文章涵盖了算法动机以及额外的反向模式消融实验。
反向成本与内存
反量化反向传播增加了一个训练侧的反量化步骤。我们与 NVIDIA 合作,在 TransformerEngine#2865 中降低了其开销。
相对于 TransformerEngine 默认的低精度反向路径,高精度和反量化反向传播还可以降低峰值内存。这两种模式都不需要生成并保留上文 MXFP8 训练部分所述低精度反向 GEMM 所使用的第二份按列量化的副本。
我们在验证 TransformerEngine#2644 中的 TransformerEngine 反向模式实现时,测量了下面的内存数据。alloc 列报告已分配内存,resrv 列报告已预留内存,所有数值均以 MB 为单位。
MXFP8 线性层内存,dtype=torch.bfloat16,input_shape=(2048, 2048),out_features=8192:
| 模式 | 前向分配 | 反向分配 | 端到端分配 | 前向预留 | 反向预留 | 端到端预留 | 前向差值 | 反向差值 | 端到端差值 |
|---|---|---|---|---|---|---|---|---|---|
| 默认 | 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:
| 模式 | 前向分配 | 反向分配 | 端到端分配 | 前向预留 | 反向预留 | 端到端预留 | 前向增量 | 反向增量 | 端到端增量 |
|---|---|---|---|---|---|---|---|---|---|
| 默认 | 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%) |
位精确量化器契约
在强化学习中,量化不一致会在权重更新过程中不断累积。如果训练侧和 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#614、miles#1054 和 miles#1261)。我们还为由此产生的混合精度 checkpoint 实现了 SGLang 支持(sglang#18742 和 sglang#20214)。具体而言,转换使用 --num-layers-at-start-in-bf16 和 --num-layers-at-end-in-bf16;Megatron 训练将这些计数与 --first-last-layers-bf16 结合;SGLang 则服务由此产生的混合精度 checkpoint。
层级精度选择
按照 NVIDIA NVFP4 预训练论文中的建议,我们将一小部分最后几层保持在较高精度。在我们的实验中,将最后 15% 的层保持为 BF16 能显著减少训练-推理不匹配并稳定梯度。
在我们的实验中,对早期层启用 BF16 并不会带来有意义的训练-推理不匹配减少:
将共享专家保持在高精度同样能减少训练-推理不匹配,且性能或内存开销很小。路由专家是稀疏选择的,其输出会经过高精度加权归约。共享专家始终处于激活状态,因此其精度误差会影响通过该模块的每一个 token。
案例研究: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 后端。
我们比较六种配置:
- 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 参考模型计算这一诊断指标的,因此该指标除了包含 RL 过程中的策略演化外,还包含了每种低精度策略与 BF16 参考模型之间的差异。不应将其单独解读为优化惩罚。
奖励
尽管诊断不匹配程度更高,但五条低精度奖励曲线都与 BF16 奖励曲线紧密吻合。
这是配方消融实验的关键结果:在这个 Qwen3-30B-A3B B200 配置中,Blackwell 原生低精度在保持所观测到的学习曲线的同时,提升了 rollout 效率。
性能
与 BF16 相比,MXFP8 和 NVFP4 都缩短了 rollout 时间:
对于 NVFP4 rollout,FlashInfer 直接在激活量化 kernel 路径内部计算在线 per-token FP32 scale,而不是将其作为单独的预处理步骤。因此,所报告的 rollout 性能包含了在线 scale 计算的开销。
在训练侧,MXFP8 变体比 BF16 更快,而 NVFP4 backward-override 变体在此处所测量的实现中更慢:
训练侧的差距来自此消融实验中所使用的实现,而非 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 变体仍可能出现偶发的梯度尖峰:
在此消融实验中,反量化反向传播减少了最大的尖峰,但并未将其消除。更先进的技术,包括 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