Beyond Precision: Training-Inference Mismatch is an Optimization Problem and Simple LR Scheduling Fixes It

TL;DR

This paper identifies training-inference mismatch in RL for large language models as an optimization issue and introduces a dynamic LR scheduling method that stabilizes training by reducing the learning rate when instability signals appear.

Contribution

It reveals the dynamic nature of training-inference mismatch as an optimization failure and proposes a simple LR scheduling fix based on response length as an early warning.

Findings

Dynamic LR scheduling stabilizes RL training.

Training-inference mismatch correlates with gradient noise.

LR decay based on response length prevents instability.

Abstract

Reinforcement Learning (RL) for training Large Language Models is notoriously unstable. While recent studies attribute this to "training inference mismatch stemming" from inconsistent hybrid engines, standard remedies, such as Importance Sampling, might fail during extended training runs. In this work, we analyze this instability through the lens of optimization, demonstrating that gradient noise and training-inference mismatch escalate in tandem as training progresses. Meanwhile, we find that the mismatch can be effectively suppressed by shrinking the update size. Taken together, we deduce that the mismatch is not merely a static numerical discrepancy, but a dynamic failure coupled with the model's optimization. Based on this insight, we propose a simple yet effective solution: a specialized Learning Rate (LR) scheduler. Instead of pre-defined decay schedule in traditional LR scheduler, our method dynamically triggers LR decay based on response length, which we identify as a reliable early-warning signal for impending instability. Empirical evidence suggests that by reducing the learning rate as gradient noise rises, we can consistently stabilize RL training and keep the training-inference mismatch at a safe level.