Where #4507's 1,606 test lines go, and the ~300 that can go without weakening the restore check

NVIDIA-NeMo/RL #4507, part 6 of 11 in the partial-checkpointing stack. All permalinks pin the PR head f6d6aa5.

Before
Round 1 found about half of the then 2,098 lines repeated existing code. The author fixed 14 of 15 comments and the PR shrank to 1,606.
This PR
Test only. It runs the same workload three times on Qwen3-1.7B (uninterrupted, killed mid-generation, restored) and checks that the restored run trains on the same rows.
Status
Proposed, not in the PR. The cuts below are review asks (C17–C22). None lets a broken restore pass. The functional test itself has never run in CI.

Three pages on this PR. This one: how much of the current head the test needs.

The test reads the rows each run trained on from the trainer's own dump (load_training_records), then checks each restored call: it finished once, it starts its row's last call, and its first tokens hash to what the restore logged (_assert_prefix_boundary). Every other row must match the uninterrupted run exactly. Those checks are worth keeping. The extra lines are checks that cannot fail, the same lookups built several times, and unit tests that repeat each other.

Lines added per file, and what each review comment removes

kept removed by a comment (about) added (C16 fix, moved jq) scale: 70 px = 100 lines unit test (new) 621−110+20 _gym_prefix_recovery_parity.py 601−150+10 _gym_prefix_recovery_snapshot.py 136−13 C21 prefix_recovery.sh 121−2+11 jq moved in (C20) turn_recovery_hook.py 57 _workplace_calendar_rows.jq 34−34 C20 (one caller) snapshot_helpers.py + L1 list 36 unit test −110 (C22): 4 repeat another test, 2 restate a few lines, 2 test code C17/C18 remove; plain imports save 14. parity script −150: C19 one row index + one cut check (−100), C18 advantages compare (−45), C17 a check that cannot fail (−5). +30 green: the C16 near-tie check (≈10) and its two unit tests (≈20) — a correctness fix, not a size cut. Total: about −300 lines, +40 back for C16 and the moved jq → about 1,300 instead of 1,606. No product function replaces the rest: nemo_rl has no reader for the trainer dump, and the restore lines exist only in logs.

Two checks that can never fail (C17, C18)

Where the fix lands: RL only, test files only. _gym_prefix_recovery_parity.py and the matching unit tests. No product change.

prefix + tail = total. The vLLM worker computes the logged total as exactly that sum (vllm_worker_async.py:1914-1931). The author removed this check in round 1, and the last commit brought it back. Advantages within tolerance. GRPO computes each row's advantage from its prompt group's rewards and masks (advantage_estimator.py:106-145), and the test already compares rewards, masks and group membership exactly. With top_p=1e-6 every advantage here is 0 anyway.

@@ _gym_prefix_recovery_parity.py L330-L336 · _assert_prefix_boundary
         if tail_tokens <= 0:
             raise AssertionError("restored generation produced no suffix tokens")
-        if prefix_tokens + tail_tokens != total_tokens:
-            raise AssertionError(
-                "restored generation token accounting has a gap or overlap: "
-                f"prefix={prefix_tokens}, tail={tail_tokens}, total={total_tokens}"
-            )
@@ _gym_prefix_recovery_parity.py L36-L38, L166-L195, L267-L274 · approximate compare
-_APPROXIMATE_FIELDS = ("advantages",)
-def _assert_nested_close(baseline, recovery, *, path, rtol, atol): ...   # 30 lines
-        for field in _APPROXIMATE_FIELDS:
-            _assert_nested_close(left.get(field), right.get(field), ...)
         # plus the rtol/atol parameters and flags (L205-206, L561-562, L595-596)

What still catches a dropped or extra token: counting back from the row's end lands the restored call at the wrong place, and its hash no longer matches (parity.py:352-364). That stays.

C17 — parity.py:332 C18 — parity.py:36

One row index and one cut check (C19)

The rows are re-indexed 7 times (_rows_by_identity), and the two functions that check each cut build their own row maps and look up the restored hash twice. Merging them into one ~45-line function keeps every check:

CheckTodayAfter
each prefix restored once and finished oncetwo list scansone _once() lookup each
restored call sits at the row's end, after an untrained tokenL352-359same lines
its first tokens hash to the restore logL360-364same
a set-aside row matches before its last callL444-450same, plus the C16 near-tie check
not every cut is set asideL452-456same

C19 — parity.py:138

So what. The test is worth having: it is the only end-to-end proof that #4506's prefix restore trains the same tokens as an uninterrupted run. But it can be about 300 lines smaller with the same checks, and one check needs to get stronger: a set-aside row should need evidence of a real near-tie (C16), or a wrong captured prefix passes for every cut but one. Before merge it also needs its first CI run.

C16 — parity.py:451 C17 — parity.py:332 C18 — parity.py:36 C19 — parity.py:138 C20 — calendar_rows.jq:12 C21 — snapshot.py:83 C22 — unit test:25