Skip to main content
Back to timeline
arXivSource publication:

BF16 catches FP16 in RL post-training: stochastic rounding stops small weight updates from sticking at low learning rates

Synopsis

At the low learning rates common in reinforcement learning with verifiable rewards, mixed-precision BF16 training converges more slowly than FP16 because small updates to FP32 shadow weights are often erased when those weights are rounded to BF16, leaving the weights unchanged for several steps. Replacing round-to-nearest with stochastic rounding in that quantization step makes BF16 follow FP16's reward trajectory on 1.5B, 14B and 30B-A3B models across two math datasets, while drifting less from the reference model than BF16 at a raised learning rate and keeping BF16's dynamic range without loss scaling.

Source-provided article image: Closing the BF16-FP16 Convergence Speed Gap in RL Post-Training with Stochastic Rounding
arXiv

Interpretation

At the low learning rates common in reinforcement learning with verifiable rewards (RLVR), mixed-precision BF16 training needs more iterations than FP16 to reach the same reward or downstream score, and the gap persists when the rollout and training backends produce bit-exact log-probabilities for the same weights and inputs. BF16 had been treated as the natural choice for RL post-training because it shares FP32's dynamic range and converges similarly to FP32 in pretraining; BF16's slower RLVR convergence had been attributed to training-inference mismatch. With vLLM's TorchInductor optimizations disabled and training-inference numerical parity enforced, FP16 still reached a given reward level in fewer RL iterations than BF16 at the same learning rate, a comparison run on 64 NVIDIA H100 GPUs.

The gap comes from weight sticking: when FP32 shadow weights are quantized to BF16 by round-to-nearest, an update smaller than half the spacing between representable values leaves the BF16 weights used for rollout generation and training unchanged, so several optimizer steps may pass before it shows up. BF16 has 7 mantissa bits against FP16's 10, so within the normal range shared by both formats the BF16 spacing is 8 times the FP16 spacing, and the same update moves BF16 weights less easily. Across the four optimizer steps of the first RL iteration on DeepSeek-R1-1.5B, the share of unchanged BF16 weights fell from 99.42% to 98.84% while FP16 fell from 95.35% to 90.79%, so by the end of the iteration FP16 had changed roughly 8 times as many weights as BF16 (9.21% versus 1.16%).

Switching the shadow-weight-to-BF16 quantization to stochastic rounding makes BF16's reward trajectory match FP16's, with less drift from the reference model than BF16 at a higher learning rate. Raising the BF16 learning rate also recovers FP16's convergence speed, but it increases the policy's KL divergence from the reference model, and such drift is commonly penalized in RL post-training. BF16 with stochastic rounding matched the FP16 reward trajectory on 1.5B, 14B and 30B-A3B models with MATH-1460 and DAPO-Math-17k; for mean rewards above 0.7, BF16 at the higher learning rate drifted about as much as FP16, while BF16 with stochastic rounding stayed close to FP16.

Stochastic rounding only needs to act on the weights used in forward passes: one rounding realization is drawn after each optimizer step, both backends use those BF16 weights for every forward pass until the next step, and all other conversions keep round-to-nearest. Several stochastic-rounding variants were compared on DeepSeek-R1-1.5B with MATH-1460, including redrawing every forward pass, also applying it to backward-pass weights or gradients, and applying it only in the rollout backend. Redrawing every forward pass gave the same reward trajectory as drawing once per step, adding stochastic rounding to backward weights or gradients brought no further speedup, and applying it only in the rollout backend left mean reward plateauing well below FP16's.

Perspective

The demonstrated scope is reinforcement learning with verifiable rewards, the listed DeepSeek-R1 distilled and Qwen3 models, the MATH-1460 and DAPO-Math-17k datasets, and a training setup built on verl, vLLM and Hugging Face Transformers; enforcing numerical parity reduces throughput, so it is reserved for the BF16-FP16 comparison and the weight-sticking measurements.

On Qwen3-30B-A3B, BF16 with stochastic rounding produces responses about 25% shorter than the other two configurations, a behavior left to future work; in preliminary FP16 experiments on an agentic software engineering task, training overflowed frequently even at a loss scale of 1024; stochastic rounding on low-precision formats such as MXFP8 is supported only by preliminary asynchronous RLVR experiments, with a detailed study still to come.