[Bug] Async oversampling checkpoints drop buffered rollouts on resume

[Bug] Async oversampling checkpoints drop buffered rollouts on resume

Severity

Medium-high. A normal checkpoint reports success and restores model/optimizer/dataloader state, but silently skips generated prompts that had not yet been trained.

Affected revision

main at 3c3be6234e0cb353e76bb8019947db9dfe99fca7 (v0.11.0).

Preconditions

• --train.async_enable
• --rollout.vllm_generate_batch_size > --rollout.batch_size

This is the documented oversampling configuration in the README.

Summary

SamplesGenerator.generate_samples() advances the stateful prompt dataloader by the full vLLM generation batch and keeps extra Experience objects in _sample_buffer. Async checkpoint client state saves only prompts_dataloader.state_dict(). On restart the dataloader resumes after every generated prompt, while _sample_buffer is initialized empty, so all generated-but-untrained rollouts are permanently skipped.

The checkpoint also lacks an explicit dataloader-exhausted bit. Recreating an iterator from a finished StatefulDataLoader can start the same episode again, especially when the dataset ends exactly on a generation-batch boundary and no final StopIteration probe was needed.

Minimal reproduction

Use a stateful loader containing IDs 0..7, rollout.batch_size=2, n_samples_per_prompt=1, and vllm_generate_batch_size=4:

1. The first generation advances the loader through ID 3.
2. The first training chunk is [0, 1]; _sample_buffer retains [2, 3].
3. Save the current loader state (position=4) as async training does.
4. Construct a new generator, load that state, and generate the next chunk.

Current result:

first_trained_chunk=[0, 1]
checkpointed_loader_position=4
live_buffer_at_checkpoint=[2, 3]
checkpointed_buffer=[]
first_chunk_after_resume=[4, 5]

Expected first_chunk_after_resume is [2, 3]. PR #1323's committed regression suite covers buffer capture, checkpoint materialization, restore, exact boundaries, and exhausted-buffer variants.

Root cause

The generator sends this checkpoint state with each produced rollout:

client_states = {
    "episode": episode,
    "total_consumed_prompts": total_consumed_prompts,
    "data_loader_state_dict": self.prompts_dataloader.state_dict(),
}

However, the state required to continue from that exact training boundary is the dataloader state, the remaining _sample_buffer, and whether the current episode's dataloader is already exhausted. GenerateSamplesActor.load_state_dict() restores only the loader cursor.

Because async generation can run ahead of training, checkpoint state must stay associated with each queued rollout. Reading the generator's current buffer only when the trainer saves would capture a later boundary.

Impact

• Resume silently changes the training data stream instead of failing.
• Up to almost vllm_generate_batch_size * n_samples_per_prompt rollouts can be lost at one checkpoint boundary.
• If the oversampled generation drained the dataset, a resumed episode can finish immediately even though most generated samples were never optimized.
• At an exact generation boundary, recreating the finished loader can instead repeat already-trained prompts from the same episode.
• Prompt accounting still includes skipped prompts, so logs do not expose the loss.

Duplicate search

I searched all issue and PR states for oversampling, vllm_generate_batch_size, sample_buffer, async checkpoint, resume, dataloader, and skipped samples.

• Merged #1152 introduced the current async sampler/checkpoint restore path, but its client state contains only the loader cursor and prompt counters.
• #1103 and merged #1124 concern progress-bar display after a correctly restored dataloader; the reporter confirmed that the loader state itself did load.
• Open #1179 proposes an earlier oversampling design but does not persist the current _sample_buffer state.
• #1297 / merged #1300 fixes a non-empty terminal batch being discarded during uninterrupted synchronous training, not state lost across async checkpoint restore.

No existing issue or PR covers loss of the current _sample_buffer across async checkpoint restore.