Dynamic Data Snoozing

Dynamic data snoozing for efficient online RL

Skip to Main Menu

Skip to Main Content

Skip to Footer

Back to Blog

-->

Back to Blog

TL;DR

When running GRPO on verifiable rewards, dynamic sampling does a lot to stabilize the training process. However, it may also slow training down dramatically. In this blog, we introduce “dynamic data snoozing” as a simple and effective strategy for cutting compute waste without breaking training stability. By applying this method, we are able to achieve up to a 3X gain in compute efficiency with no quality degradation.

Efficient batch construction for GRPO

GRPO, the de-facto RL algorithm used today due to its effectiveness and simplicity, often runs into cases where a training input results in no learning signal. This can happen when all rollouts for that input have the same reward, so the computed advantage of each rollout is 0. Ignoring this fact when running GRPO leads to training instabilities, as our effective batch size isn’t fixed and depends on the amount of examples with a learning signal.

These 0-advantage examples cause training instabilities, and this problem was addressed in the DAPO paper with the now common-practice mechanism of dynamic sampling. With dynamic sampling, groups without a learning signal which are either too-easy or too-hard are filtered and replaced with a new input. Regardless of whether 0-advantage examples are practically ignored by the training process or explicitly replaced, they significantly slow down the training process. We create wasteful generations for them which have no impact on the model’s training.

Figure 1 shows the total number of examples used for generation during the DAPO training run, vs. the number of examples actually used for gradient computation. The training inefficiency is apparent:

Initially: Training efficiency is reasonable since many examples are in the range of not-too-easy and not-too-hard.

As training progresses: Training slows considerably by up to 16X compared to an optimal example sampler, as examples become too easy.

At the end of the day, less than 15% of the generations performed in the training run are actually used for training.

Figure 1 : Training efficiency of DAPO, using dynamic sampling. We can see that as training progresses, the gap between the number of generated batches vs. the optimal amount of generated batches (one per step) increases, leading to slower and less efficient training.

To optimize this large gap in performance, we developed a few simple methods that gave us performance gains during the training of Jamba2 .

Offline difficulty stratification

We began with offline difficulty stratification, a critical preparatory step in online RL training, as it provides a grounded understanding of task complexity and model behavior before learning begins. This mapping can inform both data filtering and curriculum building.

To create this stratification, we ran multiple completions of a few baseline models on our training data. Based on the success rate of these models on each example, we were able to partition our training data based on empirical difficulty and build a curriculum for our model: Overly-difficult examples, that would have been filtered anyway, were postponed to a later stage of training. On the other hand, overly-easy examples could be removed entirely from the training process when the time was right. Using this stratification, we could also adapt the dataset weights to upsample sections of the data with the most signal for the current model checkpoint.

While useful for manually controlling the difficulty of examples seen by the model during training, this off-policy method leads to a suboptimal categorization of example difficulty. To more effectively handle the overly-easy side of the data, we turned to an online, on-policy method.

Dynamic difficulty filtering

To optimize the removal of overly-easy examples from our training sets and to simplify the process of data rebalancing, we introduced a simple but effective heuristic to our training procedure.

Assume the reward for a given dataset is in the range of [0, 1]. An example is defined as too-easy if all rollouts in its group have a reward of 1. An example like this would be filtered out since all rollouts will have an advantage of 0. Assuming training goes well, our model should continue to get a perfect score on this example, so there’s no reason to keep showing it to our model in future batches. So, once an example gets filtered out for being too easy, we keep filtering it out every time it reappears in subsequent epochs, saving the expensive rollout generation phase. The training budget can then be reallocated for appropriately-difficult examples.

Figure 2 describes the training dynamics when we apply Dynamic difficulty filtering to the GSM8K dataset. Dynamic filtering leads to a considerable 3X gain in compute efficiency. The final test scores remain similar, with all variants reaching a score of 84%.

Figure 2: Training efficiency of GSM8K. We compare DAPO dynamic sampling (red) with our proposed methods – dynamic difficulty filtering (blue) and dynamic data snoozing with a snoozing factor of 5 (green, introduced later in the post). We see both methods greatly reduce the number of batches generated during training compared to naive dynamic sampling.

Maintaining efficiency in multi-task batch construction

While dynamic difficulty filtering provides considerable efficiency gains, it is not without its flaws. On a simple dataset like GSM8K there are no tradeoffs, but things become messier when we deal with the realistic scenario of training on multiple datasets 1 . Since different tasks are learned at different rates, the rate of filtering from each dataset is different and changes over time. This makes the balance of datasets within a batch harder to control, and this can be detrimental to stability when using online difficulty filtering. We needed to find a way to maintain our efficiency gains, while maintaining stability across multiple datasets.

Task starvation and mode collapse: Risks of dynamic difficulty filtering

In Figure 3 we can see a math dataset stratified by empirical difficulty into easy and hard subsets. The two datasets are equally-weighted when sampling before batch construction. The figure describes the fraction of each dataset in the batch in each step. First, due to the different group-filtering rates, we observe how they are not...