A restored row can differ outside its restored prefix and still pass the parity test

Finding on NVIDIA-NeMo/RL #4507. All permalinks pin the PR head deb0e146. The review comment is in a pending review; its link works once it is published.

Before
No parity test existed. The base PR only checked that the restored run trained.
This PR
Requires each trained row to match an uninterrupted run, but excuses some rows of restored runs as "sampling noise" — and the excused range is wider than the restored tokens.
Status
Open — not fixed in this PR. Shown on a small hand-built case (CPU, the PR's own function). Raised as a review comment.

Three pages on this PR — this one (round 1, an older head): which differences the test lets through.

After a crash, the restored run resumes a half-finished model reply: the tokens saved before the crash (the restored prefix) are replayed, and the model generates the rest. The test then compares each trained row with the same row from a run that never crashed. When a row differs, _classify_token_divergence decides whether to excuse it. It excuses the row (pre_cut) when the first difference sits before the end of the restored prefix — but it starts counting at the row's first generated token, which is in turn 1, not at the start of the restored prefix.

One row, ten tokens

prompt turn 1: tool call tool result closing turn: list of numbers baseline 10 11 20 21 30 40 41 42 43 44 restored run 10 11 20 21 77 40 41 42 43 44 index 01 23 45 67 89 loss mask 00 11 01 11 11 excused by the code: index 2 up to 7 (first generated token … end of prefix) really restored: 2 tokens (5, 6) first difference: index 4, the tool result (30 → 77). It is inside the orange range, so the row is pre_cut. pre_cut rows compare only reward and sample mask — tokens, loss mask, advantages and log-probs are skipped. The test passes.

This is the PR's own function run on that row: {'kind': 'pre_cut', 'first_diff': 4, 'restored_prefix_end': 7, 'restored_prefix_hash_matches': True}. A changed turn-1 tool call (index 2 or 3) is excused the same way. A tool result is never sampled, so a difference there is a restore bug — for example Gym handing back the wrong state. The test lets it through as long as the reward still matches.

review comment — _gym_prefix_recovery_parity.py:516

The fix: drop the excuse — proposed, not in the PR

Nothing below is implemented, and it has not been run.

Where the fix lands: RL test files only. _compare_training_rows and its caller, plus the unit tests for the classifier. No product change.

Narrowing the range to the restored prefix would still excuse a noise source that no other row gets. Rows that finished before the crash, and the tail regenerated after it, face the same small bf16 rounding differences, and both must already match exactly. The author's run matched all 24 rows. So the simpler test is also the stricter one: every row must match.

@@ _gym_prefix_recovery_parity.py L212 · _compare_training_rows signature
     atol: float,
-    excused: frozenset[tuple[int, int, int]] = frozenset(),
 ) -> None:

@@ _gym_prefix_recovery_parity.py L256-L265 · _compare_training_rows
         left, right = baseline[identity], recovery[identity]
-        if identity in excused:
-            # Tokens differ by phase-one sampling noise, but the verifier and
-            # the row's admission must still agree.
-            for field in _EXCUSED_ROW_EXACT_FIELDS:
-                if left.get(field) != right.get(field):
-                    raise AssertionError(f"row {identity!r} field {field!r} differs")
-            continue
         for field in _EXACT_FIELDS:
             if left.get(field) != right.get(field):
-                raise AssertionError(f"row {identity!r} field {field!r} differs")
+                # Say where, so a real token flip is quick to find.
+                raise AssertionError(
+                    f"row {identity!r} field {field!r} differs at index "
+                    f"{_first_difference(left.get(field), right.get(field))}"
+                )

@@ _gym_prefix_recovery_parity.py L975 · compare()
         atol=args.atol,
-        excused=frozenset(
-            identity
-            for identity, classification in divergence.items()
-            if classification["kind"] == "pre_cut"
-        ),
     )
   … also delete the divergence report in compare() and _classify_token_divergence, _restored_prefix_hashes, _token_ids_sha256 (L440-L523),
       _EXCUSED_ROW_EXACT_FIELDS (L37) and their unit tests: ~70 script + ~115 test lines …

What still catches a bad restore after this change: a dropped or swapped prefix token makes the row differ from the baseline (exact match), and the check that the trained closing turn has exactly prefix + tail tokens stays (:356-364). _first_difference is a new 3-line helper.

So what. The excuse exists for noise that the passing run never showed, and it opens a gap exactly where a restore bug would land: in what Gym hands back between turns. Deleting it makes the test both smaller and stricter. If a real run then fails on a single flipped token, the error message says where, and that is the moment to decide whether an excuse is needed.