Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning

TL;DR

Aligning rounding locations in key components—RoPE, SwiGLU, and residual RMSNorm—between inference and training stabilizes BF16 reinforcement learning in LLMs without altering the learning algorithm or core kernels. This simple fix prevents training collapse observed in default BF16 setups, with RoPE alignment alone proving sufficient in tested scenarios.

Why It Matters

This work demonstrates that minor implementation differences in rounding can destabilize RLHF training, even with identical weights and precision. By focusing on specific rounding locations, it offers a low-cost, plug-and-play solution to stabilize training without requiring full stack alignment or precision changes.

Key Findings

• Aligning rounding locations in RoPE, SwiGLU, and RMSNorm prevents training collapse in BF16 RL.
• RoPE-path alignment alone is sufficient to stabilize training despite remaining numerical discrepancies.
• Stabilization is achieved without modifying inference engines, attention kernels, or learning objectives.