迈向Blackwell原生8位与4位强化学习:Miles中的端到端MXFP8与NVFP4强化学习
李梓昂,人类团队及Miles团队
摘要:我们在Miles中实现了两种Blackwell原生的强化学习方案:端到端MXFP8以及针对MoE专家的逐token NVFP4。这两种方案均通过检查点转换、Megatron训练、SGLang推理部署和实时权重更新中的细粒度精度控制提供支持。MXFP8覆盖了推理部署、前向传播、权重梯度通用矩阵乘法以及数据梯度通用矩阵乘法。NVFP4在其MoE专家路径上使用在线逐token激活缩放,两种格式均支持高精度或反量化反向模式。在8块B200上对Qwen3-30B-A3B方案进行消融实验时,BF16与所有五种低精度配置的原始奖励曲线高度重合,而MXFP8和NVFP4则缩短了推理部署时间。
引言
在低精度强化学习中,推理部署、训练、检查点转换和实时权重更新必须遵循统一的精度约定,否则采样器与训练器策略将产生偏差。Miles与SGLang强化学习生态系统已整合低精度方案:LMSYS FP8强化学习文章表明,在训练和采样中统一使用FP8相比BF16训练配合FP8推理部署能减少精度不匹配;INT4 QAT文章则指出,训练期间的伪量化与W4A16推理部署可使INT4具备实用性。我们通过以下方式将该工作扩展至Blackwell原生格式:在Miles中实现MXFP8与NVFP4方案,并向SGLang、TransformerEngine、FlashInfer、Megatron及cuDNN前端贡献相关支持组件。公开路线图详见miles#615。
主要贡献包括:
- 端到端MXFP8强化学习,其中推理部署、前向传播、权重梯度通用矩阵乘法以及数据梯度通用矩阵乘法均使用MXFP8。
- 针对MoE专家权重的逐token NVFP4强化学习,采用在线逐token激活缩放。
- 细粒度精度控制,使得选定张量(如最终层)可始终维持BF16精度。
- TransformerEngine与FlashInfer之间位级精确的量化器约定,确保权重更新不会引入可避免的训练-推理精度不匹配。
为何选择Blackwell原生方案?
此前的低精度方案并非围绕 MXFP8 或 NVFP4 设计。现有的 Miles 路径遵循 DeepSeek-V3 风格的块缩放 FP8 方案:权重采用 128x128 块缩放,激活值采用 1x128 瓦片缩放,缩放因子针对每个瓦片或块在线计算。这是一个优秀的 Hopper 时代方案,但在 Blackwell 上,其 FP32 缩放因子仍需在 Tensor Core 路径外围通过软件应用,而非通过原生微缩放硬件实现。
INT4 QAT 解决的是另一个问题。训练阶段使用伪量化来使模型适配 INT4 权重,而推理阶段则采用 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 推理。
- 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 缩放系数可以在不同的张量作用域中选择。这一选择属于配方决策,而非格式本身的属性,并且在强化学习场景中变得尤为重要。
配方 1:端到端 MXFP8 强化学习
MXFP8 配方是早期端到端 FP8 工作最直接的 Blackwell 原生扩展。模型推理、前向传播、权重梯度通用矩阵乘法以及数据梯度通用矩阵乘法均使用 MXFP8,而选定的张量则通过下文描述的精度控制规则保持 BF16 格式。
训练
TransformerEngine 和 Megatron 将 MXFP8 实现为一条性能优化的一等公民 Blackwell 训练路径,其中包括 deepseek-v3-gb200-optimization.md 中描述的 GB200 DeepSeek-V3 优化。在我们的 Miles 集成中,我们使用此路径作为端到端 MXFP8 强化学习的训练端基础。
与 DeepSeek-V3 FP8 配方的一个区别在于反向激活的表示方式。DeepSeek-V3 将前向激活存储为 1x128 FP8 块,并在反向通用矩阵乘法之前将其转换为反向方向。这种方法存储的 FP8 数据更少,但在反向通用矩阵乘法之前引入了反量化加再量化步骤。TransformerEngine 的 MXFP8 文档指出,按行的 1x32 块和按列的 32x1 块在数值上是不同的,必须从全精度数据独立量化。因此,TransformerEngine 在量化期间会同时生成按行和按列量化的副本。这会使用更多内存,但避免了额外的再量化步骤,并减少了反向路径中的额外量化误差。
这是强化学习中典型的系统权衡。我们使用 TransformerEngine 路径来维护一个端到端的 MXFP8 契约,而不引入另一个不匹配的来源。
模型推理
在模型推理方面,SGLang 使用来自 FlashInfer 和 Triton 的 Blackwell MXFP8 内核。我们在 FlashInfer 和 SGLang 中实现并上游了模型推理路径(flashinfer#2581, sglang#17449, sglang#19537, sglang#21576 和 sglang#28459)。
几乎所有主要的通用矩阵乘法都可以量化为 MXFP8,包括注意力投影和混合专家专家层。主要的例外是显式控制的高精度层,例如下文描述的 BF16 多头潜在注意力投影。
方案二:逐 token NVFP4 强化学习
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)和随机哈达玛变换(RHT)。在该论文中,SR 应用于梯度以减少量化偏差并产生无偏的量化梯度,而 RHT 则用于分散大幅值的块级异常值,尤其是针对权重-梯度 GEMM 的输入。
这是一个很好的起点,但强化学习具有不同的失效模式:
| 场景 | 数值情况 | 核心要点 |
|---|---|---|
| 预训练 | 梯度信号稳定,权重更新幅度大,模型权重具有自适应性,激活值和梯度对精度敏感且动态范围高。 | 保持粗略的优化方向和收敛性。 |
| 强化学习 | 梯度噪声大,奖励方差高,有用的更新信号微小且脆弱。 | 将量化噪声控制在真实更新信号之下;否则,量化噪声可能会覆盖脆弱的模型能力,导致性能崩溃。 |
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 激活值缩放不需要单独的校准缩放因子传递。
训练-推理一致性还需要匹配并行策略。如果在专家张量并行分区内按 token 计算 FP32 缩放因子,那么 SGLang 和 Megatron 应使用相同的 ETP 大小。否则,每一方可能会看到张量的不同分区,并计算出不同的缩放因子。
SwiGLU MoE 层增加了另一个关键约定。SGLang 和 Megatron 通常将门控投影和上投影融合到一个 GEMM 中,因此两个张量在转换和实时权重更新期间必须共享相同的 FP32 缩放因子,即使 Hugging Face 检查点将它们分开存储。Miles 通过在 NVFP4 导出路径中将门控/上投影对一起量化来强制执行此操作。
我们在整个技术栈中实现并向上游提交了逐 token NVFP4 方案:
- TransformerEngine 训练方案:TransformerEngine#2931
- cuDNN 前端训练内核:cudnn-frontend#251
- FlashInfer rollout 内核:flashinfer#3027
- SGLang 集成和权重更新修复:sglang#22918, sglang#22204
高精度和反量化反向传播
在高精度反向 NVFP4 变体中,前向传播和 rollout 对 MoE 专家使用 NVFP4,而反向 GEMM 则使用原始的 BF16 操作数。
反量化反向是第二种可选模式。反向 GEMM 仍以 BF16 运行,但消耗的是前向过程中产生的精确低精度操作数的 BF16 反量化值,而非原始的 BF16 值。
两种模式都避免了低精度反向 GEMM,因此这些配置不使用原始 NVFP4 预训练方案中的 RHT 或随机舍入。它们以反向吞吐量为代价换取更高精度的计算,但强化学习通常受 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:
| 模式 | 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:
| 模式 | 前向分配 | 反向分配 | 端到端分配 | 前向预留 | 反向预留 | 端到端预留 | 前向增量 | 反向增量 | 端到端增量 |
|---|---|---|---|---|---|---|---|---|---|
| 默认 | 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%) |
比特精确量化契约
在强化学习中,量化不一致性会在权重更新过程中累积。如果训练侧和推理侧对张量的量化方式不同,那么用于采样和学习的策略就不再是同一个低精度模型。因此,我们需要在强化学习栈的两端之间建立明确的契约。
我们将 FlashInfer 和 TransformerEngine 的量化器对齐到相同的 MXFP8 和 NVFP4 比特级契约。我们的 FlashInfer 单元测试针对随机数据、量化边界数据、全零张量和最大值张量(对应 init_data = ["random", "boundary", "zeros", "maxes"]),检查与 TransformerEngine 风格参考实现之间精确的字节级一致性。我们在 flashinfer#3387 中实现并上游提交了此量化器对齐。
推理服务和强化学习之间存在一个实际区别。对于纯推理服务的工作负载,FlashInfer 可能在 FP4 量化路径的某些部分使用快速数学运算以提升性能。这是合理的推理服务默认设置,但强化学习的权重更新受益于与训练侧量化器的精确一致。对于本方案,我们设置:
FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1
每个涉及推理权重的后端要么精确实现此量化契约,要么将近似行为设为可选加入。
细粒度精度控制
在实践中,单一的全局精度开关对于低精度强化学习是不够的。某些张量应保留在 BF16 中,但选择它们只是问题的一部分:同样的决策必须在 Hugging Face 检查点转换、Megatron 训练、SGLang 推理服务和实时权重导出中强制执行。
我们在 Miles 中通过基于计数和基于名称的 BF16 异常机制,在检查点转换、训练、推理 rollout 和线上导出等环节实现了这种张量级精度控制(对应 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 并未带来显著的训练-推理不一致性降低效果:
将共享专家保持在高精度也能减少训练-推理不一致性,且性能或内存开销很小。路由专家是稀疏选择的,其输出会经过高精度的加权归约。而共享专家始终处于激活状态,因此其精度误差会影响通过该模块的每一个 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"
匹配的基于名称的转换时和 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
对于 DeepSeek-V3 风格的 MLA,这个 BF16 异常的开销很小。一个形状为 32768 x 512 的 `kv_b_proj` 张量,在 61 层中以 BF16 格式存储,大约占用 1.90625 GB。
结果:Qwen3-30B-A3B 在 8x B200 上的配方消融实验
为确保对比一致性,所有实验均采用同步 Qwen3-30B-A3B RL,并使用默认的 Miles 配置在 8 块 B200 上运行。固定工作负载采用 GRPO 风格训练,数据集为 dapo-math-17k,每个提示词生成 8 个 rollout 样本,最大响应长度为 8192 个 token。这仅是一个配方消融实验设置,并非完全调优后的训练或服务基准测试。KL 路径已启用用于诊断,但其系数为 0.0,因此在此消融实验中 KL 不作为优化惩罚项。
硬件分配为 4 块 GPU 用于 rollout,4 块 GPU 用于训练。
对于所有低精度配方:
- 启用 MoE rollout 路由重放。
- 最后 15% 的层保持 BF16 精度。
- 低精度权重使用 0 权重衰减以保持稳定性。
- SGLang rollout 使用 BF16 KV 缓存和 FlashInfer TRTLLM 路由 MoE 后端来处理低精度 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 的部署,FlashInfer 直接在激活量化内核路径中计算每个 token 的在线 FP32 缩放因子,而不是将其作为单独的预处理步骤。因此,所报告的部署性能已包含在线缩放因子计算的开销。
在训练方面,MXFP8 变体比 BF16 更快,而在此处测量的实现中,NVFP4 反向传播覆盖变体则更慢:
训练侧的差距源于本次消融实验所使用的实现方式,而非 FP4 Tensor Core 的固有限制。我们的 TransformerEngine 路径将每个 token 的 FP32 缩放作为单独的 PyTorch 操作(TransformerEngine#2931)来实现,而不是采用将缩放融合到内核 epilogue 中的原生 per-token NVFP4 GEMM 路径。我们已经实现并向上游提交了融合后的 cuDNN 前端内核(cudnn-frontend#251);TransformerEngine 的集成仍在等待中。反量化反向传播增加了上述的反量化步骤。由于该强化学习工作负载以部署为主,即使在训练路径完全加速之前,部署加速仍然具有重要意义。
除了本次消融实验之外,humans& 在生产环境中使用相同的配方系列和组件,用于大规模、长上下文、多智能体异步强化学习研究。
未来工作
移除额外的 BF16 权重副本
尽管部署和训练执行相同的低精度配方,但 Megatron 仍然会保存一个额外的 BF16 权重副本。这会增加内存消耗,并限制了低精度路径的实际内存优势。
Megatron 具有 `--fp8-param-gather` 和 `--fp4-param-gather` 选项,但 Blackwell 原生路径仍在完善中。Megatron-Bridge 的跟踪问题 Megatron-Bridge#3801 反映了为支持稳健的低精度参数收集而正在进行的工作。NVFP4 的 `--fp4-param-gather` 尚不支持此配方所使用的 1D 1x16 权重布局。
偶发的梯度尖峰
高精度反向传播的 NVFP4 变体仍可能出现偶发的梯度尖峰:
反量化反向传播减少了本次消融实验中最大的尖峰,但并未完全消除。更高级的技术,包括 4/6 和链式法则一致的反向传播选择,在配套的 humans& 博文中进行了讨论。
重构权重更新接口
低延迟的 FlashInfer 后端通常需要填充、交织、混洗以及后端特定的权重布局。这些变换对于推理服务来说是自然的,但它们会使实时的 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 计算团队(Siyuan Fu、Yigong Qin、Zhongbo Zhu)、TransformerEngine 团队和 FlashInfer 团队。
我们还要感谢 Cursor 团队提出的每个 token 的 NVFP4 激活缩放思路。