[Train] Micro-batch scheduling on rollout side + variable global batch size (port slime #1926+#1930+#1933)

[aoshen02]

Ports the coupled THUDM/slime chain #1926 → #1930 → #1933 (one variable-global-batch feature line, three sequential PRs) into vime as one PR — they touch the same files interdependently (dp_schedule.py, actor.py, data.py, model.py, loss.py, ray/rollout.py), so splitting would mean reconciling the same diverged files three times and a merge-order hazard. One PR = one CI run.

What

• #1926 — move the per-step micro-batch schedule from the actor (train) side to _split_train_data_by_dp on the rollout side, computed once over the full global batch. New vime/utils/dp_schedule.py (build_dp_schedule, compute_dynamic_global_batch_size) holds the Ray/engine-free logic so it's unit-testable. Drops the DP-wide all_reduce(num_microbatches, MAX); get_data_iterator becomes a thin consumer reading num_microbatches / micro_batch_indices out of rollout_data.
• #1930 + #1933 [1/N, 2/N] — variable global batch size: per-step split by rollout id so each rollout yields a fixed training-step count regardless of how many samples it produced; CP-aware loss / metric handling.

vime reconciliation

• vime already carried --use-dynamic-global-batch-size / --disable-rollout-trim-samples (args + an inline _compute_dynamic_global_batch_size) from an earlier partial sync; #1926 removes that inline method in favor of dp_schedule.compute_dynamic_global_batch_size — applied cleanly.
• _split_train_data_by_dp keeps vime's diverged extra keys (rollout_routed_experts, teacher_log_probs, dynamic_global_batch_size).
• train_parallel_config extended with cp_size / vpp_size / microbatch_group_size_per_vp_stage (set in actor.py:92, read by build_dp_schedule).
• Two rollout.py import hunks hand-reconciled (vime's vLLM seam — no SGLangEngine import).
• Patch applied via path+import rewrite (slime/→vime/, slime.→vime.); /root/slime mount and SLIME_TEST_* envs preserved.

Tests / CI

Adds the chain's CPU tests to the always-on cpu job: test_dp_schedule, test_cp_utils, test_loss_cp_invariance, test_metric_report, test_metric_report_dist (+ tests/_cp_dist_helpers.py). j2 updated + yml re-rendered.

Validation

• py_compile clean across all 16 touched files; pinned pre-commit (ruff/isort/black) clean.
• cpu CI (1916 image): see PR comment once the run completes.
• ⚠️ This changes training micro-batch counts / bin-packing (numerics) → needs GPU validation before merge. GPU CI is currently off per maintainer instruction; cpu validation covers the shipped CPU tests only.

Upstream: THUDM/slime#1926, #1930, #1933. Part of the slime→vime sync tracked in #107 (chain registered there; not ticking boxes pending sign-off).

[gemini-code-assist]

Code Review

This pull request introduces a robust per-rollout token-weighted mean reduction mechanism for context parallel (CP) training, allowing sibling samples from a single rollout to be correctly aggregated even when split across micro-batches. It refactors the data-parallel scheduling logic into a pure-Python module (dp_schedule.py) supporting Virtual Pipeline Parallelism (VPP) and dynamic batching, accompanied by extensive unit and distributed tests. The review feedback highlights two important improvements: first, correcting the destination rank in dist.gather_object to use group-relative rank 0 instead of the global dp_src_rank to prevent crashes in multi-GPU setups; second, strengthening the step-size assertion in the DP scheduler to require at least align_to samples when VPP is enabled, avoiding downstream scheduling failures.

[gemini-code-assist]

In PyTorch distributed collective APIs (such as dist.gather_object), the dst parameter specifies the destination rank within the process group (dp_group), not the global rank.

Using the global rank dp_src_rank as dst will cause incorrect behavior or crashes (e.g., ValueError: Invalid destination rank) in multi-GPU/multi-node configurations where the DP source rank is non-zero or exceeds the size of the DP subgroup. Since the source rank is always the first rank in the DP group, its group-relative rank is always 0.

Please change dst=dp_src_rank to dst=0 to ensure correct group-relative routing.

Suggested change

	if dist.get_rank() == dp_src_rank:
	gathered = [None] * dp_size
	dist.gather_object(log_dict, gathered, dst=dp_src_rank, group=dp_group)
	reduced: dict = {}
	for key in log_dict:
	values = [d[key] for d in gathered]
	first = values[0]
	if isinstance(first, tuple) and len(first) == 2:
	total_sum = sum(v[0] for v in values)
	total_count = sum(v[1] for v in values)
	reduced[key] = total_sum / total_count if total_count else 0.0
	else:
	reduced[key] = sum(values) / dp_size
	return reduced
	dist.gather_object(log_dict, None, dst=dp_src_rank, group=dp_group)
	return None
	if dist.get_rank() == dp_src_rank:
	gathered = [None] * dp_size
	dist.gather_object(log_dict, gathered, dst=0, group=dp_group)
	reduced: dict = {}
	for key in log_dict:
	values = [d[key] for d in gathered]
	first = values[0]
	if isinstance(first, tuple) and len(first) == 2:
	total_sum = sum(v[0] for v in values)
	total_count = sum(v[1] for v in values)
	reduced[key] = total_sum / total_count if total_count else 0.0
	else:
	reduced[key] = sum(values) / dp_size
	return reduced
	dist.gather_object(log_dict, None, dst=0, group=dp_group)
	return None

[aoshen02]

This is a false positive — keeping as-is.

For dist.gather_object, dst is the global rank (PyTorch docstring: "Destination rank on global process group (regardless of group argument)"); the group-relative parameter is group_dst. The guard here is if dist.get_rank() == dp_src_rank, and dist.get_rank() returns the global rank — so dst=dp_src_rank is in the same (global) coordinate system and is correct.

Switching to dst=0 would actually introduce the bug you describe: on multi-node, the DP group's first rank is a non-zero global rank, so gathering to global rank 0 while the receiving branch is dp_src_rank would mismatch. This also matches slime upstream verbatim. No change.

[gemini-code-assist]

The assertion currently checks if len(sample_indices) >= dp_size. However, if Virtual Pipeline Parallelism (VPP) is enabled (vpp_size > 1), the schedule requires the number of microbatches to be a multiple of align_to = dp_size * mb_group.

Since each microbatch must contain at least one sample, the step must have at least align_to samples. If len(sample_indices) is between dp_size and align_to - 1, this assertion passes, but the scheduling downstream will fail with a confusing assertion error about maximal splitting or static path alignment.

Updating the assertion to check len(sample_indices) >= align_to provides a much clearer error message and prevents downstream failures.

Suggested change

	assert len(sample_indices) >= dp_size, (
	f"step {step_i}: {len(sample_indices)} samples < dp_size {dp_size}; "
	f"each step needs at least one sample per rank."
	)
	assert len(sample_indices) >= align_to, (
	f"step {step_i}: {len(sample_indices)} samples < align_to {align_to}; "
	f"each step needs at least {align_to} samples to satisfy DP and VPP constraints."
	)

[aoshen02]

Valid but minor — keeping >= dp_size for now. The downstream alignment assert (target_K/align_to) already rejects the VPP-with-too-few-samples case, just later in the call. Holding this byte-aligned with slime upstream to ease future syncs; noted as a follow-up if we want an earlier/clearer error message.

[aoshen02]

cpu CI — all green (fresh --rm container, latest cumulative image inferactinc/public:vime-vllm-cu129-sync1916, --gpus all):

group	result
test_dp_schedule	8 passed
test_cp_utils	5 passed
test_loss_cp_invariance	6 passed
test_metric_report	19 passed
test_metric_report_dist	27 passed
test_megatron_argument_validation	7 passed
plugin_contracts (rollout/runtime-hook/path-loading/generate)	6 / 10 / 14 / 4 = 34 passed
unit suite (per-file parallel)	0 failed of 11

The five tests the chain ships (dp_schedule scheduling, CP-utils, CP loss-invariance, metric reporting incl. distributed) exercise the core of the refactor on CPU. ⚠️ Still needs GPU validation before merge — this changes training micro-batch counts / bin-packing; GPU CI is off per maintainer instruction.

[aoshen02]

Closing for now (slime→vime sync being restructured). Branch is preserved — can reopen anytime.