Skip to content

fix(async-grpo): eliminate prompt skips across checkpoint save/restore - #3599

Merged
yfw merged 13 commits into
mainfrom
yifu/replay-buffer-prompt-skips
Aug 25, 2026
Merged

fix(async-grpo): eliminate prompt skips across checkpoint save/restore#3599
yfw merged 13 commits into
mainfrom
yifu/replay-buffer-prompt-skips

Conversation

@yfw

@yfw yfw commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Makes async-GRPO checkpoint resume lossless: no prompt is skipped and none is duplicated across a save/restore cycle, for both load_replay_buffer=true and false.

Problem

On a SWE E2E run, 675 of 736 skipped prompts (~30% of the reached dataset) were traced to two mechanisms:

  1. Non-atomic checkpoint. The saved dataloader cursor runs ahead of what training consumed: the cursor advances on every yield, the buffer holds only completed groups, and in-flight work is serialized nowhere. Each restore lost a ~96-prompt window. The loss is also length-biased — a group completes only when its longest rollout finishes, so the lost in-flight set is exactly the long-rollout prompts.
  2. Gap-fill tail discard. When a target needed fewer prompts than a dataloader batch holds, _process_batch sliced off what it needed and dropped the rest (61 of the 736 skips).

Fix (one commit each)

  1. Carry over gap-fill tails (b1b196a1e) — _process_batch returns the unconsumed remainder; the loop consumes it before the next pull, and it is serialized in rollouts.pt so a checkpoint can't strand it.
  2. Frontier-aligned checkpointing (4edd6976) — every yielded prompt is stamped with a monotonic stream ordinal; the collector keeps a small ring of pre-pull dataloader snapshots keyed by ordinal. Checkpoints persist the snapshot at the trained frontier instead of the live cursor, plus frontier metadata. On resume the dataloader re-yields the covered window; rows already trained or retained in the restored buffer are dropped at yield, and the existing gap-fill path regenerates exactly the remainder. Un-stampable batches and legacy (pre-frontier) checkpoints fall back to today's live-cursor behavior.
  3. checkpointing.load_replay_buffer (a72ebf48) — set false to skip the buffer restore and regenerate the whole buffered window fresh, removing the short-rollout bias of retained groups. With frontier alignment this is just the empty-retained-set case of the same resume path — no second code path, and no prompts are skipped either way. Default remains true. Supersedes feat(grpo): optionally skip replay buffer restore #3583.
  4. A/B regression test (aee3c801) — drives the legacy and frontier resume paths through the same scenario: legacy loses prompts [5, 8, 9, 11], frontier regenerates all of them with zero duplicates.
  5. Self-review fixes (a482c6c5, 8c1e170d) — frontier derived from the sampled groups' own ordinals instead of consumed_samples; dataloader snapshot and pending batch read in one actor call under the pending lock; native group ordinals captured from the rollout's input rows; shared rollout-state builder; fallback warnings, scoped docs, and rollouts.pt contract tests.
  6. Conservative checkpoint cut (3b32a8b6) — thanks to @terrykong for the failure scenario: a target refilled from later prompts (tolerated failure, or a resume gap-filling an incomplete restored target) can train past another target's in-flight groups, and a frontier-aligned save would strand them. Checkpoints now cut at min(lowest in-flight ordinal, trained frontier) and persist the ordinals already trained above the cut, so a resume regenerates exactly what was lost — no skip, no re-training. Healthy path unchanged (cut == frontier, empty trained list); every step mutation-tested.

Guarantees and caveats

  • This buys data determinism (exact prompt coverage and order), not bitwise reproducibility: regenerated trajectories are fresh samples under post-resume weights (in-flight vLLM state is not serializable). Retained groups are reused bitwise.
  • With async_grpo.max_generation_failures > 0, prompts dropped by a tolerated failure are lost during normal operation — that loss happens at failure time, not at resume, and a resume neither recovers them nor (with the ordinal-derived frontier) re-trains or double-counts anything around them. The default (0) fails fast and is unaffected.
  • Resuming from a pre-upgrade checkpoint disables frontier alignment permanently for that checkpoint lineage — the run and every run descended from its checkpoints stay on live-cursor semantics (a warning is logged at resume). Fresh runs are protected from their first checkpoint.

Validation (E2E kill/resume on a production workload)

Nemotron nano-v3.5 RLVR via NeMo-Gym (SWE/agentic + multilingual chat blend), 512-prompt slice, shuffle: false, 32 prompts/step, checkpoint every step. Jobs were scanceled mid-generation and relaunched under the same lineage; the train_data_step*.jsonl dumps were then verified against the dataset's expected stream order — every step must train exactly its 32-prompt window, each prompt exactly num_generations times.

  • Run A: 14 steps across 3 job segments with 2 mid-generation kills (6 steps re-executed after kills landed before checkpoint finalization) → 448/448 prompts, no skip, no duplicate, and every re-executed step trained the identical prompt set. The resume logs show the mechanism working: Frontier-aligned resume: dataloader rewound to ordinal 64, trained frontier 64, 13 retained prompt groupsdropping 13 already-covered prompts (ordinals 64–95), regenerating 19 — those 19 in-flight prompts are exactly what the pre-PR code silently skipped.
  • Run D (with the conservative cut, max_generation_failures=3 as in the SWE recipes): 6 steps, 2 kills → 192/192, zero Checkpoint cut lowered warnings and zero worker failures across all segments — confirming the cut is inert on healthy runs.

Baseline for the same workload class before this PR: ~30% of reached prompts skipped across restarts (675/736, from the original findings).

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • Credit to @pjin-nvidia who identified this issue and provided reference implementation. And @snowmanwwg for running related test showing impact of saving partial replay buffer.

yfw and others added 4 commits August 11, 2026 14:24
…g them

When the replay buffer needs fewer prompts than a dataloader batch holds,
the collector sliced off what it needed and dropped the rest: with 32
prompts yielded and 2 needed, 30 prompts silently vanished from training.
Measured on a SWE E2E run, these discarded tails accounted for 61 of the
736 skipped prompts (30 after step 20, 31 after step 34; see the
prompt-skip findings report).

Keep the unconsumed suffix in a pending slot instead:

- _process_batch returns the unconsumed remainder (the sliced-off tail,
  or the whole batch when no target currently needs generation) rather
  than dropping it.
- _collection_loop consumes the pending remainder before the next
  dataloader pull, so the prompt stream stays strictly ordered and
  lossless. A batch that made no progress is retried after a short yield
  instead of being discarded.
- get_rollouts_state serializes the pending remainder under
  PENDING_PROMPTS_KEY, and async_grpo_train hands it back to the
  collector on resume. Without this, a checkpoint taken while a tail was
  pending would strand prompts the dataloader cursor had already passed.

The dataloader now counts as exhausted only when the iterator drains with
no pending prompts remaining.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Async checkpoints save the dataloader cursor ahead of what training has
consumed: the cursor advances on every yield, while the buffer holds only
completed groups and in-flight work is serialized nowhere. On resume the
collector continues from the advanced cursor, so every prompt that was
yielded but not buffered is skipped for the rest of the run. Measured on
a SWE E2E run this lost a 96-prompt window at every restore plus 64-prompt
cursor jumps — 675 of 736 skipped prompts (~30% of the reached dataset).
The losses are also length-biased: a group completes only when its longest
rollout finishes, so the in-flight (lost) set is exactly the long-rollout
prompts.

Make the trained frontier (consumed_samples) the source of truth and
treat the cursor as disposable:

- Stamp every yielded prompt with a monotonic ordinal equal to its global
  position in the dataloader stream (previously only NeMo-Gym prompts
  were stamped, and only after gap-fill slicing). Slices and carried-over
  remainders keep their original ordinals, so ordinal == stream position.
- Keep a small ring of pre-pull dataloader snapshots keyed by ordinal in
  the collector. Checkpoints persist the snapshot at the trained frontier
  into train_dataloader.pt instead of the live cursor, plus frontier
  metadata in rollouts.pt. Batches that cannot carry stamps fall back to
  live-cursor checkpoints (today's behavior), as do resumes from legacy
  checkpoints, whose ordinals are not stream-aligned.
- load_from_path additionally reports the post-age-filter task indices of
  the retained groups.
- On a frontier resume the dataloader re-yields the covered window; rows
  already trained (below the frontier) or retained in the restored buffer
  are dropped at yield, and the existing gap-fill path regenerates exactly
  the remainder — including groups age-filtered during the restore. The
  filter clears itself once the stream passes the covered window.

No prompt is skipped and none is duplicated across a save/restore cycle;
regenerated trajectories are fresh samples under post-resume weights
(in-flight vLLM state is not serializable). With
async_grpo.max_generation_failures > 0 a permanently failed group shifts
the stream, so one resume may re-train a substitute prompt once; the
default (0) is unaffected.

Existing behavior updates covered by tests: the two exact-dict assertions
on load_from_path metadata gain the retained-indices key, and the gym
stamping test now stamps at the loop boundary (where stamping moved) and
verifies ordinals survive slicing and repetition.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
…ier path

Resuming from a restored replay buffer biases the next steps toward
short-rollout prompts: a prompt group completes only when its longest
rollout finishes, so the groups sitting in the buffer at a save boundary
are systematically the short ones, while the long-rollout groups were
in flight. checkpointing.load_replay_buffer=false opts out of the
restore, and with frontier-aligned checkpoints that is simply the
empty-retained-set case of the same resume path: the dataloader is
rewound to the trained frontier and the whole buffered window is
regenerated fresh under one weight version — unbiased composition at
the cost of re-rolling the window. No prompt is skipped either way.

Restoration stays enabled by default and for legacy configs that omit
the field. On a legacy (pre-frontier) checkpoint the option only skips
the restore, matching the previous cursor semantics.

The knob is declared in the exemplar YAML and the CheckpointingConfig
Attributes block, read with a None-tolerant .get (no call-site default),
and the restore block is extracted into
_maybe_restore_async_replay_buffer_checkpoint for testability. Docs:
async-grpo.md checkpointing section updated to describe frontier-aligned
saves, lossless restores, and the new option.

This supersedes the standalone load_replay_buffer change in #3583, which
skipped the buffered window on resume because the dataloader cursor
stayed advanced; here the window is regenerated instead.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Encode the measured failure mode as a permanent contrast test: 16 unique
prompts, a checkpoint with frontier 5, groups {6, 7, 10} retained, and
{5, 8, 9, 11} in flight. The legacy resume path (still in-tree as the
fallback for pre-frontier checkpoints) restores the live cursor and loses
exactly those four prompts; the frontier resume regenerates all of them
with zero duplicates. This is the miniature of the 96-prompt restore
gaps in the prompt-skip findings report.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 12, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Aug 12, 2026
@yfw yfw added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 12, 2026

@yfw yfw left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Multi-agent review of the frontier-aligned checkpointing work. 5 agents (RL/codebase, bug hunt, tests, design, adversarial verification); every finding below survived an adversarial pass, and 10 were cut as unfounded or disproportionate.

Overall: the diagnosis is careful and well evidenced, and the _process_batch return-the-remainder contract is the highest-value, lowest-risk change here — it turned "does the gap-fill tail survive?" into a pure return-value assertion. The frontier-alignment half is where the findings cluster.

Two things worth deciding before this leaves draft:

  1. consumed_samples is a counter being used as a stream position. With max_generation_failures > 0 — 13 shipped configs, including the SWE recipes that motivated this PR — a resume both re-trains already-trained prompts and permanently skips others. Details on grpo.py:5240.
  2. This may be two PRs. The _pending_batch carry-over is small, safe, well-tested, and plausibly accounts for most of the 675 skips on its own (pre-PR, _process_batch discarded an entire dataloader batch whenever no target needed generation). Frontier alignment is the complex half and carries every bug below. Landing the first now and iterating on the second separately is a real option.

CI: pre-commit currently fails — 2 pyrefly errors, both PR-introduced. Verified fix in the first comment.

Also verified and explicitly not flagged, so you don't re-derive them: the iter() hoist is behaviour-preserving (and the pause reorder is an improvement); dataloader_exhausted is correct in both directions; the snapshot ring doesn't alias mutating state; repeat_interleave preserves stamped ordinals; _ng_task_index on native runs is inert across every consumer; load_replay_buffer: false is correct end-to-end and backward compatible; discarding the serialized pending batch on a frontier resume is correct; and the config/exemplar plumbing is complete (all 12 GRPO exemplars inherit via defaults:, and the reference config was updated in lockstep).

Generated by Claude Code

Comment thread nemo_rl/algorithms/async_utils/trajectory_collector.py Outdated
Comment thread nemo_rl/algorithms/grpo.py Outdated
Comment thread nemo_rl/algorithms/grpo.py
Comment thread nemo_rl/algorithms/grpo.py
Comment thread nemo_rl/algorithms/async_utils/trajectory_collector.py
Comment thread nemo_rl/algorithms/async_utils/trajectory_collector.py Outdated
Comment thread nemo_rl/algorithms/grpo.py
Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py Outdated
Comment thread nemo_rl/algorithms/async_utils/trajectory_collector.py
Comment thread nemo_rl/algorithms/async_utils/trajectory_collector.py
Comment thread nemo_rl/algorithms/async_utils/trajectory_collector.py Outdated
Correctness:
- Derive the checkpoint frontier from the sampled groups' own stream
  ordinals (running max + 1) instead of consumed_samples. The counter
  never sees the holes left by tolerated generation failures
  (max_generation_failures > 0), so it lags the true stream position:
  a resume would re-train the lag window and treat the dropped prompts
  as covered. Initialized from the restored frontier on resume.
- Read the dataloader snapshot and the pending remainder in one actor
  call under the pending lock (get_checkpoint_state). Separate reads
  let the collection loop consume/pull in between, so a fallback
  checkpoint could pair a stale cursor with a newer pending batch and
  duplicate its prompts on resume. Frontier-aligned checkpoints now
  omit the pending batch entirely — the rewound stream re-yields it.
- Capture each native group's ordinal from the *input* rows before the
  rollout runs. Multi-turn environments may replace extra_env_info
  wholesale, so the output rows are not a reliable stamp carrier.
- Return the whole batch (not just the gap-fill tail) from
  _process_batch's exception path when no worker was started; nothing
  was consumed, and returning the tail dropped the sliced prefix.
- Compute _skip_horizon as the max of both terms; max(default=) only
  applies the frontier term when the retained set is empty.

Diagnosability:
- pyrefly: narrow re-yielded ordinals via cast; drop dead hasattr
  guards (BatchedDataDict is a UserDict).
- Warn when the snapshot ring holds nothing at or below the frontier,
  and when a legacy resume permanently disables frontier alignment for
  the checkpoint lineage.
- Make the missing-task-index ValueError explain the likely cause
  (dataset/processor changed since the save) instead of surfacing as a
  misleading exhaustion error.

Docs:
- Scope the lossless-restore claim to the save/restore itself and list
  every fallback condition and its blast radius.
- Scope load_from_path's never-re-issued docstring to legacy resumes;
  update the nemo-gym design doc for yield-time stamping and the
  deliberate ordinal re-issue on frontier resumes.

Tests:
- create_local_collector forwards the resume kwargs; frontier tests
  construct collectors instead of poking private attributes.
- Driver-level contract tests: rollouts.pt carries the frontier keys
  iff the snapshot is aligned (write side), and the keys plumb into
  the collector's resume kwargs with the pending batch dropped (read
  side). Stub buffer groups now carry ordinals so the trained-frontier
  tracker is exercised.
- Pin the aligned-but-evicted ring fallback and the one-call
  checkpoint pair shape.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw

yfw commented Aug 14, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 95844b6

@yfw
yfw marked this pull request as ready for review August 14, 2026 01:56
@yfw
yfw requested review from a team as code owners August 14, 2026 01:56
@NVIDIA-NeMo NVIDIA-NeMo deleted a comment from yfw Aug 14, 2026
@yfw
yfw force-pushed the yifu/replay-buffer-prompt-skips branch 2 times, most recently from c19fc96 to f1e7bd9 Compare August 14, 2026 22:12
…int reads

get_checkpoint_state duplicated get_rollouts_state's key-building instead
of reusing it (threading.Lock is not reentrant, so it could not simply
call the locked accessor). Extract a lock-free _build_rollouts_state that
both call under _pending_lock; get_rollouts_state stays as the standalone
accessor used by tests and diagnostics, and its docstring now says the
driver reads it through get_checkpoint_state.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw
yfw force-pushed the yifu/replay-buffer-prompt-skips branch from f1e7bd9 to ee5741a Compare August 14, 2026 22:18
@yfw

yfw commented Aug 15, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test ee5741a

…ontier

A tolerated generation failure (max_generation_failures > 0) refills the
failed target from later prompts, so training can advance the frontier
past another target's still-generating groups. A checkpoint cut at the
trained frontier then strands those never-failed prompts: they sit below
the resume base (never re-yielded) and are not in the buffer (unfinished
groups are not stored). Thanks to @terrykong for finding this and writing
up the failure scenario. The trigger is broader than failures alone: any
target refilled from later prompts interleaves ordinals — routinely
including a resume that gap-fills an incomplete restored target — so
max_generation_failures=0 configs with max_trajectory_age_steps > 1
(e.g. grpo-qwen3-30ba3b-24n8g-async-8off.yaml) are covered too.

Track the ordinals dispatched to each rollout worker: registered at
dispatch, discarded per group when it lands in the buffer, and swept at
worker exit (a failed batch's remainder is a permanent, documented loss
and must not pin the cut). get_checkpoint_state saves at
cut = min(lowest outstanding ordinal, trained frontier) and persists the
cut as the resume filter threshold — lowering only the base while the
filter stayed at the frontier would re-strand the window one layer down.

Prompts already trained at or above the cut are persisted in
rollouts.pt (TRAINED_TASK_INDICES_KEY) and folded into the resume's
covered set alongside retained groups, so the re-yielded window
regenerates only what was genuinely lost: no skip, no re-training. The
driver tracks recently trained ordinals, seeds the set from the restored
checkpoint, and prunes below the cut at each save (the cut never
decreases, so pruning is safe). With no interleaving every outstanding
ordinal is at or above the frontier, the cut equals it, and the trained
list is empty — the healthy path is byte-identical.

Tests (each verified to fail under the mutant it targets): the 24-prompt
interleaving scenario (cut at 8, not 21; resume regenerates the stranded
window); buffered groups clear outstanding before the worker exits — the
per-group discard is what stops a checkpoint landing between
last-group-buffered and worker-exit from spuriously lowering the cut on
a healthy run; a worker failing after shutdown began still sweeps in its
finally; an end-to-end healthy dispatch leaves the cut at the frontier
with no warning; the checkpoint persists a stub-lowered cut, not the
trained frontier; outstanding-set lifecycle through dispatch/buffered/
failed-worker paths including start-failure rollback.

Docs: guide, warning text, and docstrings describe the widened trigger,
the qualified no-duplication guarantee, the re-train bound, and the
load_replay_buffer=false interaction with a lowered cut (buffered-
untrained groups below the cut are covered only by the buffer).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
yfw and others added 2 commits August 24, 2026 19:55
Under target interleaving with partial completion, a target's finished
groups can sit below the cut while its stragglers keep the cut low —
until they finish too, at which point the cut rises above ordinals whose
only record is the replay buffer. A resume with
checkpointing.load_replay_buffer=false discards that buffer and the
threshold filter drops everything below the cut, silently losing those
buffered-but-untrained groups. Previously a documented caveat; closed
now that the flag is planned for production use.

Fold the buffer's held ordinals into the cut:
cut = min(lowest outstanding, lowest buffered, trained frontier).
The buffer reports them via get_held_task_indices, queried inside the
atomic get_checkpoint_state read, after the outstanding set — a group
leaves that set only once its buffer add succeeded, so every dispatched
group is visible to at least one of the two reads. Sampling removes
trained groups, so everything held is untrained and sits at/above the
frontier on healthy runs: the healthy path is unchanged, and
load_replay_buffer=true resumes cover the re-yielded groups via the
retained set exactly as before (no duplication).

Tests: the partially-buffered scenario (outstanding {13..15}, buffered
{8..12}, frontier 21 -> cut 8; mutation-verified against dropping the
buffered term), the healthy buffered case (held >= frontier -> cut
unchanged), and an end-to-end lowered-cut resume with no retained set
that regenerates the formerly-buffered window. Guide caveat replaced
with the closed semantics.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
…ompt-skips

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Comments and docstrings now describe only what the code cannot show —
locking contracts, read-ordering constraints, and the invariants the
checkpoint cut and covered set maintain — instead of narrating how the
design was arrived at or comparing against prior behavior. Also corrects
the stale FRONTIER_ORDINAL_KEY description (it holds the cut, not
consumed_samples). No code changes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw
yfw force-pushed the yifu/replay-buffer-prompt-skips branch from ee5741a to 9353dfb Compare August 25, 2026 04:24
@yfw

yfw commented Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 9353dfb

@yfw
yfw merged commit 05c6e8a into main Aug 25, 2026
147 of 149 checks passed
@yfw
yfw deleted the yifu/replay-buffer-prompt-skips branch August 25, 2026 09:56
yfw added a commit that referenced this pull request Aug 25, 2026
#3599)

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
(cherry picked from commit 05c6e8a)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants