Repository navigation
test(checkpoint): add Workplace prefix-recovery parity - #4507
macandro96 wants to merge 9 commits into
Conversation
5bccccd to
082e927
Compare
082e927 to
1a6de5c
Compare
1a6de5c to
cc5d8bc
Compare
cc5d8bc to
deb0e14
Compare
1 similar comment
deb0e14 to
24050b2
Compare
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Thanks for this. It adds the token-level parity against an uninterrupted run that was asked for on #4500 (discussion) and #4266 (discussion), which is the right check for prefix recovery.
Two comments are about what the test proves:
- restored rows are excused too widely, so a changed tool result or turn-1 token passes (parity L516);
top_p=0.000001keeps rewards out of the gradient (script L194).
Size. About 1,040 of the 2,098 added lines can go without weakening a check:
| What | Source lines | Unit-test lines |
|---|---|---|
Use the built-in log_full_train_data dump instead of the recorder (hook L57) |
~125 | ~28 |
| Ordering / timeline / Kendall-tau code (parity L691) | ~215 | ~127 |
| Hand-rolled Spearman / Pearson / percentiles (parity L548) | ~125 | ~70 |
| Restored-row excusal (parity L516) | ~70 | ~115 |
| Small repeated checks, fallbacks and copied helpers (L346, L367, L285, L304), plus the report printing and script plumbing that only feed the code above | ~165 | — |
A version with the same checks is about 330 source lines plus 260 unit-test lines.
This PR also builds on #4506, whose 25 review comments (including its ~300-line size ask) are still open. Please apply #4506's size asks first, so this PR shrinks on top of a smaller base; the prefix-recovery stack is already over 20k lines.
Gym pin. This test also needs Gym 7a254e8e98 and 994cd12404 (prefix mode and the SimpleAgent _prepare_model_request_for_turn hook) on Gym main. Same ask as the open #4500 thread (discussion).
Small nit: the PR body says "unit 37 passed", but the unit file has 38 cases.
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4507-size-and-reuse.html
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4507-excused-rows.html
Generated by Claude Code
| data.validation.data_path="$TEST_DATA" | ||
| ) | ||
|
|
||
| echo "=== Reference: run the same workload without interruption ===" |
There was a problem hiding this comment.
1 action item.
TL;DR — no GPU CI has run on this PR or #4506, and this PR turns the test into 3 launches of a 1.7B model for 3 steps (was 2 launches of 0.6B for 1 step). Please run CI and post the job time.
PR-introduced cost:
- 2 → 3 launches (this reference run is new), 1 → 3 steps, Qwen3-0.6B → 1.7B, 2 → 4 prompts per step.
- Each launch spends ~3–4 min on startup. For scale,
grpo_async_gym_single_controller.sh(10 steps, 0.6B) took 4m16s in the latest green main H100L1_Functional_Tests_SingleController_3job, out of ~41 min total. Estimated +5–10 min over feat(checkpoint): recover active vLLM generation prefixes #4506; not measured. - The only checks on feat(checkpoint): recover active vLLM generation prefixes #4506 and test(checkpoint): add Workplace prefix-recovery parity #4507 so far are copyright, secrets and submodule, so "passes locally" is the only GPU evidence.
- The L1 runner comment still says "two long-generation runs" (L1_Functional_Tests_SingleController_3.sh#L66-L68).
AI-1
Run /ok to test <head SHA>, post the L1_Functional_Tests_SingleController_3 job time in the PR body, and update the L1 comment (L1_Functional_Tests_SingleController_3.sh#L67) to say three runs. Not a one-click block: that file is not in this diff, and the rest is not a code edit.
There was a problem hiding this comment.
The L1 runner comment now says three runs (8feaa87499). Current size: 1 train step, 3 launches, Qwen3-1.7B, 4 prompts x 2 generations. Locally (1 node, 2 GPUs) each launch spends ~70 s in setup, and the whole test took 7-9 min in the last two runs (412 s and 531 s). I'll post /ok to test and the L1_Functional_Tests_SingleController_3 job time once it runs.
Note: the GPU CI run on #4506 checked out a stale merge commit (GitHub's merge ref still points at the Oct 14 head), so its unit failures are not from the current code. I'll re-trigger once that ref is refreshed.
|
Thanks for the review. The inline items are addressed in What the test proves: exact parity for every row, with one verified exception: a restored row whose phase-1 prefix already diverged before the cut. A real run needed it; details are in the excusal thread. The script now says the test checks data parity, not learning. Size: #4507 is now +1,503 / -125 lines against #4506 (was +2,090 / -73); the parity helper is 578 lines. The extra deletions are the turn-recovery script's copies, now shared. That is still well above your ~590-line estimate. Most of the remainder is the pre-cut check, the trainer-dump loader, and the #4506: its review round landed in Gym pin: acknowledged. The stack pins Nit: the description will list the current counts (22 parity unit cases). |
|
/ok to test 8feaa87 |
1 similar comment
|
/ok to test 0c77813 |
…t be read After rebasing onto #4507's fail-soft prefix restore, a call's media is restored from the same chunk fetch as its prefix, inside the same fallback. A prefix row whose media metadata cannot be read therefore regenerates the call from its turn boundary, like any other unusable prefix. Reprocessed media that does not match the staged digests is still rejected with retained_media_changed. Cover the first case and say both in the guide. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
0c77813 to
558eb4d
Compare
f66dec2 to
6986fae
Compare
6986fae to
464275e
Compare
464275e to
f6d6aa5
Compare
terrykong
added
CI:Lfast
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Make the parity test exercise many restored prefixes and separate recovery faults from run-to-run sampling noise: - Workload: four Workplace prompts per step whose closing turn lists 1..N (N = 150/250/350/450) at natural length. Group lengths are far apart, so prompt-group readiness order is deterministic. One prompt expects a different duration so rewards are not constant. Runs Qwen3-1.7B for 3 steps. - Selection/verification: require a recoverable prefix in every Workplace prompt group at the crash (Gym refuses dispatches that land during a checkpoint, so siblings may start late), and verify each cut returns to its Gym owner under attempt+1, restores its exact prefix digest, and completes once with its trained generated-token count equal to the restored prefix plus the regenerated tail. - Parity: classify token divergences relative to the restored prefix. A divergence inside the prefix is excused as phase-one noise only when the trained prefix tokens hash-match what the worker restored, and excused rows still require equal rewards and sample masks; a prompt divergence always fails. Gate vLLM generation_logprobs in aggregate over token- matched rows (bit-identical fraction, Spearman, max abs diff) since they are not batch-invariant, and require exact prompt-group readiness order. Report Kendall tau for arrival and sealing order, per-rollout timelines, and logprob statistics in parity-report.json before the strict checks run. - Hook: timestamp every event and record completion arrival separately from the ledger seal. - Relies on the Gym checkpoint test agent's natural-length closing turn (NEMO_GYM_TEST_PREFIX_MIN_TOKENS=0), pinned by the base branch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…very parity The prefix-recovery parity harness reads each restored cut from Gym's checkpoint records and reports it. NeMo Gym now names the staged prefix prefix_token_count and prefix_digest, so the harness, its report, and its unit tests use those keys. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
… test The prefix-recovery snapshot helper now imports _gym_recovery_snapshot_helpers as a sibling module. Put tests/functional on sys.path before loading it by file path, as running it as a script does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Parity data and gates: - Compare runs on the trainer's own async_rl.log_full_train_data dump (read from the run's exp_NNN directory) instead of a recorder in the hook; the hook now only probes live staging keys per train step. - Require exact parity for every row. The one exception is a restored row whose phase-1 prefix already differed from the reference before the cut (vLLM sampling is not batch-invariant): it must start with exactly the restored tokens (prefix_ids_sha256), match up to its closing call, and keep its reward and admission; at least one restored row must match outright. - Gate generation logprobs on bit-identical fraction and max |diff|. - Default to one train step and say the test checks data, not learning (top_p=1e-6 zeroes the policy gradient). Removed: the ordering/timeline/Kendall report, the hand-rolled Spearman/Pearson gate and parity-report.json, the metric comparison, the prefix+tail==total tautology, the duplicate trained-once check, the unreachable successor Gym-manifest check, the single-cut selection fallback, and the swallowed Workplace layout errors. The JSON readers and the restored-prefix log parser are shared in _gym_recovery_snapshot_helpers. The L1 runner comment now says the test performs three runs. The Workplace calendar sentinel count lives with the shared readers, and the Workplace test rows are built by one jq module (_workplace_calendar_rows.jq). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…agent Gym main has no checkpoint test agent, so the prefix-recovery parity test now runs the stock workplace_assistant_simple_agent behind the test hook's policy proxy, as the turn-recovery test does. With no hold and no follow-up tool, the proxy sends every call after the tool result with tool_choice none and leaves its length to the model, which is what the test agent did for the closing turn. Gym only cuts a call whose decoding is unconstrained, judged from the request its agent sends, so the rows leave tool_choice auto instead of forcing calendar_create_event: a forced tool would keep every turn out of the cut. The prompt asks for the call, and the snapshot selector already requires the calendar write to have happened exactly once. The rows name workplace_assistant_simple_agent, the script drops the test agent's config_paths override and environment-server removal, and the hook writes the proxy's events to the prefix test's events file. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…y active call The snapshot records only the cuts that succeeded, so a snapshot where Gym failed or skipped an active call's cut still qualified, and that call quietly regenerated after the restore. Read the trainer's save line for the snapshot and require zero failed and skipped cuts; a bug that keeps failing cuts now times out the selection with the counts instead of passing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…identical On Gym main's Workplace agent the workload generates ~9.3k tokens, and four runs scored 98.89-99.13% bit-identical with every exact token, mask and reward gate passing; three fell below the old 99% floor across three different progress-read implementations. The bf16 noise floor is a property of the smaller workload, not of the restore, so the gate now sits at 98%. The max-|diff| bound stays at 0.5. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…s end The prefix-recovery parity checks located a restored call at the start of its row's last unmasked segment. When top-k/top-p filtering is on, the trainer drops positions whose recomputed logprob is -inf from token_mask (#4154), and the rows these checks read are dumped after that pass, so a masked token inside the call made the check start the call too late, miscount its tokens, and fail a correct restore. A restored call closes its row, so the total_generation_tokens it logs on completion end the row. _assert_prefix_boundary now counts back from the row's end, requires the token before the call to be untrained, checks the completion's prefix + tail = total accounting, and requires the restored prefix at that position; a dropped or duplicated token moves the position. It returns where each call starts, which _rows_diverged_before_cut reuses to compare the reference row at the same position. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
f6d6aa5 to
855b226
Compare
What does this PR do ?
Adds an end-to-end parity test for generation-prefix recovery: an uninterrupted run and a hard-killed, prefix-restored run of the same Workplace workload must train on identical data.
Stacked on #4506. The base branch is
amahishi/ckpt-v2-prefix-recovery-pr1, so this diff shows only this PR's 9 commits. It is test code only.Issues
How it works
flowchart LR B[Reference run<br/>uninterrupted] --> P[Parity compare] P1[Phase 1<br/>checkpoint with active prefixes] -->|SIGKILL| P2[Phase 2<br/>restore prefixes, generate tails] P2 --> Pworkplace_assistant_simple_agent. Each episode makes a namedcalendar_create_eventcall, then closes by listing 1..N (N = 150/250/350/450) at natural length, so prompt groups finish far apart. Gym seestool_choice: autoon every turn, so both calls are cuttable; a policy-call proxy in the test hook (the one refactor(checkpoint): move rollout checkpointing out of the single controller #5011 uses for the turn test) sends every post-tool call to vLLM withtool_choice: none. One prompt expects a different duration, so rewards are not constant.top_p=1e-6makes the regenerated tail comparable token-for-token; training applies the sametop_p, so the policy-gradient term is 0. The test checks rollout and training-data parity, not learning.-inflogprobs, fix(grpo): drop top-k/top-p -inf from token_mask #4154) cannot shift where it starts. No restored prefix key is still live in any batch the trainer consumed.async_rl.log_full_train_datadump from all three runs:prefix_ids_sha256), match token-for-token up to its closing call, and keep its reward and sample mask. At least one restored row must match the reference outright;Changes
test(checkpoint): add Workplace prefix-recovery paritytest(checkpoint): strengthen Workplace prefix-recovery paritytest(checkpoint): read the renamed generation-cut keys in prefix-recovery paritytest(checkpoint): load the shared snapshot helpers in the parity unit testtests/functionalonsys.pathbefore the unit test loads the snapshot helper.test(checkpoint): address review of generation-prefix recovery paritytrain_data_step*.jsonldump; removed the ordering/correlation report, the metric comparison and redundant checks; JSON readers, the restore-log parser, the calendar sentinel count and the Workplace row builder (_workplace_calendar_rows.jq) are shared with the turn-recovery test.test(checkpoint): run prefix-recovery parity on Gym main's Workplace agenttool_choice: autoso Gym does not skip the cuts.test(checkpoint): select only prefix-recovery snapshots that cut every active calltest(checkpoint): lower the prefix-recovery logprob floor to 98% bit-identicaltest(checkpoint): find a restored call by counting back from its row's endLines vs #4506: unit tests +621, functional tests +985 / −87.
Testing
tests/unit/environments/test_gym_prefix_recovery_parity.py25 passed.tests/functional/grpo_async_gym_single_controller_prefix_recovery.shpasses locally on this head withflat_logprobs: true(1 node, 452 s): 8 restored prefixes (153–250 tokens), every row matches, 0 rows needed the pre-cut exception; generation logprobs 99.11% bit-identical over 9,312 tokens, max |diff| 0.212. It also passed with the defaultflat_logprobs: false(99.10%, max |diff| 0.162). The selected snapshot cut 8 of 8 calls with 0 failed and 0 skipped; the bootstrap snapshot, which reported 1 failed cut, was passed over.<tool_call>instead of1as a closing turn's first token. The restored prefix hash matched phase 1's tokens. That run's data passes with the exception (1 row set aside) and fails without it.🤖 Generated with Claude Code