From fbc9c1a153b6dd3133d343de8dfda49ce5066218 Mon Sep 17 00:00:00 2001 From: Anil Balireddy Date: Wed, 30 Sep 2026 08:45:32 -0700 Subject: [PATCH] Fix off-by-one in GRPO rollout chunk-slicing that produces empty chunks _build_chunk_slice_segment() computed chunk boundaries with ceiling division: CHUNK_SIZE = ceil(EFFECTIVE / NUM_CHUNKS) START_LINE = CHUNK_ID * CHUNK_SIZE + 1 END_LINE = (CHUNK_ID + 1) * CHUNK_SIZE, clamped to EFFECTIVE Whenever EFFECTIVE doesn't divide evenly by NUM_CHUNKS, NUM_CHUNKS * CHUNK_SIZE overshoots EFFECTIVE, so the last chunk (or several, when NUM_CHUNKS > EFFECTIVE) gets START_LINE > END_LINE -- a 0-byte chunk input, and therefore a 0-byte chunk output. E.g. EFFECTIVE=4, NUM_CHUNKS=3: CHUNK_SIZE=2, chunk 2 gets START_LINE=5, END_LINE=4. That empty-but-legitimate chunk then trips the merge step's own safety check (added specifically to catch real upstream failures -- see the "silent-success cascade" comment a few lines below in the same file), which reports "chunk N .done exists but file missing/empty -- aborting" and hard-fails the entire seed's merge, even though nothing actually went wrong upstream. This is reachable any time a rollout config chunks a sample count that isn't a clean multiple of num_chunks -- a common shape for smoke tests and small eval sets. Replace with balanced-remainder division: the first (EFFECTIVE % NUM_CHUNKS) chunks get one extra line, the rest get EFFECTIVE // NUM_CHUNKS. This covers exactly 1..EFFECTIVE with no gaps, overlaps, or empty chunks for any CHUNK_ID in [0, NUM_CHUNKS) whenever NUM_CHUNKS <= EFFECTIVE (proof: REM*(BASE+1) + (NUM_CHUNKS-REM)*BASE = REM + NUM_CHUNKS*BASE = EFFECTIVE exactly, by definition of floor division/modulo). A chunk is only empty when CHUNK_ID >= EFFECTIVE (more chunks requested than samples exist), which the merge check should still legitimately catch. Testing: this is generated bash, not directly importable Python logic, so tests/test_rollout.py's new tests actually execute the rendered segment via subprocess against real input files rather than only snapshotting the text -- a snapshot alone would not catch a wrong split point. Covers: full coverage with no gaps/overlaps across total_lines 2-12 x num_chunks 2..total_lines, no chunk ever empty in that range (the direct regression test), the exact previously-buggy 4-lines/3-chunks repro, the NUM_CHUNKS<=1 no-op path, and the max_num_samples clamp. Verified the new tests actually fail against the old ceiling-division formula before confirming they pass against the fix. _build_client_cmd() composes this segment into its own rendered output, so its five byte-exact fixture snapshots under tests/fixtures/rollout/ needed regenerating via scripts/_dump_rollout_fixtures.py; the diff in each is identical and scoped to exactly this arithmetic block. Co-Authored-By: Claude Sonnet 5 Signed-off-by: Anil Balireddy --- nvflow/lib/rl/rollout.py | 30 ++++- .../rollout/client_cmd_dual_server.txt | 14 ++- .../rollout/client_cmd_max_samples.txt | 14 ++- .../fixtures/rollout/client_cmd_no_chunk.txt | 14 ++- tests/fixtures/rollout/client_cmd_no_rcp.txt | 14 ++- .../rollout/client_cmd_policy_only.txt | 14 ++- tests/test_rollout.py | 113 ++++++++++++++++++ 7 files changed, 189 insertions(+), 24 deletions(-) diff --git a/nvflow/lib/rl/rollout.py b/nvflow/lib/rl/rollout.py index 37b50aa..ae84503 100644 --- a/nvflow/lib/rl/rollout.py +++ b/nvflow/lib/rl/rollout.py @@ -580,6 +580,22 @@ def _build_chunk_slice_segment(max_num_samples: int) -> str: No physical pre-splitting on the launcher — keeps the launcher lightweight and filesystem-agnostic. The ``if [ $NUM_CHUNKS -gt 1 ]`` runtime guard suppresses slicing in the single-chunk case. + + Chunk boundaries use balanced-remainder division (the first + ``EFFECTIVE % NUM_CHUNKS`` chunks get one extra line) rather than + ceiling division. Ceiling division overshoots ``EFFECTIVE`` by up + to a full chunk once ``NUM_CHUNKS * ceil(EFFECTIVE / NUM_CHUNKS)`` + exceeds it, which silently produces an empty trailing chunk (or + several) any time ``EFFECTIVE`` doesn't divide evenly by + ``NUM_CHUNKS`` -- a legitimately non-empty chunk range then reads + as a merge failure downstream, since the merge step correctly + treats a missing/empty chunk file as an upstream error. Balanced + division covers exactly ``1..EFFECTIVE`` with no gaps or overlaps + for every ``CHUNK_ID`` in ``[0, NUM_CHUNKS)`` whenever + ``NUM_CHUNKS <= EFFECTIVE``; a chunk is only empty when + ``CHUNK_ID >= EFFECTIVE`` (more chunks requested than samples + exist), which is a real misconfiguration the merge check should + still catch. """ return ( 'CHUNK_INPUT=""\n' @@ -592,10 +608,16 @@ def _build_chunk_slice_segment(max_num_samples: int) -> str: " if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then\n" " EFFECTIVE=$MAX_SAMPLES\n" " fi\n" - " CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS ))\n" - " START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))\n" - " END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE ))\n" - " [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE\n" + " BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS ))\n" + " REMAINDER=$(( EFFECTIVE % NUM_CHUNKS ))\n" + " if [ $CHUNK_ID -lt $REMAINDER ]; then\n" + " CHUNK_SIZE=$(( BASE_SIZE + 1 ))\n" + " START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 ))\n" + " else\n" + " CHUNK_SIZE=$BASE_SIZE\n" + " START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 ))\n" + " fi\n" + " END_LINE=$(( START_LINE + CHUNK_SIZE - 1 ))\n" ' CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl"\n' ' head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT"\n' ' echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)"\n' diff --git a/tests/fixtures/rollout/client_cmd_dual_server.txt b/tests/fixtures/rollout/client_cmd_dual_server.txt index 9c54ac9..4ff49c3 100644 --- a/tests/fixtures/rollout/client_cmd_dual_server.txt +++ b/tests/fixtures/rollout/client_cmd_dual_server.txt @@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then EFFECTIVE=$MAX_SAMPLES fi - CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) - START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) - END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) - [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS )) + REMAINDER=$(( EFFECTIVE % NUM_CHUNKS )) + if [ $CHUNK_ID -lt $REMAINDER ]; then + CHUNK_SIZE=$(( BASE_SIZE + 1 )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + else + CHUNK_SIZE=$BASE_SIZE + START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 )) + fi + END_LINE=$(( START_LINE + CHUNK_SIZE - 1 )) CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" diff --git a/tests/fixtures/rollout/client_cmd_max_samples.txt b/tests/fixtures/rollout/client_cmd_max_samples.txt index 83c5902..6a11de8 100644 --- a/tests/fixtures/rollout/client_cmd_max_samples.txt +++ b/tests/fixtures/rollout/client_cmd_max_samples.txt @@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then EFFECTIVE=$MAX_SAMPLES fi - CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) - START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) - END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) - [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS )) + REMAINDER=$(( EFFECTIVE % NUM_CHUNKS )) + if [ $CHUNK_ID -lt $REMAINDER ]; then + CHUNK_SIZE=$(( BASE_SIZE + 1 )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + else + CHUNK_SIZE=$BASE_SIZE + START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 )) + fi + END_LINE=$(( START_LINE + CHUNK_SIZE - 1 )) CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" diff --git a/tests/fixtures/rollout/client_cmd_no_chunk.txt b/tests/fixtures/rollout/client_cmd_no_chunk.txt index 4a163bc..1478f56 100644 --- a/tests/fixtures/rollout/client_cmd_no_chunk.txt +++ b/tests/fixtures/rollout/client_cmd_no_chunk.txt @@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then EFFECTIVE=$MAX_SAMPLES fi - CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) - START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) - END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) - [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS )) + REMAINDER=$(( EFFECTIVE % NUM_CHUNKS )) + if [ $CHUNK_ID -lt $REMAINDER ]; then + CHUNK_SIZE=$(( BASE_SIZE + 1 )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + else + CHUNK_SIZE=$BASE_SIZE + START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 )) + fi + END_LINE=$(( START_LINE + CHUNK_SIZE - 1 )) CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" diff --git a/tests/fixtures/rollout/client_cmd_no_rcp.txt b/tests/fixtures/rollout/client_cmd_no_rcp.txt index e6e7937..ed7ec6c 100644 --- a/tests/fixtures/rollout/client_cmd_no_rcp.txt +++ b/tests/fixtures/rollout/client_cmd_no_rcp.txt @@ -103,10 +103,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then EFFECTIVE=$MAX_SAMPLES fi - CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) - START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) - END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) - [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS )) + REMAINDER=$(( EFFECTIVE % NUM_CHUNKS )) + if [ $CHUNK_ID -lt $REMAINDER ]; then + CHUNK_SIZE=$(( BASE_SIZE + 1 )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + else + CHUNK_SIZE=$BASE_SIZE + START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 )) + fi + END_LINE=$(( START_LINE + CHUNK_SIZE - 1 )) CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" diff --git a/tests/fixtures/rollout/client_cmd_policy_only.txt b/tests/fixtures/rollout/client_cmd_policy_only.txt index 69062a3..d44c815 100644 --- a/tests/fixtures/rollout/client_cmd_policy_only.txt +++ b/tests/fixtures/rollout/client_cmd_policy_only.txt @@ -102,10 +102,16 @@ if [ $NUM_CHUNKS -gt 1 ]; then if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then EFFECTIVE=$MAX_SAMPLES fi - CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) - START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) - END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) - [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + BASE_SIZE=$(( EFFECTIVE / NUM_CHUNKS )) + REMAINDER=$(( EFFECTIVE % NUM_CHUNKS )) + if [ $CHUNK_ID -lt $REMAINDER ]; then + CHUNK_SIZE=$(( BASE_SIZE + 1 )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + else + CHUNK_SIZE=$BASE_SIZE + START_LINE=$(( REMAINDER * (BASE_SIZE + 1) + (CHUNK_ID - REMAINDER) * BASE_SIZE + 1 )) + fi + END_LINE=$(( START_LINE + CHUNK_SIZE - 1 )) CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" diff --git a/tests/test_rollout.py b/tests/test_rollout.py index be6b327..f267fd1 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -38,6 +38,7 @@ from __future__ import annotations +import subprocess from pathlib import Path import pytest @@ -46,6 +47,7 @@ from nvflow.lib.rl.helpers import LauncherFS from nvflow.lib.rl.rollout import ( + _build_chunk_slice_segment, _build_client_cmd, _build_merge_cmd, _build_port_read_preamble, @@ -832,3 +834,114 @@ def test_get_remaining_jobs_mixed_state(tmp_path: Path) -> None: remaining = _get_remaining_jobs(LauncherFS({}), str(tmp_path), seeds, chunks, rerun_done=False) assert remaining == [(1, 1)] + + +# =========================================================================== +# _build_chunk_slice_segment — runtime chunk-boundary arithmetic +# =========================================================================== +# +# These tests actually execute the rendered bash against real files rather +# than snapshotting the text, because the bug class here is arithmetic, not +# formatting. The previous ceiling-division formula gave +# EFFECTIVE=4/NUM_CHUNKS=3 a last chunk with START_LINE(5) > END_LINE(4) -- +# a silently empty chunk file -- which then tripped the merge step's own +# "chunk .done exists but file missing/empty" safety check (see the +# _build_merge_cmd tests above) as if rollout collection had failed +# upstream, when nothing had actually gone wrong. A text snapshot would +# not catch a wrong split point; only running the arithmetic does. + + +def _run_chunk_slice( + tmp_path: Path, + total_lines: int, + num_chunks: int, + chunk_id: int, + max_num_samples: int = 0, +) -> list[str]: + """Execute the rendered chunk-slice segment against a real input file. + + Returns the lines (no trailing newline) written to this chunk's + output file, or ``[]`` if no chunk file was produced (either the + ``NUM_CHUNKS <= 1`` no-op path, or -- pre-fix -- an empty chunk). + """ + tmp_path.mkdir(parents=True, exist_ok=True) + input_data = tmp_path / "input.jsonl" + input_data.write_text("".join(f"line{i}\n" for i in range(1, total_lines + 1))) + output_dir = tmp_path / "out" + output_dir.mkdir() + + script = ( + "set -eu\n" + f'INPUT_DATA="{input_data}"\n' + f'OUTPUT_DIR="{output_dir}"\n' + f"NUM_CHUNKS={num_chunks}\n" + f"CHUNK_ID={chunk_id}\n" + ) + _build_chunk_slice_segment(max_num_samples) + + result = subprocess.run(["bash", "-c", script], capture_output=True, text=True) + assert result.returncode == 0, f"chunk-slice script failed: {result.stderr}" + + chunk_file = output_dir / f"chunk_input_chunk{chunk_id}.jsonl" + if not chunk_file.exists(): + return [] + return chunk_file.read_text().splitlines() + + +def test_chunk_slice_no_gaps_or_overlaps_across_shapes(tmp_path: Path) -> None: + """For every (total_lines, num_chunks) with 2 <= num_chunks <= total_lines, + concatenating all chunks in order must reproduce the input exactly -- + every line appears, in order, in exactly one chunk. + """ + for total_lines in range(2, 13): + for num_chunks in range(2, total_lines + 1): + all_lines: list[str] = [] + for chunk_id in range(num_chunks): + case_dir = tmp_path / f"n{total_lines}_k{num_chunks}_c{chunk_id}" + all_lines.extend(_run_chunk_slice(case_dir, total_lines, num_chunks, chunk_id)) + assert all_lines == [f"line{i}" for i in range(1, total_lines + 1)], ( + f"total_lines={total_lines}, num_chunks={num_chunks}" + ) + + +def test_chunk_slice_no_chunk_is_empty_when_num_chunks_le_total_lines(tmp_path: Path) -> None: + """Regression test for the ceiling-division bug: whenever there are at + least as many samples as chunks, every chunk must get at least one + line. Previously e.g. total_lines=4/num_chunks=3 left the last + chunk empty, and that empty chunk was then reported by the merge + step as an upstream rollout failure. + """ + for total_lines in range(2, 13): + for num_chunks in range(2, total_lines + 1): + for chunk_id in range(num_chunks): + case_dir = tmp_path / f"n{total_lines}_k{num_chunks}_c{chunk_id}" + lines = _run_chunk_slice(case_dir, total_lines, num_chunks, chunk_id) + assert lines, ( + f"chunk {chunk_id}/{num_chunks} was empty for total_lines={total_lines} " + "-- this is the previously-fixed ceiling-division bug" + ) + + +def test_chunk_slice_previously_buggy_case_four_lines_three_chunks(tmp_path: Path) -> None: + """Exact repro of the bug found in review: EFFECTIVE=4, NUM_CHUNKS=3 + used to give chunk 2 START_LINE(5) > END_LINE(4), producing a 0-byte + chunk file. Balanced-remainder division instead gives chunk sizes + [2, 1, 1] (the first EFFECTIVE % NUM_CHUNKS chunks get one extra line). + """ + assert _run_chunk_slice(tmp_path / "c0", 4, 3, 0) == ["line1", "line2"] + assert _run_chunk_slice(tmp_path / "c1", 4, 3, 1) == ["line3"] + assert _run_chunk_slice(tmp_path / "c2", 4, 3, 2) == ["line4"] + + +def test_chunk_slice_single_chunk_is_a_noop(tmp_path: Path) -> None: + """``NUM_CHUNKS <= 1`` must skip slicing entirely (no chunk file is + created) -- this behavior is unchanged by the arithmetic fix. + """ + assert _run_chunk_slice(tmp_path, 5, 1, 0) == [] + + +def test_chunk_slice_respects_max_num_samples_clamp(tmp_path: Path) -> None: + """``max_num_samples`` still clamps EFFECTIVE before chunking -- lines + beyond the clamp must not appear in any chunk. + """ + assert _run_chunk_slice(tmp_path / "c0", 10, 2, 0, max_num_samples=4) == ["line1", "line2"] + assert _run_chunk_slice(tmp_path / "c1", 10, 2, 1, max_num_samples=4) == ["line3", "line4"]