- 摘要
- 摘要
- 1. 引言
- 2. 原生 NVFP4 强化学习为何会变得不稳定
- 2.1 训练过程中出现的失败
- 2.2 量化误差、转换误差与执行失配
- 2.3 激活精度、缩放粒度与 KV 缓存
- 2.4 失配必须结合更新方向来解读
- 2.5 平均值可能掩盖局部风险
- 3. TRIAGE:原生 NVFP4 强化学习的方向感知稳定化
- 3.1 机制
- 3.2 稳定性与任务性能
- 3.3 互补的算法与工程路线
- 4. miniTE:对齐数值执行并优化低精度计算
- 4.1 控制数值差异与执行成本
- 4.2 分层接口与 Transformer Engine 集成
- 4.3 精度配方与选择性量化
- 4.4 数值监控、性能分析与对齐诊断
- 4.5 训练–推理对齐与确定性执行(确定性模式)
- 5. 与 Transformer Engine 联合实现和优化 NVFP4
- 5.1 逐 token 缩放的数值与实现选择
- 5.2 连续布局与分组 GEMM 调用
- 5.3 在数据产生处计算 amax
- 5.4 内核优化及其适用范围
- 5.5 联合测试如何指导实现选择
- 5.6 探索单遍融合
- 6. 评估:数值正确性、学习与执行效率
- 6.1 实验设置与测量范围
- 6.2 正确性与学习行为
- 6.3 完整学习器更新:局部改动能否降低更新成本?
- 6.4 模块:前向更快不一定意味着前向+反向更快
- 6.5 Rollout 与在线强化学习:系统收益与剩余成本
- 6.6 实验所确立的结论
- 7. 使用与扩展 miniTE
- 7.1 通过线性接口集成现有模型
- 7.2 显式配方与精度边界
- 7.3 将诊断作为扩展开发的一部分
- 8. 讨论:从数值对齐到端到端学习
- 8.1 权重更新与同步仍是瓶颈
- 8.2 算子对齐后低精度为何仍然敏感?
- 8.3 优化更新与最终学习结果
- 附录 A. 实验设置与统计
- A.1 补充算法与训练材料
- A.2 补充的长程强化学习与检查点评估
- A.3 固定文本的 miniTE 学习器测试
- A.4 模块、服务与强化学习各阶段耗时
- 附录 B. TRIAGE 方法约定
- 参考文献
摘要
- 问题与诊断 · 原生 NVFP4 强化学习可能在训练后期变得不稳定。我们考察量化、激活缩放、KV 缓存与执行路径,并研究它们的误差如何影响更新。
- 算法与工程 · TRIAGE 通过更新方向控制残余失配风险。miniTE 降低数值执行差异并优化低精度计算。我们分别对二者进行评估。
- 实测收益 · 在 8×B300 上,采用 BF16 反向传播的 miniTE 在完整学习器更新上相对 BF16 达到约 1.33–1.37× 的加速。另一项固定输出服务测试达到约 1.64× 的吞吐。
- 局限与待解问题 · 在所报告的在线强化学习运行中,权重转换与同步抵消了采样收益。实用的 miniTE 检查点在五项任务平均上仍低于 BF16。算子匹配并取得小幅差距并不能保证达到 BF16 级别的学习效果;误差容忍度、更新规则与最终结果仍需进一步研究。
摘要
原生 NVFP4 强化学习必须在量化与执行误差下保持稳定,同时带来实际的效率收益。我们通过考察量化格式、激活缩放、KV 缓存精度以及影响策略概率的框架差异,研究训练后期的不稳定性。TRIAGE 根据失配方向与更新方向,对高风险片段中的更新进行调整,以应对残余失配,从而提升稳定性与任务性能。miniTE 将量化策略、线性层与计算后端分离,提供精度配置、对齐检查与性能分析。针对 Transformer Engine 逐 token NVFP4 路径的联合工作采用连续布局、设备端元数据、amax(最大绝对值)复用、分组 GEMM 与融合,以减少准备与调度开销。在 8×B300 上,两条采用 BF16 反向传播的 miniTE 路径在完整学习器更新上相对 BF16 均达到约 1.33–1.37× 的配对加速。另一项固定输出服务测试达到约 1.64× 的 BF16 吞吐。这些学习与系统结果来自不同的配置;更低的失配与更快的局部计算尚不能确立 BF16 级别的学习效果或端到端强化学习加速。
1. 引言
大语言模型的强化学习会反复执行 rollout、策略概率重算、梯度更新与权重同步。长回复增加了生成与训练的 token 数量,使低精度计算更具价值。Blackwell 支持原生 NVFP4 矩阵乘法,此前的预训练工作已通过缩放、混合精度与合适的舍入策略展示了低比特训练 [1]。强化学习循环还必须连接训练框架与推理框架所使用的数值实现。
我们早期的 NVFP4 运行最初提升了奖励,但随后出现熵异常、训练–采样对数概率差距扩大以及奖励下降。量化误差、检查点转换、激活缩放、KV 缓存精度以及两侧的算子执行都影响了结果。替换单个量化算子,或确认两侧都使用 NVFP4,并未解决所有这些问题。
我们首先描述这些失败,然后通过受控实验识别并定位其原因。这引出了互补的算法与工程路线。TRIAGE 改变更新规则,以控制残余失配在学习中的影响。miniTE 作用于数值执行,在减少可避免差异的同时,将低精度计算转化为实际的效率收益。
在大型训练框架中,改变缩放粒度或 GEMM 后端可能需要修改层封装、分组布局、通信和自动微分。miniTE 将量化策略与这些共享功能分离,使开发者能够在现有训练工作流中替换某条计算路径,然后评估其数值正确性、学习行为和运行时开销。
第 2 节追溯原生 NVFP4 强化学习的失败及其诊断过程。第 3 节介绍 TRIAGE 及其结果。第 4–7 节描述 miniTE、实现优化、实验和使用方法。我们分别报告算法结果和系统性能,并给出每项实验的配置。
2. 原生 NVFP4 强化学习为何变得不稳定
2.1 失败在训练过程中逐渐发展
我们在 DAPO 数学任务上使用 Qwen3-4B 和 Qwen3-30B-A3B 运行 GRPO [4,5,6]。在早期的原生 NVFP4 运行中,奖励先正常上升,随后训练开始恶化。在 TRIAGE 分析的轨迹中,4B 模型在约 300 步后变得不稳定,而 30B 模型的失配在约 700 步后开始迅速增大。诊断这一过程需要在训练时间上进行测量,而不是依赖单个初始误差值或一段较短的损失曲线。
我们跟踪奖励、熵、训练–采样对数概率差距以及下游任务得分。对于回复 iii 中的 token ttt,定义 δi,t=logπl(yi,t∣xi,yi,<t)−logπs(yi,t∣xi,yi,<t)\delta_{i,t}=\log\pi_{\mathrm l}(y_{i,t}\mid x_i,y_{i,<t})-\log\pi_{\mathrm s}(y_{i,t}\mid x_i,y_{i,<t})δi,t=logπl(yi,t∣xi,yi,<t)−logπs(yi,t∣xi,yi,<t)。其中 xix_ixi 是提示词,yi,ty_{i,t}yi,t 是当前 token,yi,<ty_{i,<t}yi,<t 是其已生成的前缀。学习器和采样器分别赋予条件概率 πl\pi_{\mathrm l}πl 和 πs\pi_{\mathrm s}πs;log\loglog 表示自然对数。负差距表示学习器赋予的概率更低,正差距表示赋予的概率更高。平均绝对差距对参与评估的 token 取 ∣δi,t∣|\delta_{i,t}|∣δi,t∣ 的平均值。
2.2 量化误差、转换误差与执行失配
量化误差将低精度计算与更高精度的参考结果进行比较。训练–采样失配则比较学习器和采样器的实际执行。检查点转换和推理后端适配可能引入额外误差。因此,在集成强化学习之前,我们比较 BF16 推理、在线 NVFP4 量化,以及使用离线转换的 NVFP4 权重进行推理。
NVFP4 以 E2M1 表示量化值:E 和 M 分别给出指数位和尾数位的数量,另有一个符号位。解码得到 x^=q,sb,sg\hat x=q,s_b,s_gx^=q,sb,sg,其中 qqq 是 E2M1 值,sbs_bsb 是由 16 个元素共享的 E4M3 块缩放因子,sgs_gsg 是 FP32 外层缩放因子。该方案决定外层缩放因子是按张量、按专家还是按行计算。
E2M1 数据占 4 位。每 16 个元素组成一个块,共享一个 E4M3 缩放因子,并另有一个外层 FP32 缩放因子。点击载荷位并更改缩放因子,可查看有限值和重建结果。示例缩放因子均为可表示的值。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
一个 Qwen3-0.6B NVFP4 预训练检查点最初出现严重的性能下降,原因是其融合 QKV 结构被错误处理。修复后,BF16 的平均 Lambda 准确率约为 0.3043,一维量化约为 0.3056,二维块缩放量化约为 0.3109。因此,在将此类失败归因于 FP4 格式之前,必须先检查转换和结构适配。
对于 Qwen3-4B,在线量化使 Lambda-base 和 Math23K 得分相对 BF16 分别下降 0.0355 和 0.047。离线转换又额外带来 0.0089 和 0.019 的下降。这些比较将格式与数值路径造成的损失,与转换和缩放因子处理引入的额外损失区分开来。
在强化学习中,即使两侧使用相同格式,也可能存在显著差异。早期配置产生了以下平均绝对对数概率差距:
| 训练 / 采样精度 | 稠密 | MoE |
|---|---|---|
| BF16/BF16 | 0.0110 | 0.0172 |
| NVFP4/NVFP4 | 0.1142 | 0.0668 |
| NVFP4/BF16 | 0.1376 | 0.0621 |
| BF16/NVFP4 | 0.1690 | 0.0747 |
表 1. 早期训练–采样比较。稠密模型为 Qwen3-4B;MoE 模型为 Qwen3-30B-A3B。这些结果在这组配置中定位了误差来源。
用在线量化替代离线 ModelOpt 转换后,稠密模型仍存在约 0.1142 的差距。两侧使用相同的量化/GEMM 算子也仍残留差异。因此,我们将调查扩展到输入组织、缩放因子计算以及量化周边的运算。
2.3 激活精度、缩放粒度与 KV 缓存
在固定 NVFP4 权重的情况下,提高激活精度可检验激活量化的影响。W 和 A 分别表示权重和激活;以下数字给出其位宽,因此 W4A16 表示 4 位权重和 16 位激活:
| 权重 / 激活 | 稠密差距 | MoE 差距 |
|---|---|---|
| W4A16 | 0.0139 | 0.0166 |
| W4A8,MXFP8 激活 | 0.0363 | 0.0227 |
| W4A4,NVFP4 激活 | 0.0903 | 0.0445 |
表 2. 使用相应 unquant 推理适配路径的激活精度比较。本实验与表 1 属于调查的不同阶段,并非单一消融实验。
W4A16 使差距接近 BF16 参考值,表明激活量化是该配置中的一个重要因素。在另一固定配置中,将 KV 缓存从 FP8 改为 BF16 使差距从 0.0721 降至 0.0583。采样器状态精度也会影响跨框架一致性。
接下来我们考察缩放域。采样器解码处理当前 token,而学习器通常一次性处理完整响应。如果激活的外层缩放使用张量级 amax,那么同一个 token 可能因批次中其他 token 的不同而获得不同的缩放。逐 token amax 将归约域限制在当前行,从而消除了这种依赖。
设 xxx 为参考向量,x^\hat xx^ 为其解码近似值,二者均含 nnn 个元素。相对 L2 误差为 ∣x^−x∣2/∣x∣2|\hat x-x|_2/|x|_2∣x^−x∣2/∣x∣2;信噪比为 SNR=20log10(∣x∣2/∣x^−x∣2)\mathrm{SNR}=20\log_{10}(|x|_2/|\hat x-x|_2)SNR=20log10(∣x∣2/∣x^−x∣2),单位为 dB;平均绝对误差为 MAE=n−1∑j=1n∣x^j−xj∣\mathrm{MAE}=n^{-1}\sum_{j=1}^{n}|\hat x_j-x_j|MAE=n−1∑j=1n∣x^j−xj∣。记号 ∣⋅∣2|\cdot|_2∣⋅∣2 表示欧几里得范数。相对误差和 SNR 假定参考向量非零。相对 L2 和 MAE 越低,或 SNR 越高,表示重建越接近。
目标 token 的 32 个输入保持不变;滑块仅改变批次中另一个 token 的 amax。采用就近舍入到偶数的有限 E2M1/E4M3 码本展示了块缩放、重建值和相对 L2 误差。这是一个数学示例,而非模型或 GPU 测量;它不包含 RHT、SR、FTZ 和填充。印刷版展示了默认幅度下构造的输入和重建值。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
缩放粒度并不能决定所有量化误差。我们在真实 QKV 输入上扫描指数位宽和尾数位宽。在 E=2 且 E4M3 缩放固定的情况下,将 M 从 1 增加到 4 使平均 SNR 从 19.88 dB 提升到 34.96 dB,并将相对 L2 从 0.10140 降至 0.01807。增加指数位宽并不总是有益,且可能与缩放的可表示范围产生不良交互。这些模拟格式用于识别敏感因素;它们并非额外的硬件原生 NVFP4 格式。
在一种配置中,逐 token amax 仅略微改变了初始差距,从 0.05532 变为 0.05501,而后续运行显示出更平滑的训练。这促使我们考察误差方向、局部结构及其在训练过程中的演化,而不是仅依赖初始平均值。
D1 · 指数与尾数扫描:真实 QKV 输入、模拟与初始 RL 测量真实 QKV 输入:改变尾数位宽
| 格式 | SNR E4M3 | SNR E8M0 | 相对 L2 E4M3 | 相对 L2 E8M0 |
|---|---|---|---|---|
| E2M1 | 19.88 | 17.11 | 0.10140 | 0.13950 |
| E2M2 | 26.23 | 23.98 | 0.04885 | 0.06322 |
| E2M3 | 31.19 | 30.31 | 0.02767 | 0.03051 |
| E2M4 | 34.96 | 36.41 | 0.01807 | 0.01513 |
表 2-1. 在真实预填充和单 token 生成过程中从 36 层收集的 QKV 输入,跨层汇总。模拟格式并不意味着新的硬件支持。
真实输入:指数范围与缩放范围
| 数据格式 | SNR(E4M3,dB) | SNR(E8M0,dB) |
|---|---|---|
| E2M1 | 19.88 | 17.11 |
| E3M1 | 21.50 | 19.7 |
| E4M1 | 14.39 | 19.7 |
| E5M1 | 0.0(最大值超出 E4M3 范围) | 19.7 |
| 数据格式 | 相对 L2(E4M3) | 相对 L2(E8M0) |
|---|---|---|
| E2M1 | 0.10140 | 0.13950 |
| E3M1 | 0.08478 | ~0.103 |
| E4M1 | 0.28937 | ~0.103 |
| E5M1 | 0.99698(最大值超出 E4M3 范围) | ~0.103 |
在测量输入范围内,E3M1/E4M1/E5M1 配合 E8M0 缩放表现相似。将 E4M3 缩放与更宽的数据范围结合可能产生缩放范围问题。因此,对于这些输入,更多指数位并不一定能提高精度。
初始 RL 差距:指数与尾数位宽
| 数据格式 | E4M3 下的初始差距 |
|---|---|
| E2M1 | 0.05800700932741165 |
| E3M1 | 0.11516289412975311 |
| E4M1 | 3.412203073501587 |
| E5M1 | nan |
| 数据格式 | E4M3 下的初始差距 |
|---|---|
| E2M1 | 0.05800700932741165 |
| E2M2 | 0.042353976517915726 |
| E2M3 | 0.03562581539154053 |
| E2M4 | 0.031815189868211746 |
表 2-2. 记录的初始差距,包括 NaN 结果。这些测量的是初始一致性;长期训练另行评估。
模拟输入与真实输入比较
| 数据格式 | 缩放格式 | MAE | SNR dB |
|---|---|---|---|
| E5M1 | E4M3 | 1.485030 | 0.00 |
| E4M1 | E4M3 | 0.132963 | 20.48 |
| E3M1 | E4M3 | 0.127780 | 20.63 |
| E2M1 | E4M3 | 0.154619 | 19.62 |
| E5M1 | E8M0 | 0.128281 | 20.18 |
| E4M1 | E8M0 | 0.128281 | 20.18 |
| E3M1 | E8M0 | 0.128515 | 20.18 |
| E2M1 | E8M0 | 0.176643 | 17.97 |
| 数据格式 | 缩放格式 | MAE | SNR dB |
|---|---|---|---|
| E2M4 | E4M3 | 0.017888 | 36.91 |
| E2M3 | E4M3 | 0.034627 | 31.58 |
| E2M2 | E4M3 | 0.070138 | 26.04 |
| E2M1 | E4M3 | 0.142116 | 20.17 |
| E2M4 | E8M0 | 0.018970 | 37.03 |
| E2M3 | E8M0 | 0.038432 | 30.93 |
| E2M2 | E8M0 | 0.079777 | 24.93 |
| E2M1 | E8M0 | 0.176643 | 17.97 |
表 2-3. 来自两个模拟输入批次的 MAE/SNR,分别对应指数和尾数扫描。它们将输入分布效应与位宽效应区分开来,并与真实 QKV 均值的重复测量不同。完整的格式汇总如下。
| 格式 | 平均 SNR | 最差 SNR | 平均 MAE | 平均相对 L2 |
|---|---|---|---|---|
| E2M4 + E8M0 | 36.41 | 35.87 | 0.007954 | 0.01513 |
| E2M4 + E4M3 | 34.96 | 29.58 | 0.007589 | 0.01807 |
| E2M3 + E4M3 | 31.19 | 28.12 | 0.014404 | 0.02767 |
| E2M3 + E8M0 | 30.31 | 29.73 | 0.016283 | 0.03051 |
| E2M2 + E4M3 | 26.23 | 25.03 | 0.027865 | 0.04885 |
| E2M2 + E8M0 | 23.98 | 23.64 | 0.034200 | 0.06322 |
| E3M1 + E4M3 | 21.50 | 17.81 | 0.045411 | 0.08478 |
| E2M1 + E4M3 | 19.88 | 19.53 | 0.058788 | 0.10140 |
| E5M1/E4M1/E3M1 + E8M0 | ~19.7 | ~19.3 | ~0.051 | ~0.103 |
| E2M1 + E8M0 | 17.11 | 16.66 | 0.077184 | 0.13950 |
| E4M1 + E4M3 | 14.39 | 0.71 | 0.067209 | 0.28937 |
| E5M1 + E4M3 | 0.03 | 0.00 | 0.577118 | 0.99698 |
2.4 失配必须结合更新方向来解读
设 AiA_iAi 表示响应 iii 的优势,衡量其相对于同组其他响应的奖励。图中将其缩写为 AAA,并将 token 的差距写作 δ\deltaδ。在 GRPO 中,正优势倾向于提高采样 token 的概率;负优势倾向于降低它们。在同步权重下,差距反映执行差异。在学习器更新之后,它还包括相对于固定采样策略的参数变化。
当学习器低估某个 token(δ < 0)时,负优势更新会进一步降低其概率,并可能放大局部差异;正优势更新则倾向于减小它。当学习器高估该 token(δ > 0)时,角色反转。
因此,干预既需要差距幅度,也需要更新方向。TRIAGE 利用优势与差距的符号来识别局部风险。这种局部分析为算法提供了动机;实际更新还取决于跨 token 耦合与优化器。
改变优势与差距的符号,以考察局部放大或收缩以及 TRIAGE 的方向条件。这说明了第 2.4 节中的局部分析。门控激活还取决于片段风险,而实际更新取决于跨 token 耦合与优化器。
A > 0,δ < 0
学习器低估,正优势提高概率,局部上趋近采样器。TRIAGE 修复则单独作用于正优势响应中的负向片段。
A > 0,δ > 0
这可能在局部放大失配。TRIAGE 并不会衰减每一个放大象限:其门控方向条件为负优势和负差距。
A < 0,δ < 0
学习器低估,负优势进一步降低概率。这满足 TRIAGE 门控方向条件;衰减还取决于片段平均差距和尾部风险。
A < 0,δ > 0
学习器高估,负优势降低概率,局部上趋近采样器。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
2.5 平均值可能掩盖局部风险
在每个模型的一条诊断轨迹中,方向性变化先于全局尾部的爆炸出现。当把响应划分为 64 token 的片段时,一小部分尾部密集片段造成了部分严重失配。包含这些片段的响应仍可能具有较小的平均差距。

实验结果
图 1. 来自 TRIAGE 补充分析的方向性分布与负尾部。上方面板显示按原始 ∣A∣|A|∣A∣ 加权的区域占比;下方面板显示负尾部深度 −q0.05(δ)-q_{0.05}(\delta)−q0.05(δ),其中 q0.05q_{0.05}q0.05 是 token 差距的第 5 百分位数,∣A∣|A|∣A∣ 是绝对优势。4B 分析使用直接差距;30B 分析使用路由重放后的残差差距。阶段定义和局部阈值见附录。
这些发现产生了两个需求:一种在现有技术栈中管理高风险更新的算法,以及用于比较规模、精度和算子路径的共享接口。接下来我们介绍 TRIAGE,然后描述 miniTE 如何将对齐、实验和高效执行整合到一个可复用的框架中。
诊断 · §2
在 NVFP4 权重固定的情况下,W4A16 使差距接近 BF16 参考水平;将 KV 缓存从 FP8 移至 BF16 可将差距从 0.0721 降至 0.0583。
3. TRIAGE:原生 NVFP4 强化学习的方向感知稳定化
3.1 机制
TRIAGE 作用于一个已包含重要性修正(如截断重要性采样(TIS)[10])的策略目标。它保留原生 NVFP4 前向计算,并通过修改优化目标来管理残差失配。该方法结合了片段诊断、方向选择性门控和有界修复。
门控使用 64 token 片段中的平均差距和深尾部 token 计数来评估风险。在具有负优势的高风险片段中,它仅对负差距 token 衰减更新权重;所有其他 token 保留其原始权重。随后在每个响应内对加权损失进行归一化,重新分配更新贡献。
修复为正优势响应中的负向片段偏移添加信号。它对 128 token 片段中的平均差距施加具有有界导数的单侧伪 Huber 惩罚,提高活跃片段中学习器的对数概率。门控减少高风险更新,而修复提供片段级纠正信号。参数和归一化在附录 B 中说明。
3.2 稳定性与任务性能
TRIAGE 的 NVFP4 基线已在两侧使用逐 token 激活缩放。实验在该原生技术栈中测量残差失配。TRIAGE 对 4B 完成 600 个训练步,对 30B 完成 1,700 个训练步。TIS 延迟但并未阻止 4B 的不稳定;对于 30B,它完成了训练,但后期性能较弱。
该图使用熵、奖励和差距比较两个模型上的四种训练方法。可切换模型、指标或步数范围。轨迹在其记录的端点处停止。默认显示原始值;可选择居中的 15/21 点均值。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
图 2. TRIAGE 训练动态。上行为 30B,下行为 4B。从左到右:熵损失、奖励和平均绝对差距。
| 配置 | 4B checkpoint | 五项数学任务均值 | 30B checkpoint | 五项数学任务均值 |
|---|---|---|---|---|
| BF16 | 600 | 58.26% | 1,700 | 72.41% |
| NVFP4 | 300† | 47.10% | 700† | 65.87% |
| NVFP4 + TIS | 300† | 50.76% | 1,700 | 63.18% |
| NVFP4 + TRIAGE | 600 | 58.49% | 1,700 | 70.96% |
表 3. AIME24、AIME25、MATH-500、AMC23 和 GSM8K [11,12,13] 上的未加权均值。† 标记在不稳定之前保存的 checkpoint。不同训练时长的最终结果与等预算比较相区分。
在共享的 4B 第 300 步 checkpoint 上,TRIAGE 得分为 53.03%,分别超过原生 NVFP4 和 TIS 5.93 和 2.27 个百分点。因此其收益在等训练预算下也显现出来。框架吞吐量在 miniTE 测试环境中评估。
3.3 互补的算法与工程路线
TRIAGE 和 miniTE 在上述问题中的不同环节进行干预。TRIAGE 在残差失配仍然存在时改变更新。miniTE 直接检查并改进量化、算子及其调用,同时降低执行成本。以下各节讨论剩余的工程问题:如何比较数值路径并降低其执行成本。我们分别评估这两条路线;将它们结合所带来的额外收益仍有待测试。
4. miniTE:对齐数值执行并优化低精度计算
4.1 控制数值差异与执行成本
NVFP4 实验反复需要局部改动:激活缩放粒度、为选定层使用更高精度、反向格式,以及量化或 GEMM 实现。当这些选择与模型模块、分布式通信和训练脚本紧密耦合时,每次实验还必须重新审视周边逻辑,使错误更难定位。
miniTE 在执行层解决这些问题,使量化规则显式化,并在匹配的数值边界上比较学习器和采样器。它提供普通线性层、LayerNorm–线性组合模块以及 MoE GroupedLinear 模块。统一的量化接口连接 Python 参考实现与优化后端。Megatron 和上层 RL 系统仍负责编排训练,而 SGLang 等推理框架执行采样 [7,8]。miniTE 专注于集成、比较和优化数值策略与算子。
图 3. miniTE 框架与 MoE 计算路径。右侧面板展开 GroupedLinear 的逐 token NVFP4 前向路径和三种反向选择,并图示了稠密打包。节点说明遵循代码中关于张量形状、缩放、布局、amax 复用、GEMM 调度和并行约束的描述。学习器–采样器检查在匹配的权重版本、缩放域、舍入和组合下比较编码与缩放、专家输出以及 token 对数概率。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
4.2 分层接口与 Transformer Engine 集成
模块层保留训练框架所需的参数、自动微分和分布式接口。量化策略层指定格式和缩放。计算后端实现量化和 GEMM。开发者可以先用参考实现验证数值行为,再将 CUDA、CUTLASS 或 TE 实现集成到相同接口之后。
| 层 | 组件 | 扩展点 |
|---|---|---|
| 模块 | Linear, LayerNormLinear, GroupedLinear | 模型集成与分布式封装 |
| 精度配置 | RecipeSpec, get_qlinear_params, QLinearParams | 命名选择;激活、权重、梯度以及 fprop/dgrad/wgrad 精度 |
| 量化接口 | quantize, group_quantize | 量化参数、缩放粒度、填充、分组 |
| 矩阵乘法 | qgemm 与后端 | 编码数据、缩放布局、GEMM 调度 |
| 诊断 | 张量统计、性能分析、跨框架收集与比较 | 数值回归与性能分析 |
表 4. miniTE 职责。新的数值策略主要在量化和后端层实现,复用共享的模块逻辑。
TE 集成是核心设计要求。miniTE 遵循现有 Linear 接口和 checkpoint 参数名,使 Megatron 封装可以替换相应模块,同时复用周边训练工作流。参数加载、额外状态和并行模式在每条集成路径上均经过验证,以平衡迁移成本与行为一致性。
配方注册表存储名称、用途和依赖元数据;参数工厂将选择转换为 QLinearParams。注册表和检查入口仅使用 Python 标准库。算子导入被延迟到相应的工厂分支,因此开发者可以在准备其 CUDA、TE 和编译依赖之前检查配置。参考路径和优化路径继续共享上层模块。
4.3 精度配方与选择性量化
第 2 节中的 W4A16、W4A8 和 W4A4 实验需要独立配置权重和激活。miniTE 还分别为前向传播(fprop)、输入梯度(dgrad)和权重梯度(wgrad)分离梯度精度和 GEMM 设置。因此,新的前向量化策略无需绑定特定的反向精度。
1.0.0 版通过 MINITE_RECIPE 或 get_qlinear_params选择命名配方,包括 nvfp4_per_token 和 nvfp4_per_token_bf16_bwd。显式层 qlinear_params 优先于环境选择器。层名覆盖和按层号排除量化支持敏感性分析。开发可以从参考量化开始:建立缩放和编码,选择 GEMM 后端,然后比较全精度和量化反向路径。
miniTE CUTLASS 和 miniTE per-token 是 MoE GroupedLinear 的两种实现。两者都使用 NVFP4 前向计算,并按 token 计算激活外层缩放。CUTLASS 命名的是矩阵乘法实现,而 per-token 命名的是缩放粒度;这两个术语并不互斥,TE 的 per-token 路径也会调用 CUTLASS 内核。主要差异如下。
| 属性 | miniTE CUTLASS | miniTE per-token |
|---|---|---|
| 激活外层缩放 | 每个 token 一个 | 每个 token 一个 |
| 权重外层缩放 | 每个专家权重张量一个 | 展平后的专家权重每行一个 |
| 量化与 GEMM 接口 | miniTE 打包与 CUTLASS V2 分组 GEMM | TE 逐 token 批量量化与分组 GEMM,支持列表/稠密打包 |
| 配方 | nvfp4_cursor2_cutlass_v2 | nvfp4_per_token |
两条路径均使用 MXFP8 反向和 BF16 反向进行了测试。反向精度单独配置,不能从路径名称推断。所报告的 BF16 反向性能使用层旁路,该旁路保存原始 BF16 操作数并跳过反向量化。专用的 nvfp4_per_token_bf16_bwd 算子是单独的实现,需要单独的性能测量。
该图展示了为两条 miniTE 路径测量的前向和两种反向选择。BF16 反向使用保存的原始 BF16 操作数;更改反向精度不会改变 NVFP4 前向路径。静态版本展示 BF16 反向。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
公共选择器通过一个注册表解析配方,按通用、参考和实验用途分组。这些分组描述用途。 describe 报告前向/反向行为和请求的布局; doctor 检查选定的运行时先决条件。每个结果记录前向和反向精度、量化覆盖范围以及源修订版本。
4.4 数值监控、性能分析与对齐诊断
低精度失效可能在训练后期出现。因此,miniTE 提供可配置的数值与性能采集。开发者可以按模块、实例、rank 和 step 检查激活、权重、输出和梯度的最小值、最大值、均值和范数,并保存张量以复现异常。性能分析器生成 Chrome trace 和内核计时统计,将计算、准备和调用之间的间隙区分开来。
比较学习器和采样器首先需要对齐它们的输入。MoE dispatch 可能改变两侧的 token 顺序;此时直接逐元素相减会把排列差异误认为数值误差。我们收集路由信息、专家 FC1/FC2 张量和 token 对数概率,对齐 token 与专家身份,然后比较输出、量化数据和缩放。
诊断区分实际内核输出、hook 边界和参考重放。当融合内核内部的中间量无法直接采集时,重放可以定位问题,但其来源仍被明确标注。重放中的一致与实际执行中的一致并不相同。此后,每项 recipe 或内核改动都可以沿输入与路由、量化、GEMM 和输出概率进行追踪。
4.5 训练–推理对齐与确定性执行(确定性模式)
低比特训练的稳定性取决于量化精度,以及算子在训练和推理中行为是否一致。即使采用相同的量化策略,算子实现、归约顺序或并行执行上的细微差异也可能累积成影响训练的数值偏差。
因此,我们在 miniTransformer 中实现了 确定性模式 :逐层替换 Megatron 算子,使其使用与 SGLang 推理相同的内核、公式和归约顺序。多个经过校验的关键层实现了 逐位精确 的一致。
这种逐层对齐与 4.4 节中的张量和路由身份校验配合使用。在原始记录中,严格配置下 TIS 被禁用,非严格参考下则启用。因此,该图保留了执行路径与算法校正的联合身份;逐位精确的层间一致性和端到端对数概率差距分别报告。
训练动态:奖励、差距和熵按训练记录在 step 0–1199 上展示,端点分别为 871/1199/1189。默认使用居中的 15 点均值。策略梯度损失见下方四面板图。严格对齐保留了原始的联合干预配置。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
奖励、差距、熵和策略梯度损失
图 4. 训练–采样对齐:奖励、平均绝对对数概率差距、熵和策略梯度损失。粗线为 15 步移动平均。对齐配置在测量区间内达到约 10−810^{-8}10−8 的差距,并在后期保持相对稳定的奖励。该比较中多项实现选择同时发生变化。
在本实验中,对齐后的训练–推理对数概率差异降至约 10−810^{-8}10−8,后期奖励稳定,扰动更小。未对齐的 NVFP4 则出现后期奖励下降,同时熵和对数概率差异上升。在该配置下,更小的训练–推理差异伴随着更稳定的 NVFP4 RL 动态。这是对不同实现配置的单次运行比较,多项选择同时发生变化。该严格对齐实验与第 6 节中的实际学习和吞吐实验相互独立。
miniTE 在提供对齐能力的同时,也提供了实现这些能力所需的配置、调试和扩展接口。新的低精度策略可以遵循相同的验证流程,从局部正确性一直到 RL 训练。
对齐 · §4.5
在严格对齐配置下,训练–推理对数概率差距保持在约 10⁻⁸,后期奖励稳定。
5. 基于 Transformer Engine 的 NVFP4 联合实现与优化
5.1 逐 token 缩放的数值与实现选择
我们围绕 TE 的 NVFP4 支持与 NVIDIA 开展联合测试和实现工作,将逐 token 配方及其接口作为重要的集成点 [9]。缩放粒度既影响数值行为,也影响 amax 计算、缩放存储和 GEMM epilogue 开销。因此,精度和速度需要一起评估。
NVFP4 以 E2M1 编码数据,对每 16 个元素组成的块使用一个 E4M3 缩放,并通过外层 FP32 缩放调整范围 [1]。设 sis_isi 为第 iii 行的外层缩放,bi,jb_{i,j}bi,j 为其长度为 16 的块。则
di,j=QE4M3 (amax(bi,j)6si),x^=sidi,jQE2M1 (xsidi,j).d_{i,j}=Q_{\mathrm{E4M3}}\!\left(\frac{\operatorname{amax}(b_{i,j})}{6s_i}\right), \qquad \hat x=s_i d_{i,j}Q_{\mathrm{E2M1}}\!\left(\frac{x}{s_i d_{i,j}}\right).di,j=QE4M3(6siamax(bi,j)),x^=sidi,jQE2M1(sidi,jx).逐张量缩放对整个张量共用一个外层缩放;逐 token 缩放则按行计算。后者降低了对其他 token 最大值的依赖,并可在各行幅值差异很大时缓解块缩放下溢。它仍保留行内动态范围和 FP4 舍入误差,并且需要更细粒度的计算支持。公开系统还将逐 token NVFP4 前向计算与更高精度的反向计算相结合 [2].
我们保留多种配方,并将数值选择与实现优化分开评估。仅降低局部量化误差并不能证明 RL 训练更稳定。额外归约的开销同样要在完整计算路径上评估。
D2 · 20k 步语言模型损失:逐 token 与逐张量缩放对比
实验结果
图 5-1. Ling MoE,单 GPU,种子 11,20,000 步。两种 NVFP4 配置均使用 RHT、SR 和 2D 权重量化。左图显示绝对损失;右图减去同一数据上的 BF16 损失。最后 1k 步的相对差距为逐张量 +1.317%、逐 token +1.060%。
在绘图尺度上绝对损失很接近;从 BF16 中减去后可以揭示更小的差异。这为缩放选择提供了长程语言模型损失证据。它与第 6.2.2 节中的 RL+TIS 轨迹不同,也不能独立证明 RL 稳定性。
5.2 连续布局与分组 GEMM 调用
原始分组接口可能将量化数据和缩放拆分到不同的专家张量中,然后为 GEMM 重建指针、步长和布局。随着专家数量增加,Python/ATen 切片、列表构建、临时分配以及主机到设备的元数据传输可能抵消 FP4 算术增益。
我们将量化输出组织到连续缓冲区中,并用偏移量描述专家边界,直接传递给稠密分组 GEMM。FP4 数据、块缩放和外层缩放通过同一接口传递,避免在量化后拆分、在计算前重新组装。设备端元数据构建和工作区复用减少了分配和同步拷贝。零 token 专家和填充在同一分组路径中处理。
此优化改变了数据组织和调用方式。稠密量化和 GEMM 接口针对操作数和输出与原始列表路径进行了核对;性能针对组合实现进行测量。修订版 b641179 后续将稠密路径、GEMM 设置和融合相结合。核心重复测试保留其原始修订标识。
5.3 在数据产生处计算 amax
单独的 amax 内核首先扫描激活,然后量化器再次读取它们。我们寻找在生产到消费路径上已经遍历数据的操作。
对于 MoE FC1,输入经过分发和置换,允许在 token 重排序期间累积逐行 amax。填充同时处理激活和 amax,因此量化可以复用外部统计量并跳过重复归约。
对于 MoE FC2,amax 更适合与激活计算融合。我们使用 _weighted_swiglu_row_amax_kernel 作为融合内核,遵循 JIT 编译实现并在遍历列时累积 amax。该统计量适用于路由加权后的 SwiGLU 输出,正是传递给 FC2 量化的张量。复用现有内存和计算路径避免了对 FC2 输入进行另一次完整扫描以计算 amax。
FC1 变换为:
Original: sort → pad(activation) → amax → quantization
Fused: sort + amax → pad(activation, amax) → quantization
FC2 在生成加权激活的同时产生 amax,并在后续量化中复用它:
Original: SwiGLU × router weight → activation → amax → quantization
Fused: weighted SwiGLU + row-amax → activation, amax → quantization
收益必须覆盖生产者、填充和消费者。跳过 amax 的模块级模拟并不总是更快,因此消除一个内核本身并不能证明完整流程的收益。第 6 节使用包含周边工作的学习器计时。
5.4 内核优化及其适用范围
内核级优化涉及跨 CTA 的 amax 归约、FC1/FC2 GEMM 流水线、缩放布局传输以及 epilogue 中的激活/amax 融合。实现细节和适用范围如下。
展开:内核级实现细节1. 行 amax:用于跨 CTA 归约的集群共享内存。 逐 token NVFP4 的 K1 阶段对每一行跨 K 瓦片归约 amax。原始路径使用 atomicMax合并瓦片结果。在优化路径中,CTA 首先计算部分 amax 并写入领导 CTA 的共享缓冲区。领导者在集群内进行归约。覆盖所有 K 瓦片的集群直接产生最终行 amax;更大的范围回退到原子操作。在记录的测试中, n_cluster=16 和 k_tile=128 覆盖 K=2048,而将瓦片增加到 256 则覆盖 K=4096。这仅替换 K1 amax;K2 转换保持独立。多集群方案未显示收益,已被移除。该优化适用于单个集群覆盖全部或大部分 K 的情况,不将其收益外推到更大的 K。
2. FC1/FC2 GEMM:根据归约长度选择瓦片和流水线。 两层都使用分组 NVFP4 GEMM,但它们的 K 长度和 epilogue 重叠条件不同。FC1 具有更长的 K 和 MMA 循环。 MMA_N=256 与 2SM 增加了每个 CTA 的工作量,并为 epilogue 工作提供更多与 MMA 重叠的机会。FC2 的 K 较短;复用 256 列与 OverlappingAccum 可能使 epilogue 读取与下一阶段的 MMA 写入竞争。因此它使用 MMA_N=128 和 2SM。TMEM 有 512 列。两阶段 256 列瓦片填满该容量,没有空间留给缩放,需要采用 OverlappingAccum 策略而非默认 TMEM 分配。这些是针对所测试 FC1/FC2 形状的配置;最优选择随形状而变化。
3. 融合 swizzle:在转换期间直接写入 GEMM 缩放布局。 下游 GEMM 需要特定的 swizzled 缩放因子布局。原始实现生成紧凑的 M-major 缩放,然后启动 nvte_swizzle_scaling_factors,增加一次内核启动和全局内存流量。融合实现向 with_swizzle 添加 launch_grouped_fused_cast_bf16路径,在 DO_ROW中直接写入 cuBLASLt 128M × 4K 瓦片布局。禁用时,紧凑布局和单独转换仍然可用。这消除了中间缩放重排序,同时保留量化公式。后续 miniTE 逐 token 测试发现,融合路径与非融合路径在量化载荷、恢复到相同布局的缩放以及最终输出上逐位一致。
4. 激活与 epilogue 融合:在 FC1 输出处计算门控和统计量。 MoE FC1 产生成对的 gate/up 张量,通常后接 SwiGLU 或 GeGLU 以及用于后续量化的输出侧 amax。融合 epilogue 执行 alpha/bias 校正、成对 gate/up 激活、路由概率加权和 D 侧 amax。根据配置,它保留可选的 pre-GLU 中间 C 写入;FP8 路径还可以生成输出缩放因子。核心数据流为:
ACC = block-scaled GEMM(A, B)
ACC = ACC × α_a[m] (optionally × α_b[n]) + bias
C = gate | up (pre-GLU, width N; optionally written)
D = act(gate, up) × prob (SwiGLU/GeGLU, width N/2)
Reduce D amax concurrently; FP8 can generate output scale factors
在原地计算激活和统计量,可减少对完整 gate/up 中间结果的写入及后续读取,同时避免单独的激活/amax 内核启动。该融合目前针对 FC1 直接馈入门控激活和量化统计的前向路径。在主性能配置中的纳入情况按每项优化分别记录;完整收益仍基于相应模块和学习器计时。
5.5 联合测试如何指导实现选择
联合工作遵循参考实现、优化实现、模块比较和模型测试的流程。在固定配方下的比较验证接口和布局变更;跨配方的比较评估精度选择及其代价。BF16、原生 TE NVFP4、我们的 CUTLASS 路径以及上游逐 token 实现承担不同的基线角色。
该过程将 TE 能力集成到 miniTE,并将 miniTE 中发现的瓶颈反馈到 TE 接口和内核。公开的逐 token 实现与讨论见 PR #3045 [9]。修订、补丁和测试范围随结果一并提供;归属和上游合并状态以实际记录为准。
5.6 探索单遍融合
我们还研究了融合 SwiGLU-FFN 路径的在线 amax 补偿,以避免物化完整的中间激活。FC2 在整行 amax 已知之前就消费激活 K 瓦片。该方法维护一个运行中的逐行 amax,并在外层缩放增大时,将先前的部分累加转换为新缩放,从而连接 FC1、SwiGLU、量化和 FC2。
该路径主要减少中间激活的内存流量和单独的内核调用。补偿在累加瓦片贡献时保持一致的缩放;先前生成的 E4M3 块缩放和 E2M1 编码仍取决于当时使用的缩放和舍入。因此我们将其保留为融合研究方向,分别评估补偿误差、与两遍执行相比的量化差异以及完整调用链收益。该原型与第 6 节中的主 B300 学习器配置分开评估。
展开在线 Amax 补偿量化算法:在当前行缩放下累加 FC2 部分输出
令 PiP_iPi 表示 FC2 K 瓦片 iii 的 SwiGLU 激活。完成相应 gate/up 切片的 FC1 归约,然后计算 SwiGLU、逐行绝对最大值以及每个 16 元素块的最大值。运行中的逐行 amax 和外层缩放为
gi=max(gi−1,rowmax(∣Pi∣)),gsi=gi2688,2688=6×448.g_i=\max\bigl(g_{i-1},\operatorname{rowmax}(|P_i|)\bigr),\qquad gs_i=\frac{g_i}{2688},\qquad 2688=6\times448.gi=max(gi−1,rowmax(∣Pi∣)),gsi=2688gi,2688=6×448.对于块 bbb,量化器产生 bsi,b=QE4M3(amax(Pi,b)/(6gsi))bs_{i,b}=Q_{\mathrm{E4M3}}\bigl(\operatorname{amax}(P_{i,b})/(6gs_i)\bigr)bsi,b=QE4M3(amax(Pi,b)/(6gsi)) 和 qi,b=QE2M1(Pi,b/(bsi,bgsi))q_{i,b}=Q_{\mathrm{E2M1}}\bigl(P_{i,b}/(bs_{i,b}gs_i)\bigr)qi,b=QE2M1(Pi,b/(bsi,bgsi))。FC2 块缩放 MMA 计算
ΔOi=(qi⊙bsi),W~2,i,Oi=Oi−1gsi−1gsi+ΔOi,Y=Ok,gsk,gsW2.\Delta O_i=(q_i\odot bs_i),\widetilde W_{2,i},\qquad O_i=O_{i-1}\frac{gs_{i-1}}{gs_i}+\Delta O_i, \qquad Y=O_k,gs_k,gs_{W2}.ΔOi=(qi⊙bsi),W2,i,Oi=Oi−1gsigsi−1+ΔOi,Y=Ok,gsk,gsW2.此处,W~2,i\widetilde W_{2,i}W2,i 包含权重块缩放,但不包含沿 K 恒定的外层权重缩放 gsW2gs_{W2}gsW2。逐行激活缩放和外层权重缩放在其对应的输出维度上广播。如果外层权重缩放随瓦片变化,则将该变化纳入每个部分结果,而不是在最后应用一个公共因子。激活 bs 在 MMA 内部逐块应用;只有外层 gs 推迟到 epilogue。
数值条件:避免饱和并不意味着对每个 gs 都不变
对于固定数据,gs≥amax(P)/2688gs\ge\operatorname{amax}(P)/2688gs≥amax(P)/2688 使未舍入的块缩放保持在 E4M3 最大值 448 以内。在线运行 amax 覆盖当前瓦片最大值,因此当前瓦片满足该条件,不会发生上界饱和。
在没有 E4M3 舍入的情况下,bs⋅gs=amax(Pb)/6bs\cdot gs=\operatorname{amax}(P_b)/6bs⋅gs=amax(Pb)/6,允许块缩放抵消外层缩放变化。在实践中,舍入后的 bs 会在该乘积中留下依赖缩放的舍入残差;较大的 gs 还可能将小块推入次正规或下溢范围。当码本保持在正规范围内、避开其边界并使用一致的舍入时,二次幂重缩放可以保持解码步长。任意缩放比例不具备此保证。
原始重尾示例使用一个 amax 约为 29 的激活行。设计缩放及其四倍均得到 rel-L2 0.099;将缩放降至设计值的四分之一会使 rel-L2 升至 0.52:
| 外层缩放 | rel-L2 | 块缩放上界行为 |
|---|---|---|
| amax / 2688 | 0.099 | 1/16 恰好达到 448,未超过 |
| 4 × amax / 2688 | 0.099 | 没有块达到上限 |
| amax / (4 × 2688) | 0.52 | 6/16 超过上限并被裁剪 |
该示例说明了饱和条件的作用以及在指定缩放倍数下误差相等。它并不确立对于下界以上的每个缩放都具有相同的编码或误差。
补偿:在共同的外层缩放中合并部分结果
ΔOi\Delta O_iΔOi 尚未乘以其瓦片的 gsigs_igsi,因此来自不同缩放瓦片的部分结果不能直接相加。递推将旧累加器转换为当前缩放。在理想实数运算中展开可得
Oi=∑j≤iΔOjgsjgsi,OkgskgsW2=∑jΔOjgsjgsW2.O_i=\sum_{j\le i}\Delta O_j\frac{gs_j}{gs_i},\qquad O_k gs_k gs_{W2}=\sum_j\Delta O_j gs_j gs_{W2}.Oi=j≤i∑ΔOjgsigsj,OkgskgsW2=j∑ΔOjgsjgsW2.该恒等式恢复每个已量化瓦片的外层缩放。它不会重新计算先前的 bs 或 q因此,它并不能确立与已知整行 amax 的两遍量化之间的按位等价性。FP32 的比值、乘法和累加同样会进行舍入,必须在数值比较中加以检查。
原始诊断使用两个 amax 分别为 1 和 100 的切片以及一个 FP32 参考值。直接相加缩放比例不同的部分结果会得到 rel-L2 2618%。在求和前恢复每个切片的缩放,以及使用顺序比值补偿,这两种方法都记录到 rel-L2 为 6.0%,处于该示例的量化误差范围内。
| 合并方法 | rel-L2(%) | 解释 |
|---|---|---|
| 直接相加两个部分结果 | 2618 | 混合了以不同外层缩放表示的输出 |
| 先恢复每个 gs,再相加 | 6.0 | 恢复每个瓦片的实际贡献 |
| 运行缩放比值补偿,再恢复 gs | 6.0 | 递归执行相同的缩放转换 |
实现流程与开销
融合原型通过 CuTeDSL 面向 SM100(Blackwell)。FC1 在 TMEM 中累加对应的 gate/up 切片,然后将其传输到寄存器以进行 SwiGLU、绝对行最大值和量化,运行中的 amax 存储在工作区中。FC2 直接使用 E2M1 数据和 E4M3 块缩放。其尾声从 TMEM 中取回部分结果,并在寄存器或 SMEM 中对其进行补偿和合并。正常数值范围内的核心流程为:
Inputs: X, W_FC1, W_FC2
Initialize gamax = 0, gs_prev = 0, O = 0
for each FC2 K-tile i:
Complete FC1 reduction for gate/up slice → SwiGLU → P_i
rowmax = max(|P_i|); gamax = max(gamax, rowmax)
If gamax = 0: this row contributes zero; skip computation
gs = gamax / 2688
Generate E4M3 bs and E2M1 q per 16-element block
FC2 block-scaled MMA → ΔO_i (includes block scales, excludes gs)
O = O × (gs_prev / gs) + ΔO_i
gs_prev = gs
Y = O × gs_prev × gs_W2; all-zero rows output zero
SwiGLU 输出可能为负,因此行最大值必须使用 max(|P|)。对于全零行或块,显式设置 q=0 以避免 0/0。如果非零块的缩放发生下溢,则应用量化器的零值规则并记录相关误差。
该路径避免了写入和读取完整的 FP32 中间激活,但 O 必须保留在 TMEM 之外,从而为每个 K 瓦片增加部分结果传输和合并。在所提供的 B200 Qwen3-30B-A3B 形状测试记录中,该原型比未融合的 TE 两级链快 4–8 倍,精度接近该路径。与融合的单级路径相比,在 512 token 的小批量下合并开销更为明显,而在 4096 token 时差距约为 5%。因此,预填充和更大的批量是值得进一步评估的工作负载。该原型形状测试与第 6 节中完整的 B300 学习器测试分开测量。
6. 评估:数值正确性、学习与执行效率
我们通过检查量化、接口、矩阵乘法和梯度来验证前述实现选择,然后评估长期学习。接下来我们测量完整的学习器更新、模块、服务和在线 RL 阶段。每个实验报告其结果以及它如何为实现选择提供依据。
正文呈现关键趋势和结果。完整的形状、梯度、重复运行和配置数据随相应章节提供。网络阅读器在其标题下提供这些内容;打印版将其收集在附录 C 中。
6.1 设置与测量范围
本节评估 miniTE 系统性能;第 3 节报告 TRIAGE 的算法结果。主要性能模型为 Qwen3-30B-A3B,在单节点 8×B300 上运行,TP/PP/CP/EP = 1/1/1/4。学习器测试使用序列长度 4096、微批量大小 1 和 GBS 128,从相同的检查点、种子和冻结的 DAPO 文本出发,每次更新处理 524,288 个 token。模块、学习器、服务和在线 RL 测试具有不同的计时范围。
| 实验 | 问题 | 测量范围 |
|---|---|---|
| 量化 / API / GEMM / 梯度 | 优化是否会改变数值?误差从何而来? | 参考内核或 FP32 黄金标准;匹配的输入 |
| 长期训练和任务评估 | 该配置能否维持有效的学习? | 训练轨迹和检查点评估 |
| 模块前向 / 直接 F+B | 局部收益和反向成本在哪里? | 量化和专家 GEMM;调度不计入 |
| 固定文本学习器更新 | 局部优化能否改善完整更新? | 前向、反向、通信、优化器、数据加载 |
| 固定输出服务 | 采样容量提升了多少? | 八引擎输出吞吐量;学习器和同步不计入 |
| 在线 RL 阶段 | 哪些系统成本抵消了计算收益? | Logprob + actor 训练以及转换/同步分别测量 |
| 配置 | 前向路径 | 反向 / 覆盖范围 | 作用 |
|---|---|---|---|
| TE BF16 | TE BF16 GroupedLinear | BF16 | 高精度参考和基线 |
| TE NVFP4 | 原生 TE 逐张量 NVFP4 | 测试的原生配方;非专家层为 BF16 | 现有 NVFP4 路径 |
| miniTE CUTLASS | 自定义 CUTLASS NVFP4 | MXFP8 和 BF16 反向分别测试 | 自定义路径和精度选择 |
| miniTE 逐 token | 基于 PR #3045 的逐 token NVFP4 | MXFP8 和 BF16 反向;仅专家 | 逐 token 路径和优化链 |
| TE 逐 token | 固定的上游逐 token 实现 | 相同的模块前向配方;全模型覆盖/反向另行记录 | 模块实现和配置计时参考 |
表 6-1. 比较对象。两条 miniTE 路径均使用 MXFP8 和 BF16 反向进行测试;主要学习器配置使用 BF16 反向。
- 模块测试使用五个独立进程,对每次运行取稳态块内的中位数。4k/16k/64k/128k 大小表示每个 rank 在调度前的 token 数。
- 核心学习器测试使用十次独立的 20 次更新运行,前五次作为预热,测量第 6–15 次更新。吞吐量和配对加速比各自在取平均前丢弃其最小值和最大值。
- 完整优化链每个配置使用一个进程和 17 次更新。在两次预热更新后,移除 15 次测量更新时间的最大值和最小值;吞吐量由剩余 13 次时间的平均值计算得出。这些更新属于同一个进程。
- 核心修订版本为 miniTE 953f17f9;优化运行使用 b641179,包含密集路径、GEMM 类型和融合。TE 的零 token 专家更新以及运行 3 中的干扰在附录 A.3 中描述。
6.2 正确性与学习行为
6.2.1 量化、API、GEMM 和梯度的独立检查
在完整训练之前,我们先验证局部数值边界。API 与布局优化在固定配方下比较编码、缩放和输出。量化计算在反量化后与 FP32 参考进行比较。反向测试在直通估计器(STE)语义下,对照 FP32 梯度检查方向、幅度和符号。这样可将格式误差与优化引入的误差区分开来。
miniTE 逐 token 密集量化和分组 GEMM 通过 28/28 位级等价检查。融合布局在七个合成和真实激活用例中未增加任何差异。对于所测试的配置,这些布局变更保持了量化结果。
D4 · 八条算子路径及其参考| 类别 | 算子 / 路径 | 测试目标 | 参考 / 判据 |
|---|---|---|---|
| miniTE CUTLASS | group_quantize_nvfp4_pertoken_swizzled | 逐 token NVFP4 激活量化 | 与逐组参考位级一致 ;反量化遵循 FP4/E4M3 缩放语义 |
| miniTE CUTLASS | group_quantize_nvfp4_tensor_global_packed | 逐张量/全局 NVFP4 权重打包 | FP4 数据与 swizzle 后的缩放与逐专家参考一致 quantize_weight_nvfp4_swizzled |
| miniTE CUTLASS | grouped_cutlass_gemm_v2 / v2_evt | FP4 分组 GEMM | 反量化后的 FP32 GEMM 黄金参考 |
| CUTLASS 端到端 | miniTE CUTLASS GroupedLinear | 模块前向与反向 | 纯 BF16 分组矩阵乘 / FP32 黄金梯度 |
| miniTE 逐 token | nvfp4_per_token_group_quantize_bulk_dense | TE 逐 token 分组量化密集 API | PR#3045 列表 API; 位级一致 |
| miniTE 逐 token | nvfp4_cutlass_grouped_per_token_gemm_dense | TE 逐 token 分组 GEMM 密集 API | PR#3045 列表 API; 位级一致 |
| miniTE 逐 token | 融合 swizzle / 非融合 swizzle | 缩放布局融合是否引入数值误差? | 融合与非融合端到端 位级一致 |
| 两条 miniTE 反向路径 | dgrad / wgrad | 反向梯度正确性 | FP32 黄金梯度;在 STE 语义下预期余弦 > 0.99 |
表 6-2. 正确性检查中使用的对象与参考。
D5 · 极端输入、次正规数与缩放下溢| 测试 | 结果 | 备注 |
|---|---|---|
| CUTLASS 激活量化与权重打包 | 50 项通过 | 包含 16 个新增极端用例;8 个基准用例被取消选择 |
| 逐 token 密集量化:极端输入 | 8 项通过 | miniTE 逐 token 密集量化的极端输入 |
| CUTLASS 权重布局参数回归 | 8 项通过 | 现有位级一致回归未受影响 |
表 6-3. 极端输入的工程回归测试。通过测试数描述的是用例覆盖范围,而非模型训练稳定性或统计显著性。
- 逐 token 激活量化: FP4 数据与缩放与逐组参考位级一致。极端幅度、接近 BF16 最大值的数值、次正规数以及混合幅度输入均不产生 NaN/Inf。次正规行的全局缩放可能下溢为零。
- 逐张量/全局权重打包: 编码与 swizzle 后的缩放与逐专家参考一致。一个次正规专家保留了约 9e-41 的输入,相对 L2 约为 1.56e-2。全局缩放检查使用 3% 的相对容差,而非对每个浮点缩放都要求位级相等。
- miniTE 逐 token 密集逐 token 量化: q_row、块缩放和 row_amax 与上游列表路径一致。极小的行可能因块解码缩放下溢而变为零。所测试的生产路径要求 m_splits 为 128 的倍数。
这些不同的归零与缩放行为有助于区分格式范围、浮点 FTZ 和接口问题。
FP32 尾声也有直接的数值支持。普通 miniTE CUTLASS 分组 GEMM 和 EVT 路径均在反量化后与 FP32 黄金参考进行比较:
| 内核 | 余弦 | max_abs | norm_rmse |
|---|---|---|---|
grouped_cutlass_gemm_v2 | ≥ 0.999994 | ≤ 0.0156 | ≤ 0.0033 |
grouped_cutlass_gemm_v2_evt | ≥ 0.999999 | ≤ 0.0078 | ≤ 4e-5 |
表 6-4. miniTE CUTLASS GEMM 精度。EVT 将 norm_rmse 上界从约 0.0033 降至 4e-5 量级,支持该路径中的高精度尾声。
与纯 BF16 前向计算相比,GroupedLinear 的余弦相似度约为 0.99095,norm_rmse 约为 0.135,显著大于 GEMM 相对于其自身反量化参考的误差。前者包含因 FP4 丢失的信息;后者检查 GEMM 实现。参考决定了精度指标所衡量的内容。
BF16 反向 dgrad/wgrad 与 BF16 参考高度吻合,最大相对 L2 约为 0.001663。MXFP8 反向的相对 L2 约为 0.03765,符号一致率约为 98.8%。这支持在第 4 节中独立选择反向精度;第 6.4 节测量其执行成本。
D6 · 两条 miniTE 路径的完整梯度结果:余弦、相对 L2、符号一致率与专家指标| 后端 | 反向路径 | dgrad 余弦(最小) | dgrad 相对 L2(最大) | dgrad 符号一致率(最小) | wgrad 余弦(最小) | wgrad 相对 L2(最大) | wgrad 符号一致率(最小) | 专家余弦(最小) |
|---|---|---|---|---|---|---|---|---|
| miniTE CUTLASS | MXFP8 反向量化 | 0.999292 | 0.037626 | 98.802% | 0.999291 | 0.037653 | 98.802% | 0.999287 |
| miniTE 逐 token | MXFP8 反向量化 | 0.999292 | 0.037626 | 98.802% | 0.999291 | 0.037653 | 98.802% | 0.999287 |
| miniTE CUTLASS | BF16 反向 | 0.999999 | 0.001660 | 100.000% | 0.999999 | 0.001663 | 100.000% | 0.999998 |
| miniTE 逐 token | BF16 反向 | 0.999999 | 0.001660 | 100.000% | 0.999999 | 0.001663 | 100.000% | 0.999998 |
| BF16 PyTorch 参考 | BF16 矩阵乘 | 0.999999 | 0.001660 | 100.000% | 0.999999 | 0.001663 | 100.000% | 0.999998 |
表 6-5. 跨种子和形状的最坏情况梯度指标。参考在 FP32 中计算 dX=dY,WBF16⊤dX=dY,W_{\mathrm{BF16}}^\topdX=dY,WBF16⊤ 和 dW=dY⊤XBF16dW=dY^\top X_{\mathrm{BF16}}dW=dY⊤XBF16。100% 的符号一致率并不意味着与 FP32 位级相等。
D7 · 密集 API 与融合/非融合缩放布局| 测试对象 | 形状覆盖 | 结果 |
|---|---|---|
nvfp4_per_token_group_quantize_bulk_dense | 小型、8 专家、32 专家、极端输入 | q / scale / amax 全部 位级一致 |
nvfp4_cutlass_grouped_per_token_gemm_dense | 小型、类 fc1、类 fc2 | GEMM 输出 位级一致, max_abs = 0 |
表 6-6. 密集接口的形状覆盖与结果。
| 用例 | 融合 vs 非融合 相对 L2 | 余弦 | max_abs | FP4 数据不匹配 | 缩放不匹配 |
|---|---|---|---|---|---|
| 合成数据,多种子 / token 数 | 0 | 1.0 | 0 | 0 / 96 组 | 0 / 96 组 |
| 真实转储 | 0 | 1.0 | 0 | 0 / 12 组 | 0 / 12 组 |
表 6-7. 融合/非融合输出与中间结果的差异。生产版 miniTE 逐 token 路径与早期实验路径具有不同的形状约束;本表报告生产路径的检查结果。
6.2.2 补充性长程学习评估
图 5 比较了 BF16+TIS、原生 NVFP4+TIS 和 miniTE+TIS。miniTE 在前 1,000 步中保持奖励增长,并在运行后期超过原生 NVFP4;BF16 在该区间内获得更高的奖励。两种 NVFP4 配置的平均对数概率差距相近,但学习轨迹不同。因此,仅凭平均差距不能替代任务性能。
这些曲线包含 6,000 个差距和奖励观测值。三种配置均使用 TIS,覆盖第 0–999 步。默认保留居中的 21 点均值;请检查原始值或缩放范围。
打开交互式图表:检查数值、切换配置、下载 SVG / CSV
图 5. 实用 miniTE 路径与基线的学习动态。每种配置在第 0–999 步上各运行一次。左图显示平均绝对对数概率差距;右图显示奖励。细线为原始值,粗线为居中的 21 步均值。所有配置均使用 TIS。这是与严格对齐参考相独立的实验。
本补充运行使用 Qwen3-30B-A3B-Base。BF16、原生 TE NVFP4 和 miniTE 均使用 TIS,并禁用路由重放。表 6-8 在各自的第 999 步检查点上评估四项数学任务和 IFEval [11,12,13,14]。每个 AIME 基准平均六次评估;MATH-500、GSM8K 和 IFEval 各平均四次。IFEval 使用提示级严格准确率。BF16 检查点使用 BF16 推理评估,两个 NVFP4 检查点均使用 NVFP4 推理评估。
| 配置 | AIME24 | AIME25 | MATH-500 | GSM8K | IFEval | 五项任务均值 |
|---|---|---|---|---|---|---|
| BF16 + TIS | 41.11% | 31.11% | 90.05% | 93.57% | 38.36% | 58.84% |
| TE NVFP4 + TIS | 26.67% | 20.56% | 86.90% | 94.71% | 43.99% | 54.57% |
| miniTE + TIS | 36.67% | 24.44% | 87.45% | 94.79% | 38.40% | 56.35% |
| miniTE 减 TE NVFP4 | +10.00 | +3.88 | +0.55 | +0.08 | −5.59 | +1.78 |
表 6-8. 补充运行第 999 步的任务得分。最后一行报告百分点差异。五项任务权重相等,与评估次数无关。对检查点的重复评估不会产生独立的训练种子。
miniTE 在全部四项数学得分上均优于原生 NVFP4。其五项任务平均值提高 1.78 个百分点,但仍比 BF16 低 2.49 个百分点。IFEval 的排序不同:miniTE 接近 BF16,但低于原生 NVFP4。这些结果支持在所测量的数学训练任务上取得改进,下游收益因任务而异。
本节补充框架的长程训练评估。其任务集、训练时长和评估协议与第 3 节中的 TRIAGE 不同;各自描述其自身配置下的学习情况。
6.3 完整学习器更新:局部改动能否降低更新成本?
表 6-9 汇总了十次运行的主测试,比较 BF16、原生 TE NVFP4 和 miniTE。两条 miniTE 路径的结果还展示了反向精度如何影响完整学习器更新的吞吐量。
| 配置 | Tokens/s,十次运行截尾均值 | 相对 BF16 的配对加速比 |
|---|---|---|
| TE BF16 | 24,192 | 1.000× |
| TE NVFP4 | 23,294 | 0.959× |
| miniTE CUTLASS,MXFP8 反向 | 27,607 | 1.148× |
| miniTE CUTLASS,BF16 反向 | 33,140 | 1.369× |
| miniTE 逐 token,BF16 反向 | 32,517 | 1.334× |
表 6-9. 完整学习器吞吐量。加速比通过先配对运行再截尾计算。计时包括前向、反向、通信、优化器和数据加载;不包括 rollout 和权重同步。
该图展示了十次运行实验和单独的单进程优化实验。主测试显示所有独立运行值和截尾均值。可切换到配对加速比或单进程优化运行。各协议保持独立;不添加置信区间。静态版本显示主测试吞吐量。
打开交互式图表:检查数值、切换配置、下载 SVG / CSV
两条采用 BF16 反向的 miniTE 路径均达到约 1.33–1.37× BF16 速度,而 TE NVFP4 接近 BF16。将 miniTE CUTLASS 反向从 MXFP8 改为 BF16 进一步提高了吞吐量,表明低位 GEMM 的收益可能被反向量化和准备工作所抵消。miniTE 逐 token 的聚合吞吐量均值与 TE NVFP4 之比约为 1.396×。1.344× 的 miniTE 逐 token/TE BF16 值同样是均值之比,与此处 1.334× 的配对估计量不同。
D8 · 十次独立运行、截尾统计与运行差异| 运行 | TE BF16 | TE NVFP4 | miniTE CUTLASS MXFP8 反向 | miniTE CUTLASS BF16 反向 | miniTE 逐 token BF16 反向 |
|---|---|---|---|---|---|
| 0 | 25545.3 | 25863.2 | 30070.0 | 31196.6 | 33463.5 |
| 1 | 24530.5 | 23653.1 | 27773.2 | 34662.4 | 31416.6 |
| 2 | 23384.9 | 23291.4 | 27501.7 | 33288.5 | 34428.9 |
| 3 | 23802.7 | 20997.8 | 24865.7 | 29218.7 | 28959.3 |
| 4 | 23437.2 | 23473.3 | 26551.5 | 34777.1 | 32164.2 |
| 5 | 25292.7 | 21553.5 | 27759.3 | 31109.0 | 32828.9 |
| 6 | 24352.1 | 22381.7 | 26508.3 | 34214.7 | 30948.3 |
| 7 | 22730.2 | 23138.2 | 25925.8 | 32078.5 | 33315.1 |
| 8 | 23191.5 | 26672.6 | 28769.0 | 33794.5 | 32159.6 |
| 9 | 26277.9 | 22993.5 | 34352.4 | 36279.4 | 33835.6 |
| 截尾均值 | 24192.1 | 23293.5 | 27607.4 | 33140.2 | 32516.5 |
表 6-10. 机器可读结果中每次运行的吞吐量,保留一位小数。正文摘要使用整数。
第 3 次运行期间,一个空闲的服务进程占用了资源。其量化变体的最小观测值被预定义的截尾规则剔除。早期和后期 TE 修订版对零 token 专家的处理方式不同;相应的检查覆盖非空组。配置按运行 ID 配对,而非同时执行。配对加速比的截尾均值与聚合均值之比是不同的估计量。
数据:十次运行记录;运行与修订版。
D9 · 早期三次运行结果与预留内存| 配置 | 早期运行 tokens/s | 峰值预留内存(原始 MB) |
|---|---|---|
| TE BF16 | 25,545 / 24,531 / 23,385 | 119,914 |
| TE NVFP4 | 25,863 / 23,653 / 23,291 | 122,672 |
| miniTE CUTLASS MXFP8 反向 | 30,070 / 27,773 / 27,502 | 144,598 |
| miniTE CUTLASS BF16 反向 | 31,197 / 34,662 / 33,289 | 137,778 |
| miniTE 逐 token MXFP8 反向 | 27,493(仅第 0 次运行) | 138,458 |
| miniTE 逐 token BF16 反向 | 33,464 / 31,417 / 34,429 | 133,398 |
表 6-11. 早期运行与峰值预留内存,保留原始 MB 单位。这些数据与其他协议下以 GiB 为单位的后期测量相互独立。采用 MXFP8 反向的 miniTE 逐 token 有一次参考运行。
miniTE 逐 token 相对于部分 miniTE 变体降低了预留内存,但在本比较中仍高于 BF16。保存的高精度操作数、量化权重缓存、工作区和填充均对总用量有贡献;仅 FP4 载荷宽度并不能决定模型内存消耗。
随后的完整优化链结合了稠密路径、GEMM 类型、置换融合和 amax 复用。采用 BF16 反向传播的 miniTE 逐 token 达到 33,540 tokens/s,即同轮 BF16 基线的 1.39 倍、TE NVFP4 的 1.47 倍。该组合测试使用了不同的修订版和协议:
D10 · 完整融合运行:配置与计时协议| 配置 | 融合运行 tokens/s | 相对同轮 BF16 的加速比 |
|---|---|---|
| TE BF16 | 24,134 | 1.000× |
| TE NVFP4 | 22,801 | 0.945× |
| miniTE CUTLASS,MXFP8 反向 | 25,505 | 1.057× |
| miniTE CUTLASS,BF16 反向 | 33,065 | 1.370× |
| miniTE 逐 token,MXFP8 反向 | 26,516 | 1.099× |
| miniTE 逐 token,BF16 反向 | 33,540 | 1.390× |
表 6-12. 单进程优化运行:每个配置一个进程,吞吐量由 15 次实测更新时间的截尾均值计算得出。所有配置均予保留。miniTE CUTLASS MXFP8 反向达到 25,505 tokens/s,低于单独核心测试中的 27,607。由于协议不同,融合不被视为一致的正面增量。
33,540 这一数值比核心 miniTE 逐 token 结果 32,517 tokens/s 高约 3.1%,但修订版、稠密/GEMM 设置和融合同时发生了变化。该差异描述的是组合运行,无法单独归因于 amax 或 GEMM 类型。此前“gemm_kind 会拖慢速度”的归因在修订版检查后已被更正。
关键结果 · §6.3
两条采用 BF16 反向传播的 miniTE 路径在完整学习器更新上均达到 BF16 速度的约 1.33–1.37 倍。
6.4 模块:前向更快并不一定意味着 F+B 更快
在完整学习器中测得增益后,模块测试解释了其来源。图 6 在 FC1/FC2 之间切换,并显示前向、直接计时的 F+B 或两者。F+B 使用直接测量的计时:
前向和完整的前向加反向(F+B)均直接计时。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
图 6. Qwen3-30B-A3B GroupedLinear 在 B300 上,热静态权重,五次运行。实线表示前向,虚线表示直接测量的 F+B。本图中两条 miniTE 路径均使用量化反向传播。横轴为预分发 token 规模;该均衡工作负载与全模型路由不同。
miniTE 逐 token 前向比 BF16 快约 1.09–1.43 倍,而图中其 F+B 更慢。BF16 反向传播使 miniTE 逐 token 的 F+B 比其自身的 MXFP8 反向路径快约 1.2–1.3 倍。这直接说明在选择反向精度前应测量完整算子成本。
选择 FC1/FC2、前向/直接 F+B、miniTE 反向精度和上游参考。数据点为五次运行的中位数;须线显示最小值和最大值,而非置信区间。连线连接四个实测的预分发 token 规模。默认视图显示两条 miniTE 路径均使用 MXFP8 反向的 FC1 直接 F+B;FC2 可选。图 6 提供前向/F+B 概览。
打开交互式图表:查看数值、切换配置、下载 SVG / CSV
D11 · 四种规模下的 FC1/FC2:完整前向与直接 F+B 延迟前向
| 模块 | 配置 | 测量 | 4k 毫秒 | 16k 毫秒 | 64k 毫秒 | 128k 毫秒 |
|---|---|---|---|---|---|---|
| FC1 | TE BF16 | 前向 | 0.504 | 0.515 | 0.801 | 1.329 |
| FC1 | TE NVFP4 | 前向 | 1.295 | 1.349 | 1.250 | 1.327 |
| FC1 | miniTE CUTLASS | 前向 | 0.336 | 0.396 | 0.895 | 1.508 |
| FC1 | miniTE 逐 token | 前向 | 0.372 | 0.424 | 0.582 | 0.932 |
| FC2 | TE BF16 | 前向 | 0.542 | 0.508 | 0.546 | 0.820 |
| FC2 | TE NVFP4 | 前向 | 1.428 | 1.190 | 1.281 | 1.244 |
| FC2 | miniTE CUTLASS | 前向 | 0.339 | 0.355 | 0.677 | 1.111 |
| FC2 | miniTE 逐 token | 前向 | 0.389 | 0.415 | 0.565 | 0.753 |
采用 MXFP8 反向的直接 F+B
| 模块 | 配置 | 测量 | 4k 毫秒 | 16k 毫秒 | 64k 毫秒 | 128k 毫秒 |
|---|---|---|---|---|---|---|
| FC1 | TE BF16 | fwb | 2.461 | 2.543 | 2.572 | 4.309 |
| FC1 | TE NVFP4 | fwb | 4.174 | 3.986 | 4.011 | 4.583 |
| FC1 | miniTE CUTLASS | fwb | 4.006 | 3.987 | 3.923 | 5.631 |
| FC1 | miniTE 逐 token | fwb | 3.712 | 3.830 | 3.771 | 4.989 |
| FC2 | TE BF16 | fwb | 2.959 | 2.994 | 2.933 | 4.592 |
| FC2 | TE NVFP4 | fwb | 4.604 | 4.790 | 4.802 | 5.414 |
| FC2 | miniTE CUTLASS | fwb | 4.147 | 3.916 | 3.881 | 5.417 |
| FC2 | miniTE 逐 token | fwb | 3.875 | 3.915 | 3.907 | 5.056 |
采用 BF16 反向的直接 F+B
| 模块 | 配置 | 测量 | 4k 毫秒 | 16k 毫秒 | 64k 毫秒 | 128k 毫秒 |
|---|---|---|---|---|---|---|
| FC1 | miniTE CUTLASS | fwb | 2.975 | 2.966 | 3.253 | 5.473 |
| FC1 | miniTE 逐 token | fwb | 2.964 | 2.992 | 3.018 | 4.869 |
| FC2 | miniTE CUTLASS | fwb | 2.908 | 2.913 | 2.958 | 4.936 |
| FC2 | miniTE 逐 token | fwb | 3.158 | 2.943 | 3.088 | 4.636 |
表 6-13. 来自 miniTE 机器可读记录的五次运行中位数,保留三位小数。反向配置单独列出。TE BF16/TE NVFP4 保留在其原始配置表中,因为本 BF16 反向补充没有对应行。
| 配置 | TE BF16 | TE NVFP4 | miniTE CUTLASS | miniTE 逐 token |
|---|---|---|---|---|
| 128k FC1 峰值预留 GiB | 7.39 | 6.91 | 8.65 | 8.76 |
这些内存值是在所述计时和权重缓存设置下的模块测试中测得的。学习器峰值内存单独测量。
原始数据:主模块记录、BF16 反向记录。
D12 · miniTE 逐 token 模块相对 TE NVFP4 的加速比| 模块 | 模式 | 4k | 16k | 64k | 128k |
|---|---|---|---|---|---|
| FC1 | 前向 | 3.48x | 3.18x | 2.15x | 1.42x |
| FC1 | F+B | 1.12x | 1.04x | 1.06x | 0.92x |
| FC2 | 前向 | 3.67x | 2.87x | 2.27x | 1.65x |
| FC2 | F+B | 1.19x | 1.22x | 1.23x | 1.07x |
表 6-14. 对应配置下 miniTE 逐 token/TE NVFP4 的加速比;>1 表示 miniTE 逐 token 更快。F+B 使用各自记录的反向路径。在 128k FC1 下,miniTE 逐 token 并非在每项比较中都领先。
在固定前向配方的情况下,miniTE 逐 token 相对上游 TE 逐 token 在专家前向上取得 2.46× 的几何平均加速,直接 F+B 为 1.10×。这一差异表明前向调用与准备工作是实现收益的主要来源。输入已预先置换,FC2 激活在计时器之外准备,且不包含分发。
D13 · 按工作负载规模的 miniTE 逐 token/TE 逐 token 及 amax 消融| 模块 | 模式 | 工作负载规模 | TE 逐 token | miniTE 逐 token | miniTE 逐 token 加速比 |
|---|---|---|---|---|---|
| FC1 | 前向 | 4k / 16k / 64k / 128k | 1.128 / 1.094 / 1.275 / 1.479 | 0.372 / 0.424 / 0.582 / 0.932 | 3.03x / 2.58x / 2.19x / 1.59x |
| FC1 | F+B | 4k / 16k / 64k / 128k | 3.873 / 3.916 / 3.846 / 5.066 | 3.712 / 3.830 / 3.771 / 4.989 | 1.04x / 1.02x / 1.02x / 1.02x |
| FC2 | 前向 | 4k / 16k / 64k / 128k | 1.193 / 1.043 / 1.555 / 1.747 | 0.389 / 0.415 / 0.565 / 0.753 | 3.06x / 2.52x / 2.75x / 2.32x |
| FC2 | F+B | 4k / 16k / 64k / 128k | 4.603 / 4.507 / 4.488 / 6.234 | 3.875 / 3.915 / 3.907 / 5.056 | 1.19x / 1.15x / 1.15x / 1.23x |
表 6-15. 五次运行的模块延迟与加速比。数值等价性检查固定前向配方;相应记录规定了环境和反向路径。
amax skip-K1 模拟在大多数规模下为中性或略慢,范围约为 0.81–1.11×。它排除了生产者收益并保留了准备开销,因此无法独立证明完整生产者–消费者融合的效果。稠密布局与其他优化一同测量;其单独贡献仍未测量。
| 模块 | 模式 | 工作负载规模 | miniTE 逐 token amax 关闭 | miniTE 逐 token amax 开启 | 开启 vs 关闭 |
|---|---|---|---|---|---|
| fc1 | 前向 | 4k | 0.372 | 0.459 | 0.81x |
| fc1 | 前向 | 16k | 0.4241 | 0.4586 | 0.925x |
| fc1 | 前向 | 64k | 0.5823 | 0.6166 | 0.944x |
| fc1 | 前向 | 128k | 0.9319 | 0.8365 | 1.114x |
| fc1 | fwb | 4k | 3.7122 | 3.8659 | 0.96x |
| fc1 | fwb | 16k | 3.8303 | 3.8542 | 0.994x |
| fc1 | fwb | 64k | 3.7713 | 3.8175 | 0.988x |
| fc1 | fwb | 128k | 4.9888 | 4.9196 | 1.014x |
| fc2 | 前向 | 4k | 0.3894 | 0.4017 | 0.969x |
| fc2 | 前向 | 16k | 0.4145 | 0.4407 | 0.941x |
| fc2 | 前向 | 64k | 0.5655 | 0.5927 | 0.954x |
| fc2 | 前向 | 128k | 0.7534 | 0.7842 | 0.961x |
| fc2 | fwb | 4k | 3.8746 | 3.9556 | 0.98x |
| fc2 | fwb | 16k | 3.9151 | 3.995 | 0.98x |
| fc2 | fwb | 64k | 3.9069 | 3.8958 | 1.003x |
| fc2 | fwb | 128k | 5.0562 | 5.1263 | 0.986x |
表 6-15b. miniTE 逐 token 模块级 skip-K1 模拟。统计遵循模块测试;不包含生产者收益。
数据与运行:模块消融记录。
D14 · 上游 TE 逐 token 全模型计时与配置范围| 配置 | tokens/s(3 次试验) | 相对 TE 逐 token 的配对加速比 | 最大预留(MB) |
|---|---|---|---|
| TE 逐 token | 23125 / 23150 / 22346 | — | 122572 |
| TE BF16 | 25545 / 24531 / 23385 | 1.06x | 119914 |
| TE NVFP4 | 25863 / 23653 / 23291 | 1.04× | 122672 |
| miniTE 逐 token(mxfp8 反向) | 27493(第 0 次运行参考) | 1.19×(第 0 次运行配对) | 138458 |
| miniTE 逐 token(bf16 反向) | 33464 / 31417 / 34429 | 1.45× | 133398 |
表 6-16. 三次运行的全模型对比。约 1.45× 是这些配置之间的计时比值,而非同一配方下仅由实现带来的增益。
TE 逐 token 也会量化非专家层,并使用其原生的逐 token NVFP4 反向路径。miniTE 逐 token 仅针对专家层,并使用 MXFP8/BF16 反向进行测试。TE 逐 token 的运行还记录到约 1e8–1e11 的梯度范数,裁剪后实际更新量非常小。计时与学习行为均予保留;该速度比并不描述等价的学习路径。
参见 TE 逐 token 的全模型记录。
性能剖析解释了从局部到全模型性能的转变:主机元数据准备、分组调用之间的间隙、专家通信以及非 GEMM 操作都会影响更新时间。连续布局降低了组织开销,而模型测试包含真实路由及周边开销。
D15 · 代表性剖析:多张量调用与分组 GEMM调度与路由如何影响下游 GEMM 调度。 在补充记录中第 8 步的第一个微批次里,基于真实权重和输入的路由更为均匀,而基于随机权重和输入的模拟路由则较为集中。两者均使用每 rank 4,096 个输入 token、top-k=8,并在 EP 组内路由 131,072 个 token。我们首先在这一共同 token 预算下比较专家分布:
| 负载指标 | 真实 | 模拟 | 解读 |
|---|---|---|---|
| 有 token 的专家数 | 126/128 | 24/128 | 模拟将 token 集中在更少的专家上 |
| 零 token 专家数 | 2 | 104 | 参与专家数量不同 |
| GroupedLinear M 的 CV | 0.19 | 0.66 | 记录的真实 GEMM 行数更为均匀 |
M 是某个专家 GEMM 的 token 行数。CV 是在相同统计定义下的标准差除以均值;我们保留记录的 GroupedLinear M 统计量,而非从全局专家计数重建。更多参与的专家会产生更多需要准备和调度的矩阵乘法问题。主机 API 与内核剖析有助于解释这一开销:
| 调用指标 | 真实 | 模拟 | 解读 |
|---|---|---|---|
| cuBLASLt 启发式调用次数 / 累计时间 | 148,000 / 6.26 秒 | 32,000 / 1.02 秒 | 真实的调用次数约为 4.6× |
| 每次启发式调用的中位时间 | 2.9 µs | 2.9 µs | 中位数相近,调用次数不同 |
| nvte_multi_tensor_gemm 调用次数 / 中位时间 | 4,608 / 2.29 毫秒 | 4,590 / 0.44 毫秒 | 调用次数相近;真实的中位数约为 5.2× |
| nvjet 内核调用次数 | 114,000 | 4,600 | 真实的调用次数约为 25× |
第一张表确定了该微批次的路由负载;第二张表保留了随附剖析中的调用次数、累计时间和每次调用中位数。累计时间不被视为该单个微批次的墙钟时间,也不直接加到 GPU 内核时间上。
相同的 token 预算被拆分为更多的 GEMM 问题。 对于匹配的专家矩阵形状,主要 GEMM 运算量随专家 token 行数之和缩放。真实与模拟的 token 总数相同,但真实将其分布在 126 个活跃专家上,而模拟分布在 24 个上。因此真实的每个专家平均 M 更小,非空问题条目更多。其较低的 CV 表明分布更均衡,而准备与调度还取决于问题数量和粒度。许多更小的 GEMM 可用于摊销固定准备和启动开销的计算量更少。
cuBLASLt 算法选择在主机上反复运行。 原始剖析将 AlgoGetHeuristic 识别为主要的 BF16 瓶颈。多流、多张量路径仍会执行重复的算法选择,并为专家问题准备形状、布局和调用参数。两个启发式中位数均为 2.9 微秒,但真实的调用次数约为 4.6×,表明重复工作驱动了额外开销。外层 nvte_multi_tensor_gemm 计数相近,而每次调用内部的专家问题和准备工作不同。因此真实的中位调用时间为 2.29 毫秒,约为模拟的 5.2×。
更细粒度的内核延长了提交链。 真实的 nvjet 内核数量约为模拟的 25×,表明在相同 token 预算下 GPU 调用组织方式不同。如果 CPU 算法选择、参数准备和启动提交成为瓶颈,额外的流可能无法掩盖 GPU 等待下一批任务的时间。这就是该路径中需要检查的 CPU 气泡。计数直接支持了准备和调度工作的增加;GPU 间隙的持续时间仍需相应的时间线。
稠密分组 GEMM 针对这一准备与调度链。 第 5 节的连续缓冲区、设备端偏移与元数据以及可复用分配将专家问题置于单一分组调度路径上,减少了主机端切片、列表组装、参数构建和重复组织。这些变化影响 GEMM 周边的工作和完整调用时间,因此观察到的差异应与实现路径一并解读。实际速度比仍基于可比较的模块或学习器计时,BF16 和低精度评估均使用真实权重与路由。

记录的剖析

记录的剖析
图 7. 说明调用结构的代表性剖析。截图未提供形状和计时器匹配的配对计时,因此不以 3.594 毫秒和 156 微秒相除来报告加速比。性能来自上文的模块/学习器测量。
6.5 Rollout 与在线 RL:系统收益与剩余成本
采样器首先在固定输出负载下测试。八个引擎使用真实 DAPO 提示,生成恰好 1024 个 token 并忽略 EOS,并在针对每种配置进行扫描后独立选择并发度 64:
| 配置 | 聚合输出 token/s | 平均 TPOT | 吞吐比 |
|---|---|---|---|
| BF16 | 49,119.1 | 10.21 毫秒 | 1.000× |
| NVFP4 | 80,675.8 | 6.05 毫秒 | 1.642× |
表 6-17. 固定输出服务。BF16 使用 Triton MoE;NVFP4 使用 FlashInfer TRT-LLM/ModelOpt FP4 及逐 token 缩放。这测量的是采样栈,而非隔离 miniTE 训练算子。
八引擎固定输出测试采用与表6-17相同的测量方式。可在吞吐量与TPOT之间切换。由于没有重复运行间隔数据,因此未绘制误差条。静态版本展示输出吞吐量。
D16 · 八引擎请求速率、单引擎吞吐量与服务协议| 变体 | 最佳 并发度 | 最终输出token总数/秒 | 单引擎平均输出 | 请求总数/秒 | 平均TPOT | 对比同类 |
|---|---|---|---|---|---|---|
| NVFP4 | 64 | 80675.8 | 10084.5 | 78.78 | 6.05毫秒 | 1.64× 对比BF16 |
| BF16 | 64 | 49119.1 | 6139.9 | 47.97 | 10.21毫秒 | 基线 |
表6-18. 完整服务结果。在固定输出长度下,请求速率与输出token吞吐量描述的是同一工作负载。
静态内存占比为0.6,分段CUDA图已禁用。请求使用聊天模板、无填充输入、无速率限制、缓存刷新,上下文限制为8192。单引擎扫描使用64个提示、并发度1–96;随后选定的并发度64在八个引擎上以512个提示进行测试。单引擎峰值吞吐量约为NVFP4的10206.8 tokens/s和BF16的6063.5,与八引擎聚合主结果分开。
完整的RL还将学习器更新和权重同步与生成连接起来。对于无重叠的顺序工作流,
Titer=Trollout+Tlearner+Tconvert+sync+Torchestration.T_{\mathrm{iter}}= T_{\mathrm{rollout}}+T_{\mathrm{learner}}+ T_{\mathrm{convert+sync}}+T_{\mathrm{orchestration}}.Titer=Trollout+Tlearner+Tconvert+sync+Torchestration.当各阶段重叠时,以实测墙钟时间为准。表6-19报告了BF16、原生TE NVFP4和miniTE per-token的已记录阶段:
| 配置 | 稳态步数 | 训练秒数 | 转换/同步秒数 | 两阶段总计秒数 |
|---|---|---|---|---|
| BF16 | 11 | 47.04 | 9.94 | 56.98 |
| TE NVFP4 | 10 | 44.70 | 50.88 | 95.58 |
| miniTE per-token,BF16反向 | 10 | 42.56 | 55.26 | 97.82 |
表6-19. 在线RL各阶段平均时间。训练包括对数概率重计算和actor训练。两阶段之和并非完整迭代。TE NVFP4使用其原生反向配置;miniTE per-token显式禁用反向量化。
堆叠条形图使用表6-19中实测的阶段均值:BF16为11个稳态步,TE NVFP4/miniTE per-token为10个。此图不包含rollout时间。滑块计算在假设无重叠且其他成本不变的情况下,需要节省多少rollout时间才能覆盖已记录的两阶段开销。这是条件性算术,并非完整RL的速度预测。打印版仅显示实测阶段。
miniTE per-token的训练速度约为BF16的1.11倍,TE NVFP4约为1.05倍。转换/同步分别约慢5.56倍和5.12倍。miniTE per-token在这两个阶段合计增加40.84秒。在本次运行中,rollout收益未能抵消同步成本,使得转换与传输成为进一步优化的重要目标。
这一负面结果很有价值:更快的低精度算子并不能保证更低的系统集成成本。独立的服务比率不能替代在同一运行窗口内测得的rollout时间,将这两个阶段相加也不能得出完整的RL加速比。
D17 · 三种RL配置、测量窗口与相对速度所有配置均采用单台8×B300节点上的同址部署,每个采样引擎使用一块GPU。性能ID 0–1被丢弃。BF16对第2–12步取平均(11步);miniTE per-token/TE NVFP4各自对第2–11步取平均(10步)。数值为算术平均值。
| 配置 | 相对BF16的训练加速比 | 同步时间 / BF16同步时间 | 前向 / 采样 | 反向 |
|---|---|---|---|---|
| BF16 | 1.000× | 1.000× | BF16,Triton MoE | BF16 |
| TE NVFP4 | 1.052× | 5.119× | 原生TE专家NVFP4;ModelOpt/FlashInfer采样器 | 原生配方,DBQ未设置 |
| miniTE per-token | 1.105× | 5.559× | miniTE专家per-token;相同的NVFP4采样器 | DBQ=1,BF16 |
表6-20. 比率具有明确方向。此前的0.90×值是miniTE per-token/BF16的训练时间比;速度是其倒数,并非完整RL结果。TE NVFP4和miniTE per-token使用相同类型的NVFP4 rollout配置;服务表并未针对每种配置独立重新测量。
关键结果 · §6.5
在八引擎固定输出测试中,miniTE per-token的服务吞吐量约为BF16的1.64倍。
6.6 实验所确立的结论
这些实验共同解释了miniTE的实现选择。算子与梯度检查支持接口、布局和反向精度。模块前向与直接F+B计时定位了局部成本。真实文本学习器测试确定这些选择是否降低完整更新成本。服务与在线RL测试揭示了权重转换与同步带来的额外约束。
因此,低比特训练收益必须在多个层面得到验证:数值正确性、持续学习、更快的局部计算以及可控的集成成本。我们在每个层面保留核心测量和完整的配置对比,并使用miniTE测试记录作为性能结果。
7. 使用与扩展miniTE
7.1 通过线性接口集成现有模型
实现包为 miniTransformer, 而线性层是miniTE的基本集成单元。 Linear, LayerNormLinear, 以及 GroupedLinear 提供普通投影、归一化与线性计算组合,以及MoE专家计算。现有的Megatron封装可以连接到这些模块,用户也可以直接传入由命名配方构建的参数:
import torch
from miniTransformer.module.config import get_qlinear_params
from miniTransformer.module.linear import Linear
params = get_qlinear_params("nvfp4_per_token")
layer = Linear(
1024, 1024,
bias=False,
params_dtype=torch.bfloat16,
device="cuda",
qlinear_params=params,
)
配方选择决定数值路径;模块集成决定哪些模型层使用miniTE。MoE启动器为其miniTransformer后端同时设置两者。 GroupedLinear , 施加输出维度和分组布局约束,且不支持序列并行。集成时应检查具体模型和并行配置。
miniTE保留TE/Megatron检查点参数命名和兼容接口。当前的额外状态接口无法恢复所有TE内部量化状态。参数加载兼容性与完整内部状态兼容性分别验证。
7.2 显式配方与精度边界
版本1.0.0使用描述数值行为的名称。用户选择带有 MINITE_RECIPE 的训练配方,或用 get_qlinear_params:
| 命名配方 | 前向与缩放 | 反向与使用 |
|---|---|---|
nvfp4_per_token | 逐token NVFP4;TE权重缩放 | MXFP8反向;首次使用示例 |
nvfp4_per_token_bf16_bwd | 逐token NVFP4 | 专用TE风格BF16反向算子 |
nvfp4_cursor2_cutlass_v2 | 逐token激活、逐专家张量权重amax;CUTLASS V2 | MXFP8反向;自定义分组路径 |
nvfp4_block | NVFP4块缩放 | NVFP4反向;块缩放配方 |
nvfp4_per_token_w_per_tensor | 逐token激活、逐专家张量权重amax | 权重缩放消融实验 |
unquantized | 数据类型由层参数控制 | 高精度比较时禁用量化 |
表9. 版本1.0.0中选定的命名条目。性能实验保留各自的源修订版本、前向/反向精度和量化覆盖范围;名称用于选择当前软件配置。
对于集成训练程序:
MINITE_RECIPE=nvfp4_per_token python train.py
此处 train.py 表示用户的训练入口点。通过当前命名接口表达报告中miniTE的逐token BF16层旁路选择,得到:
MINITE_RECIPE=nvfp4_per_token DISABLE_BACKWARD_QUANT=1 python train.py
该层旁路和专用 nvfp4_per_token_bf16_bwd 算子分别实现并记录。第6节中两条miniTE路径的速度属于原始实验路径。专用算子需要其自身的性能测量。历史复现还使用原始修订版本的启动配置,并固定采样器量化、检查点转换、仅专家覆盖范围和TE修订版本。
7.3 将诊断作为扩展开发的一部分
配置检查在构建模型之前开始。请从 miniTransformer 仓库根目录运行以下命令:
python -m miniTransformer recipes
python -m miniTransformer describe nvfp4_per_token
python -m miniTransformer doctor --recipe nvfp4_per_token
bash develop/script/moe_training_e2e.sh \
--backend mini_nvfp4_per_token --profile off --dry-run
recipes 和 describe 可在不使用 Torch 或 TE 的情况下暴露名称、依赖关系和反向覆盖。 doctor 检查可见的 CUDA GPU、TE 量化器参数、选定的 dense/list 分组接口以及算子导入。 --dry-run 打印配方、相关环境和完整命令,而不启动训练或创建日志/缓存目录。实际运行会在 run_meta.json.
中记录配方名称、前向/反向设置、布局请求和配置来源。
本次接口更新的全部 21 项 CPU 兼容性检查均已通过,涵盖解析、冲突选择器、无 GPU 依赖的检查、后端别名、试运行和元数据。新增的 GPU 等价性测试涵盖独立配置、受支持的 Dense/GroupedLinear 路径上的前向/反向行为,以及两种 BF16 反向选择。本次更新未在本地运行这些测试。第 6 节中的数值和性能结果保留其原始实验标识。在扩展低精度路径时,我们先检查算子,再检查模块和跨框架行为,最后完成训练。张量钩子收集激活、权重、输出和梯度统计信息。性能分析定位内核开销、调用间隙和内存分配。转储/比较脚本检查 MoE 中间结果和 token 对数概率。新配方在注册表中声明其用途和依赖关系,通过现有参数工厂连接其算子,并遵循这些数值和运行时检查。
仓库 README 涵盖安装配置、层集成和启动示例;双语配方指南记录了选择与覆盖规则。模型转换、采样和评估需要各自固定的版本与设置。当前软件接口与历史实验配置分别记录。
8. 讨论:从数值对齐到端到端学习
实验通过两条路径应对低精度强化学习:TRIAGE 在存在失配时改变更新,而 miniTE 减少执行差异并降低计算成本。两者都无法解决所有问题。接下来的问题是,更快的计算是否能带来端到端节省,以及更接近的数值一致性是否能带来更好的学习效果。
8.1 权重更新与同步仍是瓶颈
在这里, update weight 意味着将训练好的权重转换、传输并加载到采样器中。在我们测量的配置中,这仍然是一个主要瓶颈。更快的学习器更新和 rollout 并不意味着完整 RL 迭代时间会成比例减少。在第 6.5 节报告的运行中,转换和同步成本抵消了更快采样带来的收益。
进一步的优化应统筹考虑参数表示、量化与打包、传输、加载以及同步频率。可能的方向包括:在同一权重版本内复用低精度表示、将转换与传输相结合,以及在显式版本控制下让同步与采样重叠进行。降低同步频率可以节省传输时间,但会加剧策略陈旧。这些选择应在匹配的训练预算下,用完整的迭代时间和任务性能来评估。在线 amax 与更深的算子融合也需要在这一端到端范围内进行评估。
8.2 为什么算子对齐后低精度仍然敏感?
匹配算子和量化配方消除了一些执行差异。然而,仅匹配它们的名称并不能保证在不同输入形状、批次和并行配置下得到完全相同的浮点计算。浮点加法不满足结合律:归约顺序、融合边界或分块方式的变化都可能改变结果。固定的执行路径可以是确定性的。剩下的问题是:当执行条件发生变化时,扰动从何而来,以及它们如何影响后续更新。
我们的对齐实验和长期评估揭示了这个问题的两个方面。严格的对齐配置将差距缩小到约 10⁻⁸,而 miniTE 的长期学习和检查点评估并非在每次比较中都匹配 BF16。这些结果使用了不同的配置和测量范围,但合在一起说明了为什么数值一致性和最终学习质量必须分开评估。为什么在我们测试的低精度配置中,看似微小的扰动可能更具破坏性?
候选机制包括持续存在的误差偏置、粗量化边界附近的舍入转换、MoE 路由变化,以及梯度裁剪与优化器状态的累积响应。这些仍属假设。更具信息量的实验应固定模型和数据,控制扰动的位置、方向和相关性,并比较 BF16 与低精度下后续更新和长期奖励,而不是依赖单步的平均差距。
8.3 优化更新与最终学习成果
平均对数概率差无法线性转换为奖励。相似的平均差距可能出现在不同的 token、优势符号和响应片段上,并对更新产生不同影响。较小的差距也不能保证有用的探索或梯度信号。TRIAGE 的方向性分析提供了一个起点。进一步的工作可以直接考察误差如何改变更新,并在匹配的数据、更新和实际耗时预算下,通过奖励和评估分数检验修正效果。
社区的方法也在不断拓宽。Together AI 的 Score Centering(SC)加入了一项修正,可消除策略梯度中由失配引起的漂移,并且可以与重要性采样结合使用[17]。Qwen 团队的 TRACE 利用 rollout 侧的量化结果来指导训练侧的 FP4 舍入,并压缩引导信息[18]。SC 改变了梯度估计器;TRACE 协调量化决策。TRACE 仍以对齐为目标,但这两项工作都将可用的干预手段扩展到了匹配算子实现之外。
对 miniTE 而言,这意味着要支持围绕更新规则、精度分配和采样策略的联合实验,并配合受控且可比的数值执行。我们希望社区能够研究如何保留与任务相关的学习信号、抑制在多次更新中累积放大的偏差,并将这些改进转化为端到端的收益。目标是在给定资源预算内实现更好的学习效果,而不仅仅是获得更小的一致性指标。
附录 A. 实验设置与统计
A.1 补充算法与训练材料
TRIAGE 的方法与结果遵循发布稿及其诊断附录。检查点采用表 3 中的设置,在五项数学基准上等权。更新后的 4B 均值为 BF16 的 58.26% 和 TRIAGE 的 58.49%;在共享的 step-300 检查点上,TRIAGE 得分为 53.03%。
表征使用独立的诊断运行:4B 使用更新 1–324,30B 使用更新 1–724,最大补全长度分别为 16,384 和 20,480。4B 分析使用学习器与采样器之间的直接差距;30B 使用路由重放后的残差差距。在健康阶段,对有效 token 的固定尾部阈值分别为 −0.24208 和 −0.03949。
这些诊断运行使用 DAPO-Math-17K、学习器种子 1234、rollout/数据种子 42、每次迭代 16 个提示、每个提示 16 个响应以及 GBS 256。它们使用 RHT 和梯度随机舍入;这些数值设置与 miniTE 的实际路径不同,在后者中这些选项被禁用。详细的窗口、长度敏感性、单步更新审计和随机化控制见诊断附录。
TRIAGE 材料用于算法结果。其镜像未调整到与 miniTE 性能环境相同的水平,因此其吞吐量不被用作 miniTE 的上限或算法开销基线。
A.2 补充的长程强化学习与检查点评估
图 5 和表 6-8 使用以下训练配置和评估协议:
| 项目 | 设置 |
|---|---|
| 模型与数据 | Qwen3-30B-A3B-Base;DAPO-Math-17K,数学奖励 |
| 目标 | GRPO;非严格配置下的逐 token TIS |
| 种子 | 训练种子 1234;rollout 洗牌种子 42 |
| 批次 | 每次迭代 32 个提示,每个提示 16 个响应;GBS 512 |
| 生成 | 温度 1.0,top-p 1.0,top-k 禁用,最大响应长度 20,480 |
| 并行性与分配 | 8 个节点,每个节点 8 块 actor GPU;TP/PP/CP/EP/ETP = 1/1/1/4/1;actor 与采样器同址部署 |
| 优化器 | Adam,学习率 10−610^{-6}10−6,无预热;权重衰减 0.1,β1=0.9,β2=0.98,ϵ=10−8\beta_1=0.9,\beta_2=0.98,\epsilon=10^{-8}β1=0.9,β2=0.98,ϵ=10−8,梯度裁剪 1.0 |
| 内存策略 | 完全激活重计算、动态微批次;分布式精度感知优化器、CPU 卸载与重叠传输 |
| TIS 与裁剪 | 逐 token TIS 裁剪至 [0,2];策略更新比率 [0.80,1.28];路由重放已禁用 |
| 评估 | 第 999 步;温度 0.7,top-p 1.0,最大生成长度 20,480 |
| 重复与指标 | 每个 AIME 评估六次;MATH-500、GSM8K、IFEval 各评估四次;IFEval 提示级严格模式 |
表 A1. 补充的 miniTE 长程协议。与固定文本学习器测试不同,本协议在单个 8×B300 节点上运行。重复的检查点评估不会增加训练种子数量。
三种主要配置覆盖第 0–999 步,每种配置各有一条轨迹,并采用居中的 21 步均值。图 4 保留四面板图,采用 15 步平滑;配套论文还包含一个相关对齐实验的双面板展示,采用 21 步平滑。每个图注保留各自的约定。严格对齐会同时更改多个执行组件并移除 TIS;其计时与学习器性能测试分开进行。
A.3 固定文本 miniTE 学习器测试
| 项目 | 设置 |
|---|---|
| 模型与硬件 | Qwen3-30B-A3B,单个 8×B300 节点 |
| 并行性与形状 | TP/PP/CP/EP = 1/1/1/4;序列长度 4096;微批次 1;GBS 128 |
| 工作负载 | 每次更新 524,288 个 token;累计 16 个微批次 |
| 数据与初始化 | 657 条冻结的 4097-token DAPO 序列;相同的 BF16 检查点;随机种子 20260913 |
| 软件 | PyTorch 2.11.0+cu130、CUDA 13.0.88、NCCL 2.28.9、Megatron-Core 0.16.0rc0 |
| Megatron | 修订版 52153937,用于该实验的集成补丁 |
| 核心修订版 | miniTE 953f17f9;十个独立进程,无排列融合 |
| 优化修订版 | miniTE b641179;密集、GEMM 类、排列融合与 amax 复用 |
| TE | 早期 eaca54e9,随后 8ea5a92d 处理零 token 专家;检查非空组路径 |
| 每个 token | 690ffea4 ,包含 CUDA 13 构建修复、独立环境 |
表 A2. 固定文本性能协议。数据构造包含约 4.3 倍重复。激活重计算、卸载和 CUDA 图均被禁用。量化路径的专家填充开销已计入计时。
核心协议在每个进程中运行 20 次更新,使用 5 次预热更新,并测量第 6–15 次更新。每次运行的吞吐量等于 524,288 除以更新耗时中位数。在十次运行中,先去掉最大值和最小值再取平均。加速比首先按运行 ID 配对,然后以相同方式修剪。
优化后的协议为每个配置使用一个进程,并进行17次更新。在两次预热后,去掉其余15次更新耗时中的最大值和最小值。吞吐量等于524,288个token除以保留耗时的平均值;加速比使用同一轮的基线。各估计量分别报告。核心运行3受到一个空闲服务进程的干扰,受影响的极小值观测值已按预定义的修剪规则移除。全部十条记录及修订差异均保留在miniTE摘要中。
表6-9使用来自十次独立运行的配对加速比。1.334×值是各次运行比值的修剪均值;1.344×是聚合吞吐量均值之比。后续融合实验使用不同的修订版和配置,并单独报告。
A.4 模块、服务与强化学习阶段的时序
模块测试使用一块 GPU、32 个本地专家、四种预分发 token 大小,以及五次独立进程运行,并采用热静态权重。分发被排除在外,FC2 激活事先准备好。模块前向增益与学习器加速不同。上游 TE 的每 token 在完整模型量化覆盖范围和反向精度上与 miniTE 不同;其完整模型对比不用于隔离实现增益。
服务时使用固定输出长度 1024, ignore_eos=true上下文限制为 8,192 个 token,静态内存占比为 0.6,并禁用了分段 CUDA 图。单个引擎使用 64 个提示,对并发度 1、2、4、8、16、32、48、64 和 96 进行扫描。两种配置均选择 64,然后在八个引擎上处理 512 个提示。
共置 RL 测试使用一个 8×B300 节点,每个引擎使用一块 GPU。BF16 统计数据涵盖性能步骤 2–12;miniTE 涵盖 2–11。学习器计时包括对数概率计算和 actor 训练。权重更新包括 Megatron→SGLang 同步以及 miniTE 的 BF16→NVFP4 转换。两阶段总和不包括完整 rollout 和其他编排工作。
API 与融合布局检查在固定形状和配置下比较编码、缩放和输出。这些边界上的逐位一致性并不能确立全模型批次不变性。此前的 swizzle 失败与生产环境 miniTE 的逐 token 检查使用不同的形状约束,需分别解读。
附录 B. TRIAGE 方法约定
我们采用发布稿中的响应级归一化:在每个响应内计算加权 token 损失的均值,然后对各响应取平均。给定有效 token 掩码 mi,tm_{i,t}mi,t 和分离的门控权重 wi,tw_{i,t}wi,t,
\mathcal L_{\mathrm{PG},i} =\frac{\sum_t m_{i,t}w_{i,t}\ell^{\mathrm{base}}_{i,t}} {\sum_t m_{i,t}w_{i,t}},\qquad \alpha_{i,t}=\frac{w_{i,t}}{\bar w_i},\quad \bar w_i=\frac{\sum_t m_{i,t}w_{i,t}}{\sum_t m_{i,t}}.LPG,i=∑tmi,twi,t∑tmi,twi,tℓi,tbase,αi,t=wˉiwi,t,wˉi=∑tmi,t∑tmi,twi,t.固定门控使用 64 个 token 的分段。对于负优势分段,当平均差距低于 −1.5 时,均值层权重为 0.1;当平均差距在 [−1.5, −0.5) 区间内时为 0.3;其他情况为 1.0。如果该分段中有 kkk 个有效 token 满足 δ<−6\delta<-6δ<−6,则尾部权重为 1−0.9min(k/2,1)1-0.9\min(k/2,1)1−0.9min(k/2,1)。两个分段权重中较小的那个仅应用于负差距 token。其他 token 保持权重 1。
修复在正优势响应内使用 128 token 的分段,并在负边界之外施加单侧伪 Huber 惩罚。掩码、分段归一化和导数在 TRIAGE 方法附录和主方法章节中描述。
组件消融从完整 TRIAGE 的 step-300 检查点分支,并比较继续训练的结果。完整图表见 TRIAGE 消融部分。这些是续训对比,而非从零开始训练的消融实验。
这些图表和交互式演示借鉴了 四苦教训 以及 LMSYS Blackwell 低精度强化学习文章 [15,16]。构造示例和条件计算单独识别。
参考文献
[1] NVIDIA。 使用 NVFP4 预训练大型语言模型。2025;2026年更新。 arXiv:2509.25149.
[2] Cursor Research。 Composer 2 技术报告. 2026. arXiv:2603.24477.
[3] 光明生等。 HybridFlow:一个灵活高效的 RLHF 框架EuroSys 2025。 arXiv:2409.19256.
[4] 邵志宏等。 DeepSeekMath:推动开放语言模型数学推理的极限. 2024. arXiv:2402.03300.
[5] 杨安等人。 Qwen3 技术报告. 2025. arXiv:2505.09388.
[6] Qiying Yu 等。 DAPO:一个大规模开源 LLM 强化学习系统. 2025. arXiv:2503.14476; DAPO-Math-17K 数据集.
[7] Lianmin Zheng 等 SGLang:结构化语言模型程序的高效执行NeurIPS 2024。 arXiv:2312.07104.
[8] Mohammad Shoeybi 等。 Megatron-LM:使用模型并行训练数十亿参数的语言模型. 2019. arXiv:1909.08053.
[9] NVIDIA Transformer Engine。 添加 NVFP4 逐 token 量化配方,PR #3045。2026。 实现与讨论.
[10] 邱兆鹏等。 FP8-RL:一套实用且稳定的 LLM 强化学习低精度技术栈. 2026. arXiv:2601.18150.
[11] Karl Cobbe 等。 训练验证器以求解数学应用题. 2021. arXiv:2110.14168; 官方 GSM8K 仓库.
[12] Dan Hendrycks 等。 用 MATH 数据集衡量数学问题求解能力. 2021. arXiv:2103.03874。关于 MATH-500 划分,见 Hunter Lightman 等人, 让我们逐步验证, 2023, arXiv:2305.20050,以及 官方划分说明.
[13] 美国数学协会。 官方 AMC 页面 和 AIME 简介结果使用实验中记录的 AIME 2024、AIME 2025 和 AMC23 评估版本。
[14] Jeffrey Zhou 等。 面向大型语言模型的指令遵循评估. 2023. arXiv:2311.07911.
[15] Ziang Li 等。 4-bitter 的教训humans&,2026-07-10。 文章与交互示例.
[16] Ziang Li、humans& 与 Miles 团队。 迈向 Blackwell 原生的 8 位与 4 位强化学习:Miles 中的端到端 MXFP8 与 NVFP4 强化学习LMSYS,2026-07-29。 文章与格式说明.
[17] Martin Marek 和 Max Ryabinin。 分数居中稳定了离策略强化学习。 2026. 论文.
[18] 王鑫等。 TRACE:面向 MoE 语言模型 FP4 强化学习的、由 rollout 引导的量化感知训练。 2026. 论文.