The parity test is 2,098 lines; about half of them check something

Size review of NVIDIA-NeMo/RL #4507, stacked on #4506. Line counts are from git diff 834241af deb0e146. All permalinks pin the PR head deb0e146. Review comments are in a pending review; the links work once it is published.

Three pages on this PR — this one (round 1, an older head): which lines do the checking, and which repeat code that already exists.

The test runs the same small Workplace workload twice and requires every trained row to match token for token. The checks that can fail the run are a short list: rows match, each restored prefix trained exactly once with the right length, no old prefix keys left behind, rewards are not all the same, and generation log-probs agree. Most of the remaining code prints reports, re-checks something already checked, or rebuilds a training-data dump that NeMo-RL already writes.

Where the lines go

The biggest reuse, as a diff — proposed, not in the PR

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

Where the fix lands: RL test files only. prefix_recovery.sh, the hook, and the parity script. No product change: the dump already exists. It is written at the advantage stage with token ids, loss mask, advantages, generation log-probs, rewards and sample mask, and a retried step replaces its own file (train_data_dump.py:95). The author's earlier version of this harness (#4320) read these same files.

@@ grpo_async_gym_single_controller_prefix_recovery.sh L162 · COMMON_OVERRIDES
 COMMON_OVERRIDES=(
     policy.model_name="$MODEL_NAME"
+    # Built-in per-step dump of every trained row: what the parity check compares.
+    ++async_rl.log_full_train_data=true

@@ _single_controller_turn_recovery_hook.py L57-L139 · _TrainingPayloadRecorder
-class _TrainingPayloadRecorder:
-    """Test-only TQPolicy wrapper that records the rows sent to training."""
   … 81 lines: read rows back from the data plane, serialize, append …
+class _StagingKeyProbe:
+    """Record the live staging keys per train step; the dump has everything else."""
   … ~20 lines: wrap begin_train_step, append list_sample_ids(staging) …

@@ _gym_prefix_recovery_parity.py L65-L141 · serialize / append / _tensor / _generation_identity
-def serialize_training_batch(
   … 76 more deleted lines …
+def _load_dump(log_dir: Path) -> list[dict[str, Any]]:
+    # One file per optimizer step, written only after the step finished.
+    return [row for path in sorted(log_dir.rglob("train_data_step*.jsonl"))
+            for row in _read_jsonl(path)]

Each run (reference, phase 1, phase 2) gets its own log folder, so phase 1's files never mix with phase 2's. Today phase 1 and phase 2 share one (:234, :285).

review comment — turn_recovery_hook.py:57