Fix off-by-one in GRPO rollout chunk-slicing that produces empty chunks

[Anilreddy2309]

Summary

_build_chunk_slice_segment() (nvflow/lib/rl/rollout.py) computes each Slurm job's head|tail chunk boundaries with ceiling division:

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

Whenever EFFECTIVE doesn't divide evenly by NUM_CHUNKS, NUM_CHUNKS * CHUNK_SIZE overshoots EFFECTIVE, so the last chunk (or several, if NUM_CHUNKS > EFFECTIVE) gets START_LINE > END_LINE — a 0-byte chunk. Concretely: 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 a few lines later in the same file (added specifically to catch real upstream failures — see its "silent-success cascade" comment), 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.

Fix

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 whenever NUM_CHUNKS <= EFFECTIVE — by construction, REMAINDER*(BASE_SIZE+1) + (NUM_CHUNKS-REMAINDER)*BASE_SIZE = REMAINDER + NUM_CHUNKS*BASE_SIZE = EFFECTIVE exactly. 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 — that behavior is unchanged.

Testing

This is generated bash, not directly-callable Python logic, so I didn't rely on a text snapshot alone (a snapshot wouldn't catch a wrong split point, only a wrong render). tests/test_rollout.py's new tests actually execute the rendered segment via subprocess against real input files:

• Full coverage with no gaps/overlaps, total_lines 2–12 × 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 (unchanged)
• The max_num_samples clamp (unchanged)

I verified the new tests actually fail against the old ceiling-division formula (reproducing the exact reported bug: chunk 2/3 was empty for total_lines=4) before confirming they pass against the fix — see commit message for the exact repro.

_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 file is identical and scoped to exactly this arithmetic block (verified with a byte-for-byte diff comparison across all 5).

• ruff check / ruff format --check pass
• mypy passes
• License header check passes
• Full pytest tests/ suite passes, including the previously-skipped test_rollout.py (398 passed, 1 skipped, run against a minimal nemo_skills stub since the real package needs the full NeMo-RL/torch stack); the lightweight CI path (test_rollout.py skipped via pytest.importorskip, as designed) is unaffected — 339 passed, 3 skipped, matching the pre-change baseline exactly

🤖 Generated with Claude Code