TLDR: Reinforcement learning (RL) fine-tuning of large language models (LLMs) often suffers from instability due to a numerical mismatch between training and inference. This paper reveals that the root cause is the low precision of the BF16 floating-point format, which introduces significant rounding errors. The authors demonstrate that simply switching to FP16, with its higher numerical precision, effectively eliminates this mismatch. This change leads to more stable optimization, faster convergence, and stronger performance across diverse tasks, algorithms, and models, without requiring complex algorithmic or engineering modifications. The findings advocate for reconsidering FP16 as a foundational option for robust RL fine-tuning.
Reinforcement learning (RL) is a powerful technique used to fine-tune large language models (LLMs), helping them achieve better reasoning capabilities. However, this process often faces significant instability, making it challenging to reliably improve model performance. A key reason for this instability is a fundamental discrepancy known as the training-inference mismatch.
Modern RL frameworks typically use different computational engines for fast inference (generating responses) and for training (calculating gradients). While these engines are designed to be mathematically identical, subtle numerical differences, often due to precision errors and hardware optimizations, cause them to produce slightly different outputs. This seemingly minor mismatch can lead to biased gradients during training and a performance gap when the model is deployed.
Previous attempts to solve this problem have included complex algorithmic corrections, such as various forms of importance sampling, and engineering alignments. These solutions, however, come with their own set of drawbacks. Algorithmic fixes can be computationally inefficient, requiring extra processing steps, and they don’t fully close the ‘deployment gap’ – meaning the model optimized for training might not be truly optimal for real-world inference. Engineering solutions, on the other hand, demand deep technical knowledge and significant effort, and their effectiveness might not generalize across different systems.
This research paper, titled Defeating the Training-Inference Mismatch via FP16, takes a step back from these complex fixes and investigates the root cause of the numerical mismatch: floating-point precision. The authors, Penghui Qi, Zichen Liu, Xiangxin Zhou, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin, identify that the widely adopted BFloat16 (BF16) format is the primary culprit. While BF16 is excellent for pre-training LLMs due to its wide dynamic range, its lower precision makes it highly susceptible to rounding errors. These errors accumulate during the autoregressive sampling of tokens, causing the training and inference policies to diverge.
The Simple Solution: Reverting to FP16
The paper’s key finding is remarkably simple: by switching from BF16 to Float16 (FP16) during RL fine-tuning, the training-inference mismatch can be virtually eliminated. FP16 offers higher numerical precision because it allocates more bits to the mantissa (the part of a floating-point number that determines precision) compared to BF16. This higher fidelity means that the outputs of the training and inference engines are much more likely to be numerically identical, creating a buffer that absorbs minor implementation differences and prevents rounding errors from accumulating.
While FP16 has a more limited dynamic range and can be prone to issues like gradient underflow, these challenges are effectively addressed by mature techniques like loss scaling, which are standard components in modern training frameworks. Enabling FP16 typically requires only a few lines of code changes, making it a straightforward and robust solution.
Also Read:
- FALQON: Speeding Up LLM Fine-tuning with Merged Low-Bit Adapters
- Smarter LLM Finetuning: Introducing BOTS for Dynamic Task Selection
Empirical Evidence and Broad Impact
The researchers conducted extensive experiments to validate their findings across diverse settings:
- Offline Analysis: FP16 significantly reduced the mismatch in token and sequence-level probability distributions compared to BF16.
- Sanity Tests: In controlled environments, FP16 training runs were dramatically more stable, converged faster, and achieved substantially higher rewards and evaluation scores across various RL algorithms (including GRPO, GSPO, and policy gradient methods). In contrast, BF16 methods often collapsed early in training.
- Ablation Study: Using FP16 for both training and inference yielded the best results in terms of stability and efficiency, outperforming combinations that used BF16 or FP32 for inference (where FP32 was stable but impractically slow).
- Generalization: The benefits of FP16 extended to more complex scenarios, including Mixture-of-Experts (MoE) RL, Low-Rank Adaptation (LoRA) RL, and fine-tuning on larger dense models and different model families like OctoThinker. In all these cases, FP16 led to greater stability and consistently higher performance.
The paper concludes that for RL fine-tuning, the extreme dynamic range of BF16 is less critical, while the precision it sacrifices becomes a dominant drawback. By trading BF16’s unnecessary range for FP16’s critical precision, the gap between training and inference is effectively closed without complex algorithmic or engineering workarounds. This work suggests a broader reconsideration of precision trade-offs in RL fine-tuning, advocating for FP16 as a powerful and often more suitable alternative for stabilizing this crucial phase of LLM development.
