Training goes too off-policy#

Off-policy refers to the training regime where the model/policy is different from the policy used to sample the training examples, aka the rollout batch. One might wonder why the training needs to go off-policy in the first place. After all, if we stay on-policy, the training would be a lot easier and stable. I think there are mainly two reasons:

• Sampling rollouts are expensive due to the decoding/generation stage of the LLM inference. Unlike training or the prefill stage, this stage is almost always memory bandwidth bound in the hardware and hyperparameter regime I operate at (inference on one H100). We want to make better use of the generated responses.
• For hard tasks like Countdown, the high reward rollouts (stronger learning signals) are hard to come by initially. Having multiple training steps on them means better rollout usage and training performance.

But obviously, there are great disadvantages to off-policy training too:

• When too off-policy, the ratio between the probabilities of the latest policy and the old policy $\pi_{\theta} / \pi_{\mathrm{old}}$ can get too big and end up getting clipped in the GRPO loss function, which leads to no gradient and no training.
• When too off-policy, the ratio, which is the importance sampling factor, can get very noisy. This can make the training very unstable and collapse.

To address this noise source, one needs to carefully monitor the KL divergence and clip ratio. I also tried some heuristics for controlling how far off-policy the training can go. For example, I tried to put a threshold on KL divergence. If the training in one GRPO step reaches that threshold, that GRPO step will be stopped early. However, in practice, I noticed this didn’t really help much. Once the training becomes unstable, it almost always only takes one training step for the model to reach the KL divergence threshold, and stops the training there, effectively making the training on-policy. And at that point, even if the training stays on-policy, the policy will still collapse.

In the original DeepSeek V3 and R1 papers, there is an extra KL divergence term in the GRPO loss function, controlling how off-policy the model can get. However, it seems to have been dropped in most recent models in literature and practice. The rationale seems to be that we do want the model to move significantly away from the original model by self-improving using verifiable rewards. And the training can also be stabilized using other methods as I show below (tuning epochs, batch sizes, and learning rate schedule).

In my experiments, the most effective ways of avoiding policy collapse due to too off-policy is by controlling the number of epochs per rollout batch, and the number of training steps per epoch. For simple tasks like GSM8k, there doesn’t seem much need for the training to go too off-policy, since the model already has a fairly good success rate. In other words, the training already has plenty of good training signals and there is not much need to repeat on any particular rollouts. In my experiment, I set 2 epochs per rollout and one training step per rollout (using gradient accumulation). The training can already progress pretty smoothly.

GRPO eval reward curve for GSM8K.

The training above is based off the checkpoint from the expert iteration training, and took about 35 minutes on two H100 to go from 73% to 82%. No learning rate schedule is needed either. More details about the experiment can be found here. Here is an example question and model response

1
You are an helpful Assistant having a conversation with a User. The user asks you a question, and your job is to solve it.
2

3
You will first think about the reasoning process step by step and then provides the User with the final answer. Your reasoning process should be enclosed within the XML tags `<think>` and `</think>`, and the finnal answer within the XML tags `<answer>` and `</answer>`. In other words, format your response in the following way:
4

5
```
6
<think>
7
Your thinking process goes here
8
</think>
9
<answer> Your final answer goes here </answer>
10
```
11

12
Here is the Conversation:
13
User: In a dance class of 20 students, 20% enrolled in contemporary dance, 25% of the remaining enrolled in jazz dance, and the rest enrolled in hip-hop dance. What percentage of the entire students enrolled in hip-hop dance?
14
 Please give the final answer as a single number.
15
Assistant:
16
<think>

1
First, calculate the number of students enrolled in contemporary dance:
2
20% of 20 students = 0.20 * 20 = 4 students
3
Next, calculate the number of students remaining after accounting for the students in contemporary dance:
4
20 - 4 = 16 students
5
Now, calculate the number of students enrolled in jazz dance:
6
25% of 16 students = 0.25 * 16 = 4 students
7
Next, calculate the number of students remaining after accounting for the students in jazz dance:
8
16 - 4 = 12 students
9
These remaining students are enrolled in hip-hop dance:
10
12 students
11
Finally, calculate the percentage of students enrolled in hip-hop dance:
12
(12 / 20) * 100 = 60%
13
</think>
14
<answer> 60 </answer>

As we can see, the model’s CoT is much more clearly explained and logical.

For harder tasks like Countdown, however, going off-policy is necessary. The training performance starts pretty bad (about 0.7%). Therefore it would be very helpful if there could be more training steps on successful examples. This means the training becomes less stable, especially later in the training run when the performance is about to or has already plateaued.

Performance already plateaued#

In this stage, most of the gradients the model receives during training are from noise. If the learning rate is still too high, it can very often take too big a step in the noise direction, which then leads to policy collapse. The most effective way of controlling this that I learned from these experiments is to reduce the learning rate. One wants to reduce the learning rate enough to balance out the increased noise-to-signal ratio. I found typically the learning rate needs to be at most about 10% of the max learning rate. A cosine learning rate schedule worked for me well.

Another thing I noticed is that increasing the microbatch size also seems to reduce noise, though it probably doesn’t make much intuitive sense. Even with a microbatch size of 1 or 2, the training is probably already compute bound, so increasing microbatch size is probably not from a system perspective. On one H100, it seems the maximum microbatch size is 4 for Countdown. I started with 2 and noticed that going to 4 makes the training more stable, but I didn’t dig deep in this direction. It would be interesting to study this a bit more in the future.

Here is the learning curve for one RL run using Countdown

GRPO eval curve for Countdown.

It takes less than an hour to go from 0.6% to 51.4% on two H100, and then stabilizes there. As we can see, this is pretty much similar to the benchmarks posted by the person who created the dataset. Throughout the training run, all training log metrics stays stable, such as KL divergence and gradient norm

Approximate KL divergence during stable training. Clip fraction during stable training. Gradient norm during stable training.

For more details of this experiment, see here.