Rounding alignment improves stability in large language model training
Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning
Machine Learning
Summary
Training large language models with reinforcement learning can become unstable because small differences in how numbers are rounded during calculations make the updates behave unexpectedly. The authors found that carefully matching where rounding happens during training and inference helps keep the training stable, especially by focusing on three specific parts of the model’s calculations. This approach avoids changing big parts of the system or the learning rules and still stops the model training from collapsing.
What this means in practice
- •For machine learning engineers: Prevent stability issues in reinforcement learning training of large language models by aligning rounding steps in crucial model components.
- •For hardware acceleration developers: Improve compatibility and reliability between inference and training hardware by controlling rounding behaviors in specific tensor operations.
Authors
Yijing Ke, Tianle Cai, Ruosong Wang
Abstract
Reinforcement 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. We ask whether aligning selected parts of the execution stack can stabilize BF16 RL while other implementation differences remain. We align training-side rounding locations with inference in three regions---RoPE, SwiGLU, and residual RMSNorm---by retaining selected intermediates in FP32 until their BF16 outputs. Across four Qwen-family models, aligning all three regions prevents the collapse observed under the default BF16 baseline; RoPE-path alignment alone also suffices in the studied runs, despite remaining numerical discrepancies. The interventions keep the inference engine, matrix-multiplication and attention kernels, and learning objective unchanged.