Aligning BF16 rounding locations stabilizes LLM reinforcement learning
Original title: Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning
Keeping RoPE, SwiGLU, and residual RMSNorm intermediates in FP32 prevented collapse on four Qwen-family models.
The paper studies how separate inference and training stacks can turn an intended on-policy BF16 update off-policy even with identical weights. Instead of aligning whole stacks or changing the learning algorithm, the authors match training-side rounding to inference in RoPE, SwiGLU, and residual RMSNorm by holding selected intermediates in FP32. Across four Qwen-family models, aligning all three regions prevented the collapse seen with the default BF16 baseline. RoPE-path alignment alone also sufficed in the studied runs, while the inference engine, kernels, and objective stayed unchanged.
Unconfirmed: The paper does not state whether RoPE-only alignment holds outside the studied runs.
- Only selected rounding locations are aligned, not global precision.
- RoPE alignment alone was enough in the runs reported.
- Matrix-multiplication and attention kernels remain unchanged.
- Other numerical discrepancies between stacks can remain.
From the original article
1 min read at arxiv.orgReinforcement learning for large language models often uses separate systems for inference and training. Even with identical weights and BF16 tensors, implementation differences can make an intended on-policy update off-policy and destabilize training. Existing remedies align the entire inference and training stacks, change global precision, or modify the learning algorithm.
By Yijing Ke · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.