[DO NOT MERGE] trainer ft

[fzyzcjy]

ci-sglang-pr: #28524
ci-megatron-pr: #21

sgl-project/sglang#28524
radixark/Megatron-LM#21

[gemini-code-assist]

Code Review

This pull request introduces support for independent data parallelism (DP) by allowing DP replicas to function as independent Megatron instances. It refactors the Ray training group logic into cells, adds configuration for delayed data splitting, and implements consistency checks for dynamic global batch sizes. Feedback identifies a bug in the data splitting utility where the dynamic batch size was not propagated, suggests avoiding hardcoded configuration values, and recommends using strict zip operations for safer actor-critic connections.

[gemini-code-assist]

There are a couple of issues in split_train_data_by_dp:

1. The rollout_data variable initialized on lines 8-11 is unused because it's redefined inside the loop on line 24. This initial block can be removed for clarity.
2. More importantly, dynamic_global_batch_size is not being propagated to the split data dictionaries. It is added to the data dict in RolloutManager._convert_samples_to_train_data, but it's not one of the keys that are copied to each rollout_data split. This will cause issues downstream where this value is expected. It should be added to the list of keys that are copied to each split, similar to raw_reward and total_lengths.

Here is a suggested fix that addresses both points:

def split_train_data_by_dp(args, data: dict[str, Any], *, dp_size: int) -> list[dict[str, Any]]:
    """Split the train data by data parallel size."""
    total_lengths = [len(t) for t in data["tokens"]]
    data["total_lengths"] = total_lengths

    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",
        ]:
            if key not in data:
                continue
            val = [data[key][j] for j in partition]
            rollout_data[key] = val
        # keys that are not partitioned and should be present in all splits
        for key in [
            "raw_reward",
            "total_lengths",
            "dynamic_global_batch_size",
        ]:
            if key not in data:
                continue
            rollout_data[key] = data[key]
        ans.append(rollout_data)
    return ans

[gemini-code-assist]

The quorum_id is hardcoded to 0. According to repository guidelines, model parameters and configuration values should be retrieved from the model configuration rather than being hardcoded. This improves flexibility and ensures consistency across different environments.

References

1. Model parameters and configuration values should be retrieved from the model configuration rather than being hardcoded.

[gemini-code-assist]

The zip function is used with strict=False. However, in RayTrainGroup.connect, zip is used with strict=True when iterating over cells. For consistency and safety, it's better to use strict=True here as well. This will ensure that an error is raised if the number of actor and critic handles within a cell do not match, which would indicate a configuration problem.

Suggested change

	for actor, critic in zip(self._actor_handles, critic_group._actor_handles, strict=False)
	for actor, critic in zip(self._actor_handles, critic_group._actor_handles, strict=True)

To fix this, bind the temporary probe socket to loopback instead of all interfaces.
General approach: replace wildcard bind addresses ("" / "0.0.0.0") with a dedicated local interface, here 127.0.0.1, since this utility is used with MASTER_ADDR=localhost.

Best single change (without changing behavior): in tests/fast/dist_utils.py, update find_free_port() line 11 from:

• s.bind(("", 0))
  to:
• s.bind(("127.0.0.1", 0))

No new imports, methods, or dependencies are required.