Stabilizing Reinforcement Learning with LLMs: Formulation and Practices

Abstract: This paper proposes a novel formulation for reinforcement learning (RL) with LLMs, explaining why and under what conditions the true sequence-level reward can be optimized via a surrogate token-level objective in policy gradient methods such as REINFORCE. Specifically, through a first-order approximation, we show that this surrogate becomes increasingly valid only when both the training-inference discrepancy and policy staleness are minimized. This insight provides a principled explanation for the crucial role of several widely adopted techniques in stabilizing RL training, including importance sampling correction, clipping, and particularly Routing Replay for Mixture-of-Experts (MoE) models. Through extensive experiments with a 30B MoE model totaling hundreds of thousands of GPU hours, we show that for on-policy training, the basic policy gradient algorithm with importance sampling correction achieves the highest training stability. When off-policy updates are introduced to accelerate convergence, combining clipping and Routing Replay becomes essential to mitigate the instability caused by policy staleness. Notably, once training is stabilized, prolonged optimization consistently yields comparable final performance regardless of cold-start initialization. We hope that the shared insights and the developed recipes for stable RL training will facilitate future research.

• The paper introduces a novel formulation that optimizes sequence-level rewards via surrogate token-level objectives to enhance RL stability in LLMs.
• It employs importance sampling corrections and clipping to address training-inference discrepancies and policy staleness in Mixture-of-Experts models.
• Empirical studies with a 30B MoE model validate that techniques like Routing Replay effectively accelerate convergence while maintaining consistent performance.

Stabilizing Reinforcement Learning with LLMs: Formulation and Practices

Abstract

This paper introduces a new approach to reinforcement learning (RL) in the context of LLMs, explaining conditions under which sequence-level rewards can be optimized using a surrogate token-level objective through policy gradient methods like REINFORCE. By minimizing both training-inference discrepancies and policy staleness, the paper elucidates why techniques such as importance sampling and clipping are vital in stabilizing RL training. Empirical studies using a 30B Mixture-of-Experts (MoE) model indicate that basic policy gradient algorithms with importance sampling corrections achieve improved training stability in on-policy scenarios. When off-policy updates are integrated to hasten convergence, clipping and Routing Replay prove essential to counteract instability induced by policy staleness. As established through experimentation, stabilizing training promotes consistent performance across different initializations and optimization extends over time.

Introduction

Reinforcement Learning (RL) is pivotal for enhancing the capabilities of LLMs in handling complex problems. The contextual nature of language entails employing sequence-level rewards in RL settings, whereas common RL algorithms operate with token-level optimization objectives. This discrepancy raises concerns regarding training stability. The study presents a novel formulation for RL with LLMs, positing that optimizing sequence-level rewards through token-level objectives is justified under minimized training-inference discrepancy and policy staleness.

Significant contributions include demonstrating that critical techniques like Routing Replay can reduce the impact of expert routing discrepancies in MoE models, and elucidating empirical findings on stabilizing RL training across various scenarios, thereby paving the way for more stable and scalable RL applications.

Reinforcement Learning Formulation

Sequence-Level Reward Optimization

The primary goal is to maximize the expected sequence-level reward, described by $$\mathcal{J}^\text{seq}(\theta) = \mathbb{E}_{ x, y \sim \pi_\theta } \left[ R(x, y) \right],$$ a function representing the reward given a complete model response. Direct optimization of this objective poses challenges due to large numerical ranges and high variance.

Token-Level Objective as Approximation

The surrogate token-level objective is expressed as $$\mathcal{J}^\text{token}(\theta) = \mathbb{E}_{ x, y \sim \mu_{\theta_\text{old}} } \left[ \sum_{t} R(x, y) \log \pi_\theta(y_t | x, y_{<t}) \right],$$ based on a first-order approximation to the sequence-level objective. The validity of this approximation hinges on the alignment of the target and rollout policies.

Conditions for Approximation Validity

The paper identifies two critical factors: training-inference discrepancy and policy staleness. Expert routing in MoE models complicates these conditions, as routing decisions can differ between training and inference engines, further impacting the objective's validity.

Empirical Analyses

Empirical studies involve extensive experiments with a 30B MoE model, employing hundreds of thousands of GPU hours. Notable findings include the efficacy of importance sampling correction in maintaining stability during on-policy training, as seen in the results of Figure 1.

Figure 1: Results of on-policy training with ( \text{gbs (global batch size)}=\text{mbs (mini-batch size)}=1,024.

Challenges and Solutions for MoE Models

Expert routing exacerbates the training-inference discrepancy, and policy staleness due to dynamic expert selection. The Routing Replay approach is proposed to stabilize this by fixing the routed experts during policy optimization.

Routing Replay Variants

• Vanilla Routing Replay (R2): Targets policy staleness by replaying experts within the training engine.
• Rollout Routing Replay (R3): Mitigates both training-inference discrepancy and policy staleness by using experts determined in the inference engine.

Conclusion

This research advances the theoretical framework and practical methodologies for stabilizing RL in LLMs. The proposed formulation relying on token-level objectives, importance sampling, and Routing Replay, delivers empirical stability across diverse training paradigms, demonstrating consistent convergence to optimal performance limits in extended RL training phases. This paper thus lays a foundation for future explorations into more efficient and scalable RL training processes.