Why resuming an async PPO run throws away a whole batch of prompts instead of two

One question inside NVIDIA-NeMo/RL #3410 — feat(ppo): support async ppo. All permalinks pin the PR head 6e0362d.

Async PPO generates rollouts in the background while training runs. A background thread pulls prompt batches from the dataloader (_collection_loop) and puts finished rollouts into a replay buffer, each stamped with the training step it is meant for. A checkpoint saves two things side by side: the replay buffer, and the dataloader’s place — read straight off that same running loader (trajectory_collector.py:707, saved at ppo.py:2903). That place is already past every prompt handed to generation, including the ones whose rollouts have not come back yet, so a resume can never fetch those prompts again. This PR adds a flag that decides how many more prompts get thrown away on top of them.

What one checkpoint costs you

Rows 2 and 3 are the same crash and the same checkpoint. The only difference is one boolean.

The three lines that decide it

On resume the buffer counts how many prompt groups survived for each training step, and calls a step “incomplete” if it has fewer than a full batch.

# resuming at training step 11, batch = 7 prompt groups _prepare_for_training_step(current_step=11, num_prompts_per_step=7) target_counts = {11: 5} # 5 rollouts made it into the checkpoint incomplete = {11} # 5 < 7

replay_buffer.py:513 · replay_buffer.py:519

Then the flag picks a branch. One deletes those 5 rollouts; the other keeps them and prints a note.

if incomplete and self._drop_incomplete_targets_on_restore: self.trajectories = [] # true -> the 5 finished rollouts are gone else: print(" Incomplete target 11: 5/7") # false -> keep them, fill in the rest

replay_buffer.py:529

Either way the collector then asks the buffer how many more it needs, and the loader supplies them starting from the place the checkpoint saved — #19.

get_trajectories_needed(target_step=11, num_prompts_per_step=7) # false -> 7 - 5 = 2 -> loader hands out #19, #20 # true -> 7 - 0 = 7 -> loader hands out #19 ... #25 # either way it starts at #19: the saved place is restored as-is load_dataloader_state(dataloader, last_checkpoint_path, data_config)

replay_buffer.py:640 · ppo.py:466

Prompts #12–#18 are behind #19, so nothing can bring them back. Deleting the 5 finished rollouts therefore does not cause them to be regenerated — it causes them to be replaced by the next 7 prompts in the stream, and the compute spent on them is spent again.

Both branches end up training on 7 prompt groups, because get_trajectories_needed just asks for whatever is missing. So deleting the rows does not buy a better batch — the batch is the same size either way. It only changes which prompts fill it, and how much generation gets paid for twice.

One step, or several? Step 11 is not the only step that can be half-full when the checkpoint lands. Rollouts enter the buffer one prompt group at a time (one add() per group), and the collector keeps several future steps on the go at once — it only stops when every step in its look-ahead window is finished or already claimed (_should_pause_for_generation_limits). Nothing stops or drains generation while the checkpoint is written. The delete runs over every half-full step it finds, not just the one being resumed, so a single checkpoint can bin several batches at once — as many as the look-ahead window is wide.

Who asks for which branch

caller	value	matches `main`?
async GRPO — grpo.py:4143	false	yes
async PPO — ppo.py:131 (the default; no recipe sets it)	true	no — new here

On main, this resume path never deleted incomplete steps — it only printed them. Deletion existed, but only on the branch taken when no training step is supplied, which resume does not take (git show origin/main:nemo_rl/algorithms/async_utils/replay_buffer.py, lines 451–462). So true is behaviour this PR introduces, for PPO only.

Has anyone resumed one of these runs?

No. There is no run posted anywhere for this path — not in the PR description, not in any review thread, and no nightly test resumes an async PPO job from a checkpoint. The unit test that builds a replay buffer for PPO passes false (test_ppo.py:2361), so the shipped default is not the branch being exercised there.

The thing to ask for before merging is one async PPO run stopped mid-generation and resumed, with the resume log line reported: it prints either Dropping incomplete restored targets: 11=5/7 or Incomplete target 11: 5/7, which says directly which branch ran and how much was thrown away.

What this means. Every async PPO resume today takes the true branch, because nothing sets the flag. Each such resume skips whole batches of prompts rather than just the handful still generating, and bins rollouts already paid for. Both branches train on the same number of prompt groups, so the deletion buys nothing — whatever it keeps is a subset of what the other branch keeps. The loss repeats on every restart, so a job preempted often trains on less of its dataset than the epoch count suggests. false is what async GRPO and main both do, and it shrinks the loss to only the rollouts genuinely in flight — the smallest it can be until the saved loader place is rewound to cover them.