fix: pair --log-correct-samples rewards with the DP-local samples

[keepkeen]

What

--log-correct-samples crashes with IndexError on DP > 1, and produces wrong numbers before it crashes. Fixes #1784.

Root cause

RolloutManager._split_train_data_by_dp ships two kinds of fields to the trainer:

# per-sample keys -> already sliced to this rank's samples
rollout_data[key] = [data[key][j] for j in partition]
# keys that need to be splited at train side
for key in ["raw_reward", "total_lengths"]:
    rollout_data[key] = data[key]

process_rollout_data then re-splits total_lengths but leaves raw_reward global. So in log_rollout_data, the --log-correct-samples block walks the global reward list while indexing this rank's local lists:

for i, raw_reward in enumerate(raw_rewards):   # len == rollout_batch_size * n_samples_per_prompt
    if raw_reward == 1:
        correct_response_lengths.append(response_lengths[i])   # len == that / dp_size

Two failure modes:

1. IndexError once i passes the end of this rank's lists — the traceback in [Bug] IndexError in log_rollout_data when --log-correct-samples enabled with DP > 1 #1784.
2. Silently wrong metrics before that point: sample i's reward is attributed to whichever sample happens to sit at local position i. This also affects dp_size == 1, because first-fit packing returns a permuted partition, not the identity.

Why not just partition raw_reward

raw_reward has to stay global — log_passrate reshapes it into [rollout_batch_size, n_samples_per_prompt] groups, which only works on the full rollout batch. So this keeps raw_reward as-is and adds local_raw_reward, the DP-local view, for the metrics that pair a reward with per-sample tensors.

Reproduction

Driving the real process_rollout_data with the partition layout _split_train_data_by_dp produces (8 samples, dp_size=2, odd samples correct):

--- dp_rank 0 (owns global samples [0, 2, 4, 6]) ---
  len(raw_reward)       = 8   <- global
  len(response_lengths) = 4   <- DP-local
  IndexError: list index out of range
  partial correct_response_lengths = [102, 106]
  expected correct lengths for this rank = []     # rank 0 owns no correct samples at all

Test

New CPU unit test tests/test_process_rollout_data.py, registered in the cpu-unittest CI job. It pins the three-way contract (total_lengths DP-local, raw_reward global, local_raw_reward DP-local and positionally aligned) and replays the --log-correct-samples selection loop across dp1-permuted / dp2 / dp4 partitions. It fails on main (6 of 8 cases) and passes with this change.

ray.get is monkeypatched to the identity so the test stays single-process — process_rollout_data only uses Ray to deref the per-rank Box.

Note for maintainers (not addressed here)

While tracing this I noticed the same block computes correct_response_lengths / correct_length/p* / correct_entropy and stores them into rollout_data after the only gather_log_data("rollout", ...) call in log_rollout_data, and never calls gather_log_data itself (unlike its siblings log_multi_turn_data / log_passrate). Its local log_dict is unused. So those metrics appear to never be emitted anywhere. That has been true since the block was introduced in #1192, and reducing them across ranks needs (sum, count) weighting because each rank has a different number of correct samples — a behaviour change I left out of this crash fix. Happy to send a follow-up PR if you'd like it fixed.

[keepkeen]

Hi @zhuzilin, could you please take a look when you have time? This fixes #1784 while preserving the global raw_reward grouping required by pass@k and adding a positionally aligned DP-local view for correct-sample metrics. The PR includes CPU regression tests across permuted DP1, DP2, and DP4 partitions; GitHub currently reports it as mergeable, and all CI checks pass. Thanks!