TRIAGE:原生 NVFP4 强化学习的方向感知失配稳定化 TRIAGE: Direction-Aware Mismatch Stabilization of Native NVFP4 Reinforcement Learning
方向感知的门控与修复让原生NVFP4强化学习稳定训练且提速2.3倍
前置知识
GRPO(组相对策略优化)
GRPO 是推理模型 RL 的主流算法:对每个 prompt 采样一组 $G$ 个回答,用组内奖励归一化得到优势 $A_i = (r_i - \mu_r)/(\sigma_r + \epsilon)$,并赋给该回答的所有 token。训练目标是对 token 级重要性比率 $\rho_{i,t} = \pi_\theta/\pi_{old}$ 做裁剪(clip)后的策略梯度:正优势提升 token 概率,负优势降低它。
本文的方向判据完全建立在 GRPO 的结构上:单步更新内每 token 的梯度系数 $g_{i,t}$ 就等于优势 $A_i$,因此「优势符号 × 失配符号」决定了更新会放大还是收缩失配,这是划分 $\mathcal{A}^-$/$\mathcal{A}^+$ 区域的基础。
NVFP4 量化格式与 W4A4 执行
NVFP4 是 NVIDIA Blackwell 支持的 4 位浮点格式:E2M1 FP4 数值、每 16 个值共享一个 E4M3 FP8 微缩放因子 $s_b$、外加 FP32 张量级二级缩放 $S$,即 $\hat{x}_j = S\, s_b\, Q_{E2M1}(x_j / S s_b)$。W4A4 指权重和激活都是 4 位,由 Tensor Core 原生执行,峰值矩阵算力约为 BF16 的 4 倍。
论文的核心立场是在采样器和学习器两侧都保留原生 W4A4 前向(而非 fake quant 模拟)。理解微缩放如何改善局部数值范围,才能明白量化扰动为何会在逐 token 概率上产生 learner–sampler 差异。
学习者–采样器失配与离策略修正
现代 RL 把生成(SGLang/vLLM 等采样器,策略记为 $\pi_s$)和梯度计算(Megatron-LM 等学习器,策略记为 $\pi_l$)解耦,两条执行路径即使权重相同也有概率差:$\delta_t = \log \pi_l(y_t) - \log \pi_s(y_t)$。精确的离策略修正需要权重 $\pi_\theta/\pi_s = (\pi_\theta/\pi_l)\cdot\exp(\delta)$,失配误差全部由 $\exp(\delta)$ 这一项承载。
$\delta$ 是全文的核心统计量:诊断段统计的是它的均值和尾部占比,门控按 $\delta$ 的符号选择性衰减,修复项压的是段级平均失配。看不懂 $\delta$ 就无法理解 TRIAGE 的任何组件。
截断重要性采样(TIS)与幅度式干预
TIS 是 FP8-RL 等工作采用的稳定化手段:当 token 级重要性权重超过阈值 $C$(本文取 2.0)时截断,另有裁剪(clipping)和掩码(masking)等变体。它们的共同点是只根据 learner–sampler 分歧的「幅度」决定是否干预。
这是本文要超越的基线:论文论证幅度只记录分歧多远,不预测下一步梯度会把分歧推向何处。实验也显示 TIS 在 30B 上训满 1700 步不发散,但平均分 63.18% 反而低于崩溃前 naive NVFP4 的 65.87%。
伪胡贝尔损失(pseudo-Huber penalty)
伪胡贝尔损失 $\phi_\beta(q) = \frac{\beta^2}{2}\left(\sqrt{1+(q/\beta)^2} - 1\right)$ 是 Huber 损失的光滑近似:$q$ 小时 behaves 近似二次(梯度温和),$q$ 大时近似线性(梯度有界为 $\beta$ 量级),因而对大误差不会产生爆炸性梯度。
TRIAGE 的修复项用它对超过边界 $\tau_{rep}$ 的段级平均失配 $q_R = [\tau_{rep} - \bar{\delta}_R]_+$ 施加「有界」矫正压力,这是修复不会反过来破坏训练稳定性的关键设计。
研究动机
强化学习已成为提升大模型推理能力的核心手段,但随着长链推理、代码和 agent 任务兴起,生成轨迹越来越长,rollout 成为训练循环中占主导的时钟瓶颈。NVIDIA Blackwell GPU 提供原生 W4A4 Tensor Core,峰值矩阵算力约为 BF16 的 4 倍,使 NVFP4 成为长序列 rollout 极具吸引力的精度选择。然而现代 RL 系统解耦了 rollout 与训练:推理引擎负责采样,训练框架重算前向概率求梯度,两条执行路径即使权重同步也存在逐 token 概率差,低精度执行会进一步放大这种 learner–sampler 失配,形成硬件效率与策略优化稳定性之间的直接冲突。现有两条补救路线都不理想:其一,量化感知训练(QAT)试图让学习器对齐低精度采样器,但它保留高精度算术,无法保住学习器侧的原生 W4A4 前向;其二,TIS、裁剪、掩码等按失配幅度干预,可幅度只说明双方分歧多远,并不能预测随后的策略梯度更新会缩小还是扩大分歧。实测后果严重:朴素 NVFP4 在 Qwen3-4B-Base 上约 300 步后崩溃(奖励骤降、失配激增),在 Qwen3-30B-A3B-Base 上约 700 步后失配快速增长直至发散。
本文的目标是本文的目标是在采样器和学习器两侧都保留原生 NVFP4 W4A4 前向执行的前提下,让 RL 训练在整个训练区间内保持稳定,并在五个数学推理基准(AIME24、AIME25、MATH500、AMC23、GSM8K)上恢复到 BF16 级性能,同时保住 NVFP4 带来的 2 倍以上 rollout 吞吐优势。更具体地,作者要做三件事:建立一个能在失配全面恶化之前提前预警的诊断量(而不是等到边缘分布爆炸才发觉);设计一个只干预「会放大失配的更新」的选择性稳定机制(而不是无差别压制所有失配);把整套方法验证为可直接跑通稠密模型(Qwen3-4B-Base)和 MoE 模型(Qwen3-30B-A3B-Base)的端到端原生低精度 RL 配方——全部实验总计超过 34,500 个 B300 GPU 小时。
与已有工作不同的是,本文的独特切入是把失配从「幅度问题」重构为「方向问题」。作者对失配能量 $D = \frac{1}{2}\sum \delta_{i,t}^2$ 做一步更新后的一阶展开,得到 $\Delta D = \eta\sum \kappa_{i,t}\,\delta_{i,t}\,g_{i,t} + \eta R_{cross} + \mathcal{O}(\eta^2)$,指出 $\delta g > 0$ 的更新自我放大、$\delta g < 0$ 的更新自我收缩——同样大的失配,一个可能正在自我修复,另一个正在被梯度系统性放大,幅度统计完全无法区分两者。在此之上还有两层实证新发现:其一,在边缘分布仍然「健康」时($|\delta| < 0.05$ 的 token 占比在 4B 上仍不低于 83.0%、30B 上不低于 91.9%),两个理论上对称的放大区已经严重失衡($\rho_{asym}$ 在 4B 上从 1.02 升到 2.32,30B 上从 1.34 升到 1.85),即方向性失衡是先兆、负尾爆炸才是终末期;其二,尾部失配 token 集中在少数 64-token 段内(top 5% 的段含约四分之一尾部 token,约为均匀分布的 5 倍),而响应级平均完全掩盖了这种局部结构。这共同引出「段级诊断 + 方向选择性干预」的全新稳定化原则。
核心方法
直觉上,TRIAGE 不试图消灭所有失配,而是识别并只控制那些会被当前更新进一步放大的部分——就像分诊(triage)只优先处理会恶化的病人。它通过修改优化目标实现,保持采样器和学习器的原生 NVFP4 W4A4 前向完全不变。技术路线是把每个 token 的有效更新系数写成 $g_{i,t} = \alpha_{i,t}\, g^{base}_{i,t} + \psi_{i,t}$,两项按优势符号分工:$\alpha_{i,t}$ 在负优势响应内部重分配策略梯度质量(乘性门控,第 4.1 节),$\psi_{i,t}$ 在正优势响应上对残余严重失配施加有界矫正压力(加性修复,第 4.2 节)。底层损失 $\ell^{base}$ 可以是任何带重要性修正的现有 token 损失,因此 TRIAGE 是一个与具体 RL 框架正交的插件式稳定层。全部操作只依赖训练中本来就要计算的逐 token log 概率差 $\delta_{i,t}$,不引入 fake quant,也不需要低比特学习器前向,这是它与 QAT、QUADS、QaRL 路线的根本不同。
核心创新是「方向判据 + 段级诊断 + 响应内重归一化」的三位一体。与 TIS 按幅度截断的本质区别在于选择性:TRIAGE 只衰减 $\mathcal{A}^- = \{A_i < 0, \delta_{i,t} < 0\}$ 区域——对负优势响应中本已低于采样器概率的 token 继续压概率,只会越推越偏;而同一危险段内满足 $A_i < 0, \delta_{i,t} > 0$ 的 token 天然收缩失配,完全不被门控。同样刻意的还有「不做什么」:正优势放大区 $\{A_i > 0, \delta_{i,t} > 0\}$ 不被门控,因为衰减奖励为正的回答上的概率提升会直接抑制学习信号(第 3 节实证显示方向性失衡集中在负优势侧)。诊断尺度选在 64-token 段而非 token 或响应级:段能捕捉局部持续风险,避免 token 级噪声,又不被响应级平均掩盖(97.2% 的含重尾段响应其平均 $|\bar{\delta}_i| < 0.05$)。最后,门控是乘性的、只能把放大更新压向零而不能反向,因此配一个加性、有界的修复项:只加在正优势响应上,对负 gap token 有 $\delta_{i,t}\psi_{i,t} < 0$,天然贡献收缩项。
方法步骤详情
第一步,失配统计:利用训练中已有的学习器与采样器逐 token log 概率,得到 $\delta_{i,t} = \log\pi_l - \log\pi_s$,无需额外前向。第二步,段诊断:把每个负优势响应切成 $W = 64$ token 的不重叠段 $S$,对每段计算平均位移 $\bar{\delta}_S = \sum_{t \in S} m_{i,t}\delta_{i,t} / \sum_{t \in S} m_{i,t}$ 和深尾占比 $f^{tail}_S = \sum_{t \in S} m_{i,t}\,1[\delta_{i,t} < \tau_{tail}] / \sum_{t \in S} m_{i,t}$。第三步,门控权重:把两个统计量映射为 $w_S = \mathcal{G}(\bar{\delta}_S, f^{tail}_S) \in [w_{min}, 1]$,位移越负、尾越重衰减越强,操作阈值在训练前固定(附录 C)。第四步,方向选择性施加:仅当 $A_i < 0$ 且 $\delta_{i,t} < 0$ 时取 $w_{i,t} = w_S$,否则为 1。第五步,归一化加权损失:$\mathcal{L}^{PG,i}_{TRIAGE} = \sum_t m_{i,t} w_{i,t} \ell^{base}_{i,t} / \sum_t m_{i,t} w_{i,t}$,相对原响应均值诱导 $\alpha_{i,t} = w_{i,t}/\bar{w}_i$,把更新质量从目标子集重分配给响应内其余 token,而非整体削弱该响应。第六步,修复项:对正优势响应的修复段 $R$,定义单侧越界 $q_R = [\tau_{rep} - \bar{\delta}_R]_+$,在总损失上加 $\frac{\lambda_k}{N_{rep}} \sum_R 1[A_i > 0]\,\phi_\beta(q_R)$,其中 $\phi_\beta$ 为伪胡贝尔罚;修复自固定步 $k_{rep}$ 起启用($\lambda_k = \lambda_{rep}\,1[k \ge k_{rep}]$),对应每 token 加性系数 $\psi_{i,t} = \frac{\lambda_k}{N_{rep}} \cdot \frac{m_{i,t}}{n_R} \cdot \phi'_\beta(q_R)\, 1[A_i > 0, q_R > 0]$。
技术新颖性
理论层面,论文首次对低精度 RL 的 learner–sampler 失配给出显式一阶分解(附录 A 推广到带优化器残差和预条件的更新),把「失配要不要管」化归为 $\delta_{i,t} g_{i,t}$ 的符号问题,并据此提出两阶段失效观:方向性失衡($\rho_{asym}$ 从 1.02/1.34 升至 2.32/1.85)是可预警的先兆,负尾爆炸($q_{0.05}(\delta)$ 达 $-42.75$/$-10.94$)才是终末期。实证层面,用四个 25-update 滑窗(healthy/drift/pre-terminal/terminal)系统刻画了 native NVFP4 失配的时空演化,包括 top 5% 段集中约 25% 尾部 token、重尾段邻接持续性从 2.6 倍升至 3.8–4.0 倍等定量观察。方法层面,与 ACRL 的响应级四象限调整不同,TRIAGE 由失配能量分解出发并在段级定位风险;与 QUADS/QaRL 不同,它不需要 fake quant 或低比特学习器前向,纯靠目标函数修改保住原生 W4A4。系统层面,这是首个在稠密与 MoE 模型上都跑通的端到端 native NVFP4 RL 配方,并精确测量了方法自身开销:中位 learner 更新开销仅 0.84%(绝对 +0.225 秒)。
实验结果
训练稳定性方面:Qwen3-4B-Base 上朴素 NVFP4 约 300 步后崩溃,奖励骤降且平均绝对失配激增;TIS 只延迟不阻止失败;TRIAGE 在全部 600 步内保持稳定,相当于把可观测的稳定训练区间约翻倍。Qwen3-30B-A3B-Base 上朴素 NVFP4 约 700 步后失配激增;TIS 在 1700 步内未发散,但奖励轨迹在后半程持续低于 BF16 和 TRIAGE。基准性能方面(Table 1):30B 上 TRIAGE 五基准平均 70.96%(AIME24 51.67%、AIME25 38.33%、MATH500 91.60%、AMC23 77.72%、GSM8K 95.47%),与 BF16 的 72.41% 仅差 1.45 个百分点,显著高于崩溃前 naive NVFP4 的 65.87% 和 TIS 的 63.18%;4B 上 TRIAGE 在 600 步达 58.49%,略超 BF16 的 58.26%。关键对照是 4B 同为 300 步检查点时,TRIAGE 比 naive NVFP4(47.10%)高 5.93 个百分点、比 TIS(50.76%)高 2.27 个百分点,说明增益并非单纯来自训练更久。效率方面(Table 2):TRIAGE 的 rollout 吞吐为 6,745 token/s(batch 256,2.27 倍)和 8,717 token/s(batch 512,2.30 倍),几乎保住朴素 NVFP4 的 6,201/8,942(2.08/2.35 倍);端到端迭代时间从 BF16 的 208.2/254.5 秒降至 172.1/195.8 秒,即 1.21 倍和 1.30 倍加速;TRIAGE 自身中位 learner 开销仅 0.84%(+0.225 秒)。消融方面(Figure 5,全部从 300 步 TRIAGE 检查点分支):完整 TRIAGE 奖励最高;去掉修复(Gate+TIS)后失配仍低但优势驱动的更新失衡未纠正、奖励下降;单独 TIS 从 TRIAGE 检查点续训竟然稳定,与从头训练时崩溃形成鲜明对照,印证失配动态是阶段依赖的;朴素 NVFP4 则失配急增并在后期崩溃。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数学推理五基准平均(Qwen3-30B-A3B-Base) | 平均准确率(step 1700) | NVFP4+TRIAGE:70.96% | BF16:72.41%;NVFP4:65.87%(700 步崩溃前);NVFP4+TIS:63.18% | 距 BF16 仅 -1.45pt,比 naive NVFP4 高 5.09pt,比 TIS 高 7.78pt |
| 数学推理五基准平均(Qwen3-4B-Base) | 平均准确率(step 600) | NVFP4+TRIAGE:58.49% | BF16:58.26%;NVFP4:47.10%(300 步崩溃前);NVFP4+TIS:50.76% | 超过 BF16 0.23pt;同步骤 300 时比 NVFP4/TIS 分别高 5.93/2.27pt |
| Rollout 吞吐(30B,global batch 512) | 输出 token/s | 8,717(NVFP4+TRIAGE) | BF16:3,798 | 2.30 倍(朴素 NVFP4 为 2.35 倍,几乎无损保留) |
| 端到端 RL 迭代耗时(30B,global batch 512) | 秒/迭代 | 195.8 s(NVFP4+TRIAGE) | BF16:254.5 s | 1.30 倍端到端加速(batch 256 时为 1.21 倍) |
局限与改进
作者承认的局限:其一,规模受限,因算力约束未验证 100B+ 模型或显著更长的训练区间,主实验在轨迹收敛或呈稳定晚期行为后即终止;其二,只做了同步 RL 实验,虽然方向判据理论上不依赖同步采样,且在策略陈旧性叠加失配的异步 RL 中可能更有价值,但未实验验证;其三,端到端加速约 1.3 倍,明显低于 rollout 侧的 2.3 倍,大量时间花在 GEMM 之外的操作(训练器与采样器间参数同步、Amax 计算等量化相关归约和数据变换)。我自己的观察:方法引入大量固定超参——段宽 $W=64$、诊断阈值 $\tau_{ana}$、尾阈值 $\tau_{tail}$、门控下界 $w_{min}$、修复边界 $\tau_{rep}$、启动步 $k_{rep}$、强度 $\lambda_{rep}$、伪胡贝尔宽度 $\beta$——虽声称训练前固定,但跨模型、跨任务(尤其非数学任务)的敏感性没有消融;评估全部集中在数学推理,代码生成或 agent 类长轨迹任务上「段级尾部集中」现象是否同样存在未经检验;理论依据是一阶展开加对角项,cross-token 交互与优化器动量的影响只做了范围审计(附录 B.3),$\delta g$ 只是局部风险指标而非逐 token 的无条件预测器;此外门控质量依赖 $\delta$ 计算本身,若两条执行路径的实现差异较大,方向判据可能被实现噪声污染。
独立分析的弱点
第一,超参体系庞大且耦合:诊断、门控、修复三套阈值($\tau_{ana}$、$\tau_{tail}$、$w_{min}$、$\tau_{rep}$、$k_{rep}$、$\lambda_{rep}$、$\beta$)都需逐模型设定(论文明确说 $\tau_{ana}$ 是 model-specific),迁移到新模型或新任务时要重调,改进方向是把阈值改为基于在线分位数的自适应校准,或用小规模代理任务元学习。第二,只治理 $\mathcal{A}^-$ 区域:论文的干预理由是实证上方向性失衡集中在负优势侧,但如果某任务(如偏好对齐、过程奖励)中 $\mathcal{A}^+ = \{A_i > 0, \delta_{i,t} > 0\}$ 失衡占主导,现方案会系统性漏防,改进方向是对称监测两个放大区的 $\rho_{asym}$ 并按主导侧自适应选择门控象限。第三,修复项作用在段级平均位移上:段内若深尾 token 少但极深、或多个中等深度 token 相互抵消使 $\bar{\delta}_R$ 不越界,修复不会触发,可补充 token 级兜底规则或用分位数而非均值刻画段风险。第四,端到端加速只有约 1.3 倍:瓶颈在参数同步和量化归约等系统开销,方法层的收益已接近天花板,继续提升需要融合量化 kernel 和更高效的权重同步协议。第五,评估面窄:只有数学推理、只有 GRPO、只有两个 Qwen3 模型,方法在其他 RL 目标(如 RLOO、ReMax)和其他能力域的泛化性完全未知。
未来方向
作者明确提出的方向:其一,扩展到异步 RL,联合刻画量化失配与策略陈旧性两个失配源——方向判据不依赖同步采样,异步场景下 sampler 概率更陈旧,方向性分析可能更关键;其二,在 100B+ 规模稀疏模型和显著更长的训练区间上检验同样的方向动态是否成立;其三,与更紧耦合的低精度 kernel 和 RL 框架协同设计,包括更高效的参数同步和融合量化路径,以兑现 Blackwell/Vera Rubin Tensor Core 的全部潜力;其四,超越统一修复规则,发展按失配幅度、局部性和持续性自适应的细粒度修复机制。基于成果可延伸的方向:把 $\rho_{asym}$ 和段级尾部集中度做成训练监控的通用预警指标(早停或自动降精度回退的触发信号);将方向判据推广到 FP8 及其他格式、以及 DPO 等 off-policy 目标;研究门控与 TIS 的阶段自适应切换(消融已暗示失配动态是阶段依赖的,前期用 TIS、后期切 TRIAGE 或反之可能更优);在多模态或 agent 长轨迹任务上检验段级诊断是否捕捉到类似的局部风险结构。
复现评估
开源情况较好:代码在 GitHub(InfiXAI/TRIAGE)公开,训练数据为公开的 DAPO 数学数据集,模型为公开的 Qwen3-4B-Base 和 Qwen3-30B-A3B-Base,五个评测基准均为社区标准;附录 C/D/E 分别给出 TRIAGE 操作规则、训练设置和系统/吞吐细节,且所有阈值声称在训练前固定,可复现性设计意识明确。主要障碍在硬件与软件栈:结果依赖 B300(Blackwell)GPU 的原生 NVFP4 Tensor Core,需要 Transformer Engine 的 W4A4 kernel、Megatron-LM 学习器与 SGLang/vLLM 采样器的特定组合,非 Blackwell 硬件无法复现 native W4A4 主张。算力门槛很高:论文总计 34,500+ B300 GPU 小时,其中 30B 主实验(1700 步、20,480 最大响应长度、多设置对照)单实验室难以完整复现;4B 复现(600 步)成本约低一个数量级,是验证核心主张的现实起点。综合评估:文档与代码基础扎实,但完整复现属于高难度,建议先复现 4B 稳定性曲线(Figure 4)和 300 步消融(Figure 5)。
论文图表
对 Qwen3-4B-Base 和 Qwen3-30B-A3B-Base 的 native NVFP4 GRPO 训练,在 25%/50%/75%/100% 处取 25-update 滑窗(healthy/drift/pre-terminal/terminal 四阶段),对比两个指标:上图为两个放大区 $\mathcal{A}^-$/$\mathcal{A}^+$ 的优势加权 token 占比(directional mass),下图为失配 gap 的 5% 分位深度 $-q_{0.05}(\delta)$(对数轴)。结果显示 $\rho_{asym}$ 在 drift 到 pre-terminal 间从 1.02 升至 2.32(4B)、从 1.34 升至 1.85(30B),而此时 $|\delta|<0.05$ 的 token 仍占 83.0%/91.9% 以上;直到 terminal 阶段负尾才爆炸至 $q_{0.05}(\delta) = -42.75$(4B)和 $-10.94$(30B)。
这张图支撑了论文最重要的实证主张:方向性失衡是可预警的先兆信号,幅度指标在先兆期完全看不到异常。它解释了为什么必须换掉 TIS 这类幅度式干预,也定义了 TRIAGE 门控应介入的时间窗口。
两部分:(a) 上图显示把负优势响应切成 64-token 段后,按段级尾部占比排序的 top 5% 段所包含的放大尾 token 份额——从 healthy 到 pre-terminal,虽全局重尾段频率从约 4.8–4.9% 降至 2.8–3.6%,top 5% 段仍集中约四分之一(25.0%/25.7%)的尾部 token,约为均匀放置的 5 倍,且相邻重尾持续性从约 2.6 倍升至 3.8–4.0 倍;(b) 下图显示含重尾段的响应中有多少的平均失配 $|\bar{\delta}_i|$ 仍小于 0.05——pre-terminal 时高达 97.22%(4B)和 79.94%(30B),terminal 时降至 2.61% 和 0.36%。
它证明了干预粒度选择的核心依据:风险是「少部分段内的持久局部聚集」而非全局扩散,响应级平均在恶化前完全失明。没有这张图,就无法论证为什么 TRIAGE 选择段级诊断而不是 token 或响应级。