BF16在强化学习后训练里追平FP16:随机舍入让低学习率下的收敛不再被权重卡住
核心概要
在可验证奖励强化学习常用的低学习率下,BF16混合精度训练比FP16收敛更慢,原因是FP32影子权重舍入到BF16时,小幅更新常被舍入抹掉,权重连续几步不变。把这一量化步骤从就近舍入改为随机舍入后,BF16在1.5B、14B与30B-A3B模型和两个数学数据集上的奖励轨迹与FP16一致,同时相对参考模型的KL散度低于提高学习率的BF16,并保留BF16的动态范围、无需损失缩放。
深度剖析
在可验证奖励强化学习(RLVR)常用的低学习率下,BF16混合精度训练达到同一奖励或下游分数所需的迭代次数多于FP16,这一差距在训练与推理后端对相同权重和输入给出逐位一致的对数概率后依然存在。 此前BF16因与FP32动态范围相同、预训练收敛接近FP32而被视为强化学习后训练的当然选择;BF16在RLVR中收敛偏慢此前被归因于训练-推理数值不匹配。 在关闭vLLM的TorchInductor优化并强制训练-推理数值一致后,相同学习率下FP16仍在更少的RL迭代内达到给定奖励水平,该对照在64块NVIDIA H100上完成。
差距来自权重卡住:FP32影子权重按就近舍入量化到BF16时,小于半个最小可表示间隔的更新不会改变用于生成与训练的BF16权重,要累积若干优化器步才体现出来。 BF16只有7位尾数,FP16有10位,在两者共同的正规数值范围内BF16的最小间隔是FP16的8倍,因此同样的更新更难推动BF16权重。 在DeepSeek-R1-1.5B第一个RL迭代的四个优化器步中,BF16权重保持不变的比例从99.42%降到98.84%,FP16从95.35%降到90.79%,到迭代结束时FP16改动的权重约为BF16的8倍(9.21%对1.16%)。
把影子权重到BF16的量化改为随机舍入后,BF16的奖励轨迹与FP16一致,且相对参考模型的漂移小于提高学习率的BF16。 提高BF16学习率也能追平FP16的收敛速度,但会加大策略相对参考模型的KL散度,而这一漂移在RL后训练中通常被惩罚。 在1.5B、14B与30B-A3B模型和MATH-1460、DAPO-Math-17k两个数据集上,BF16加随机舍入与FP16的奖励轨迹吻合;平均奖励高于0.7时,提高学习率的BF16漂移与FP16相当,而BF16加随机舍入接近FP16。
随机舍入只需作用在前向传播所用的权重上:每个优化器步后抽一次舍入实现,两个后端直到下一步前的所有前向都用这份BF16权重,其余转换仍用就近舍入。 在DeepSeek-R1-1.5B与MATH-1460上比较了多种随机舍入变体,包括每次前向重新抽样、对反向权重或梯度也施加随机舍入、以及只在rollout后端施加。 每次前向重新抽样与每步抽一次得到相同奖励轨迹,对反向权重或梯度再加随机舍入没有进一步加速,而只在rollout后端施加随机舍入时平均奖励停在明显低于FP16的水平。
启示与展望
结论的演示范围是可验证奖励强化学习、所列DeepSeek-R1蒸馏与Qwen3模型、MATH-1460与DAPO-Math-17k数据集,以及verl、vLLM与Hugging Face Transformers组成的训练配置;强制数值一致会降低吞吐,因此只用于BF16与FP16的对照和权重卡住测量。
Qwen3-30B-A3B上BF16加随机舍入生成的回答比其他两种配置短约25%,这一行为留待后续研究;FP16在智能体软件工程任务的初步实验中即使损失缩放取1024也频繁溢出;低精度格式如MXFP8上的随机舍入只有初步的异步RLVR实验支持,详细研究尚未展开。
