fix(async-grpo): eliminate prompt skips across checkpoint save/restore - #3599
Conversation
…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>
yfw
left a comment
There was a problem hiding this comment.
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:
consumed_samplesis a counter being used as a stream position. Withmax_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 ongrpo.py:5240.- This may be two PRs. The
_pending_batchcarry-over is small, safe, well-tested, and plausibly accounts for most of the 675 skips on its own (pre-PR,_process_batchdiscarded 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
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>
|
/ok to test 95844b6 |
c19fc96 to
f1e7bd9
Compare
…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>
f1e7bd9 to
ee5741a
Compare
|
/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>
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>
ee5741a to
9353dfb
Compare
|
/ok to test 9353dfb |
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=trueandfalse.Problem
On a SWE E2E run, 675 of 736 skipped prompts (~30% of the reached dataset) were traced to two mechanisms:
_process_batchsliced off what it needed and dropped the rest (61 of the 736 skips).Fix (one commit each)
b1b196a1e) —_process_batchreturns the unconsumed remainder; the loop consumes it before the next pull, and it is serialized inrollouts.ptso a checkpoint can't strand it.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.checkpointing.load_replay_buffer(a72ebf48) — setfalseto 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 remainstrue. Supersedes feat(grpo): optionally skip replay buffer restore #3583.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.a482c6c5,8c1e170d) — frontier derived from the sampled groups' own ordinals instead ofconsumed_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, androllouts.ptcontract tests.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 atmin(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
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.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 werescanceled mid-generation and relaunched under the same lineage; thetrain_data_step*.jsonldumps were then verified against the dataset's expected stream order — every step must train exactly its 32-prompt window, each prompt exactlynum_generationstimes.Frontier-aligned resume: dataloader rewound to ordinal 64, trained frontier 64, 13 retained prompt groups→dropping 13 already-covered prompts (ordinals 64–95), regenerating 19— those 19 in-flight prompts are exactly what the pre-PR code silently skipped.max_generation_failures=3as in the SWE recipes): 6 steps, 2 kills → 192/192, zeroCheckpoint cut loweredwarnings 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
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information