Delay splitting train data by DP until actor-side processing

[fzyzcjy]

Extract the DP split into a witness-aware split_train_data_by_dp_raw
helper (with unit tests; the key list also carries seq_witness_ids) and
use it to split the training data on the actor side when
delay_split_train_data_by_dp is set, deferring the DP split from the
rollout side to actor-side processing. split_train_data_by_dp stays a
thin wrapper that ray.puts each partition.

• miles/ray/rollout/train_data_conversion.py (+ tests), miles/utils/data.py, actor_group.py, rollout_manager.py.

[gemini-code-assist]

Code Review

This pull request introduces a --delay-split-train-data-by-dp option to delay splitting training data by data parallel size until the training side, rather than doing it immediately during rollout generation. This involves moving the splitting logic to a new utility file miles/utils/data_utils.py and updating the rollout manager, actor group, and data processing pipeline accordingly. Feedback on the changes highlights several issues in the newly added split_train_data_by_dp function, including redundant variable initialization, unsafe in-place mutation of the input dictionary, and a lack of validation for sequence-level and token-level tensors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

There are three main issues in the current implementation of split_train_data_by_dp:

1. Dead Code / Redundant Initialization: The initialization of rollout_data and the assignment of rollout_data["prompt"] on lines 8-11 are completely redundant. rollout_data is re-initialized as an empty dictionary inside the loop on line 24, discarding any values set here. Furthermore, "prompt" is already correctly handled and sliced inside the loop (line 38).
2. In-place Mutation of Input Argument: Mutating the input data dictionary in-place (data["total_lengths"] = total_lengths on line 14) is a bad practice that can lead to side effects or runtime errors (e.g., if the dictionary or its values are read-only when retrieved from Ray's object store).
3. Missing Validation for Sequence/Token-Level Tensors: When processing lists of sequence-level or token-level tensors (such as rollout_log_probs and teacher_log_probs), we should validate that the list lengths match up front and perform per-sample shape checks to prevent silent mismatches or broadcast failures (e.g., GRPO-style scalar-broadcast traps).

We can resolve these issues cleanly by:

• Removing the redundant initialization on lines 8-11.
• Avoiding the in-place mutation of data.
• Directly assigning total_lengths to rollout_data["total_lengths"] inside the loop.
• Adding up-front validation of list lengths and per-sample shape checks for the sequence/token-level tensors.

    total_lengths = [len(t) for t in data["tokens"]]

    tensor_keys = ["rollout_log_probs", "teacher_log_probs"]
    active_tensor_keys = [k for k in tensor_keys if k in data]
    if active_tensor_keys:
        base_len = len(data["tokens"])
        for key in active_tensor_keys:
            if len(data[key]) != base_len:
                raise ValueError(f"Length mismatch for {key}: expected {base_len}, got {len(data[key])}")
            for idx, (t, tok) in enumerate(zip(data[key], data["tokens"])):
                if hasattr(t, "shape") and hasattr(tok, "shape"):
                    if t.shape[0] != tok.shape[0]:
                        raise ValueError(f"Shape mismatch at index {idx} for {key}: {t.shape} vs {tok.shape}")

    if args.balance_data:
        partitions = get_seqlen_balanced_partitions(total_lengths, dp_size, equal_size=True)
    else:
        partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)]

    ans = []

    for i in range(dp_size):
        rollout_data = {}
        partition = partitions[i]
        rollout_data["partition"] = partition
        for key in [
            "tokens",
            "multimodal_train_inputs",
            "response_lengths",
            "rewards",
            "truncated",
            "loss_masks",
            "round_number",
            "sample_indices",
            "rollout_log_probs",
            "rollout_routed_experts",
            "prompt",
            "teacher_log_probs",
            "seq_witness_ids",
            "weight_versions",
        ]:
            if key not in data:
                continue
            val = [data[key][j] for j in partition]
            rollout_data[key] = val
        for key in [
            "raw_reward",
            "dynamic_global_batch_size",
        ]:
            if key not in data:
                continue
            rollout_data[key] = data[key]
        rollout_data["total_lengths"] = total_lengths
        ans.append(rollout_data)
    return ans

References

1. When processing lists of sequence-level or token-level tensors (such as advantages, student log probabilities, and teacher log probabilities) in RL or distillation pipelines, validate that the list lengths match up front, and perform per-sample shape checks to prevent silent mismatches or broadcast failures (e.g., GRPO-style scalar-broadcast traps).

[Shi-Dong]

A few nits.

[Shi-Dong]

Hmm I could be wrong but it seems that the previous rollout_data["prompt"] = data["prompt"] is shadowed by this line?

[Shi-Dong]

I think it'd be better to add type annotation for rollout_data_ref.

[Shi-Dong]

Let's avoid relative import whenever possible.

[yueming-yuan]

this seems to be duplicated with the original split_train_data_by_dp? should we unify

[yueming-yuan]

rollout_indexer_topk is missing