4-Bitter的教训：在NVFP4 RL中平衡稳定性与性能

# 四点式教训：NVFP4强化学习中的稳定性与性能平衡 来源：https://humansand.ai/blog/nvfp4-rl ## 引言 **RL训练模拟器。** 该模拟器模拟了一个异步强化学习系统，其中采样器持续生成轨迹，而训练器则更新策略。策略失配源于**离线策略性**（陈旧轨迹）和**量化误差**，这些误差累积成策略漂移，最终如果超出优化器的修正能力，就会导致奖励下降。 **RL与效率旋钮。** 离线策略、权重同步、批次大小和视界控制着异步程度和策略陈旧度。MXFP8和NVFP4提升了训练和轨迹生成效率，而反量化反向传播、BF16最后15%以及共享专家则提升了数值稳定性。 **如何使用。** 通过改变算法和系统旋钮来探索吞吐量与稳定性之间的权衡。增加异步性或降低精度以提升利用率，然后观察稳定化技术如何恢复稳定性裕度。目标是找到能够最大化吞吐量、同时将策略失配保持在临界阈值（超过该阈值漂移加速且奖励崩溃）以下的配置。 强化学习的本质是通过行动与后果来训练模型；在**humans&**，我们使用强化学习来训练模型，使其理解与人交互的长期影响。在强化学习循环¹中（¹具体来说是一种策略梯度算法），模型采取行动，获得奖励，然后更新行动的可能性。然而，在实际的LM强化学习训练中，我们希望一旦观测到数据就立即更新策略，即使其他轨迹仍在采样中。对于实现我们使命至关重要的长视域多人轨迹，我们甚至可能在模型完成一次轨迹的同时进行数十次训练步骤。 这就引入了强化学习中吞吐量与稳定性之间的拉锯战。一方面，我们希望尽可能快速地在大量样本上进行训练²（²使每个训练步骤更快、每个轨迹更快，或者让训练器和采样器更重叠）。另一方面，大多数提升吞吐量的技术会导致采样策略与训练策略之间的差异，可能减缓学习并破坏训练稳定性。量化就是这种权衡的典型例子：低精度格式虽然在硬件上能实现更快的通信和计算，但会损害稳定性。像NVFP4这样快速且精确的低精度格式，配合硬件支持³（³在NVIDIA Rubin GPU上，每秒操作数比16位训练多9倍。来自NVIDIA HGX平台的密集Tensor Core PFLOPS每GPU规格（B200/B300来自8-GPU HGX系统，Rubin来自HGX Rubin NVL8表）：GPU BF16 FP8 FP4；B200 2.25 4.5 9；B300 2.25 4.5 13.5；Rubin (NVL8) 4 17.5 35），在模型训练和推理领域各自都带来了巨大的吞吐量提升。 然而，目前尚无稳定且硬件原生的4位⁴（⁴虽然有INT4 QAT方案来模拟4位量化，但它们使用16位激活（Kimi K2 Thinking），而DSv4中的MXFP4方案则使用了MXFP8激活（DeepSeek-V4 on Day 0）。我们的方法对权重和激活均采用NVFP4量化。）强化学习配方在开源社区中可用，主要原因是采样和训练的不稳定性在强化学习中会相互叠加。通过与开源社区的长期合作，我们开发并**分享**了https://github.com/radixark/miles/issues/615一个低精度强化学习配方，该配方保留了高精度训练的动态特性。在这个配方中，我们需要解决因策略量化误差导致的前向传播不稳定性、因梯度不匹配导致的反向传播不稳定性，以及因一小部分特别敏感的权重而出现的前向与反向交叉问题。下面，我们将解释如何解决每个问题并验证最终配方。 *注：这项工作离不开我们在RadixArk和NVIDIA的出色合作者，以及他们在训练、推理和强化学习栈上的工作。* ## 基线：一个具有稳定训练动态的起始配方 为了一致比较，除非另有说明，本报告中的所有实验均使用在DAPO-math-17k数据集上以8k序列长度训练的Qwen3-30B-A3B模型。 我们从使用NVFP4格式的基线配方开始。NVFP4是一种4位浮点格式，可在较新的NVIDIA GPU架构上提升吞吐量和内存效率。它采用分层块缩放：每16个值的块有一个FP8 E4M3缩放因子，而整个张量则有一个全局FP32缩放因子。这些缩放因子共同可以恢复更高精度的值。 ### 为什么NVFP4预训练配方不足以用于强化学习？ NVIDIA的NVFP4预训练配方是这项工作的自然起点。它采用混合精度策略，其中大多数操作以FP4精度执行，而数值敏感部分则保持在更高精度。 然而，预训练和强化学习有不同的失败模式。 在预训练中，梯度信号密集且在大规模token上重复平均。主要目标是避免会扰乱优化方向的量化偏差。随机舍入通过使梯度量化近似无偏来提供帮助。 强化学习环境有不同偏差-方差权衡。策略梯度本身就是一个有噪估计器，因为它依赖于采样轨迹、优势估计、奖励估计、KL正则化和策略陈旧度。因此，量化方法仅仅无偏是不够的。量化噪声必须足够小，以免降低每次更新的真实策略梯度信号。 因此，我们优先考虑那些能减少策略量化误差并提高梯度计算精度的干预措施。 ### 基线配方 对于基线配方，我们仅量化MoE层，并将所有其他层保持在较高精度的BF16。在DeepSeek-V3风格的架构中，MoE专家约占参数总量的97%⁵（⁵每个专家包含3个大小为7168x2048的投影矩阵，每个专家约4400万参数。若每个MOE层有256个专家外加1个共享专家，则每层约113.2亿参数。58层总计约6560亿参数，约占全部参数的97.8%）。因此，对MoE层进行激进量化可带来大部分内存收益。 我们以前向传播使用NVFP4，同时反向传播保持BF16精度。这是一个保守的起点，既能获得FP4在轨迹生成和内存方面的好处，又避免了在反向传播中使用FP4。 对于权重，我们使用标准的NVFP4格式，包含FP8每块缩放因子和一个FP32全局缩放因子。 对于激活，我们不使用基于整个张量计算的FP32全局缩放因子。正如Cursor Composer 2技术报告所指出的，全局缩放因子可能造成两个问题： - 同一个token可能因批次中其他token的不同而受到不同量化。 - 如果张量包含同一序列的多个位置，后续token可能影响前面token使用的共享缩放因子，从而形成从未来token到过去token的泄漏路径。 为了解决这个问题，类似于Cursor Composer 2，我们使用每token激活缩放，其中每个token沿隐藏维度计算自己的FP32激活缩放因子。这使得量化局部化到每个token，并避免了单独的校准步骤，因为缩放因子可以直接从前向传播中观察到的激活计算得出。此外，这种细粒度的FP32缩放因子相比整个张量的FP32缩放因子有助于降低量化误差。 在轨迹生成期间，这个每token缩放因子的计算被融合到激活量化内核中。对于每个token行，内核计算FP32缩放因子、16值块的FP8 E4M3缩放因子，并将这些量化值打包成FP4。这种融合实现减少了额外的内存移动和内核启动开销。 ## 改进梯度稳定性 在使用我们开源每token NVFP4配方的训练过程中，我们观察到了梯度范数尖峰。这个配方有意在反向传播中保持精度保守：尽管前向传播使用NVFP4操作数，但反向传播将权重和激活保持在BF16精度。这比通过粗粒度NVFP4运行完整反向路径要稳定得多，但它引入了前向与反向传播之间的不匹配，导致偶尔的梯度尖峰。 问题在于反向传播不再微分前向传播所使用的同一函数。例如，在一个简单的线性层中，前向传播计算`y = x · Q(w_bf16)`，其中Q(·)表示量化函数⁶（⁶量化函数的行为类似于阶梯函数，在阶跃变化值处不可微。量化通常涉及裁剪和取整，即clip(round(w/s), qmin, qmax)）。反向传播则表现得如同它在微分`y = x · w_bf16`。这样的决定使反向传播更稳定，但也意味着反向传播不了解量化函数中进行的裁剪和取整决策。 为了解决这种链式法则不一致性，同时保留BF16反向传播的稳定性，我们采用了反量化反向传播。反向传播不再使用BF16精度，而是使用前向传播中实际使用的量化张量的BF16反量化值。以线性层为例，这意味着`y = x · DQ(Q(w_bf16))`，其中DQ表示反量化。 此处，反向传播仍然使用BF16操作数，但这些操作数反映了与前向传播相同的NVFP4量化决策。请注意，我们的配方对权重和激活都应用了反量化反向传播操作。 ### 结果 **MXFP8验证：** 为了验证我们的更改，我们首先在MXFP8配方上测试了这个修复。下图显示了在三种不同配方下梯度范数的行为： - a) 量化前向和反向，表示为mxfp8-mxfp8 - b) 量化前向和BF16反向，表示为mxfp8-high-precision-bf16 - c) 量化前向，并通过在反向传播中使用反量化权重来修正链式法则违规，表示为mxfp8-dequantized-bf16 我们的方法在整个训练阶段产生的梯度比MXFP8反向更干净、噪声更小。与高精度BF16反向相比，我们的方法有助于防止梯度范数增加，并且始终稳定。 **NVFP4验证：** 与MXFP8相比，NVFP4使训练动态变得更加复杂。一方面，我们的方法通过使反向传播与量化前向匹配，提供了偏差更小的梯度，有助于稳定性。另一方面，由于其更粗糙的表示网格，NVFP4量化可能增加梯度方差。 为了从经验上理解这一点，我们的第一个设置使用SGD优化器而不是Adam，因为SGD更容易隔离我们的方法对真实梯度的影响。 我们观察到，我们的方法导致比高精度反向基线更快的学习。这可以解释为梯度与实际被训练的量化函数更对齐。 然而，我们也观察到这个配方稳定性较差，可能发散。如前所述，这很可能是因为NVFP4较粗糙的表示导致梯度具有更高的方差。 为了复制真实强化学习运行的设置，我们切换回Adam优化器，并检查动量、自适应二阶矩等属性是否减少了我们方法引入的额外方差。 我们观察到，与在反向传播中使用更高精度的基线方法相比，我们的方法产生的梯度大尖峰要少得多。 ## 强化学习权重与激活的四分之六方案 另一个挑战与NVFP4本身有关。通常，我们可以将参数表示为0、±0.5、±1、±1.5、±2、±3、±4、±6的倍数。然而，这意味着在表示接近范围5/6的值时，误差可能高达范围的1/6！如果我们将最大值映射到±4，最大可能误差就变成范围的1/8。这种方法称为四分之六，或4/6，源自论文《Four Over Six: More Accurate NVFP4 Quantization with Adaptive Block Scaling》。关键的是，仅当能减少量化误差时才引入缩小的范围。在该论文中，这种技术仅应用于预训练且仅用于激活；然而，我们将其也应用于权重。这对强化学习更为重要，因为我们从已经训练好的权重开始，其量化误差会产生不成比例的影响⁸（⁸我们为权重和激活（或两者）添加了4/6的支持）。 对于自适应缩放，我们需要计算调整后的权重，并确定每个块的误差是否减少。 然而，计算调整后的权重⁹（⁹如参考实现）的参考方法和最直接的方法开销很大，即： `candidate = fp4_value * fp8_scale * amax / (E2M1_MAXVAL * E4M3_4OVER6_MAXVAL)` 为了选择我们想要的FP4最大值（4或6），我们： - 将每个存储的FP4值转换为FP32 - 乘以块缩放因子和张量/token最大值 - 通过除以FP4/FP8最大值来归一化 - 计算4和6的块级量化误差 - 比较块量化误差¹⁰（¹⁰均方误差（MSE）或平均绝对误差（MAE））来选择4或6 问题在于所有这些操作都很复杂...