[Sync][A] rollout data-model: micro-batch scheduling rollout-side + variable global batch + group_ids (slime #1926/#1930/#1933/#1941/#1959/#1965/#1969/#1984/#1962)

[aoshen02]

PR-A — rollout data-model + train-split (RFC #107)

Heaviest-logic mega-PR. Ports the coupled rollout-data-model feature line from THUDM/slime.

Included

• #1926 move micro-batch scheduling train→rollout side (vime/utils/dp_schedule.py, first-fit packing in seqlen_balancing.py, index-only DataIterator).
• #1930 + #1933 variable global batch size + per-token-loss reduction fix (cp_utils.reduce_train_step_metrics / get_sum_of_sample_mean(sample_denoms=...); loss denominator → step_global_batch_size).
• #1959 forge_load replay (vime/rollout/forge_load.py, --load-forge-rollout-data).
• #1965 group_ids fall back to range() not sample.index. #1969 --save-hf raw mode. #1941 multi-sample fanout test. #1984 rename rollout_ids→group_ids (non-agent files; agent half → PR-D). #1962 lint.

Validation (CPU)

• py_compile clean.
• test_metric_report_dist + test_loss_cp_invariance + test_cp_utils + test_metric_report = 57/57 incl. real multi-process gloo distributed across CP/DP ∈ {1,2,4} — these assert CP-invariance and rollout-side per-rollout-mean == train-side train_one_step for the #1930/#1933 numerics. test_dp_schedule 8/8, test_sample 16/16.
• GPU validation: stack-4 full matrix (this PR changes train numerics → high-info matrix).

Cross-PR note

test_sample reads args.vllm_speculative_config; relies on PR-C #1938's getattr guard (types.py). The integration stack order F+C+B+A satisfies this.

sglang→vLLM

Translated per the arg-map (--sglang-*→--vllm-*, sglang_rollout→vllm_rollout, vllm_speculative_config).

🤖 Generated with Claude Code

---

Fidelity ledger (audited 2026-06-04 vs slime 8ef1fb47..7a7aba4, via agent_run/scripts/completeness_oracle.sh)

op	detail
Ported	variable global-batch + per-token-loss + group_ids data model; CP per-token-loss reduction (cp_utils); the 6 numeric unit tests (dp_schedule/cp_utils/metric_report{,_dist}/loss_cp_invariance/sample)
Translated (sglang→vllm)	_split_train_data_by_dp(data, dp_size)→(data) reading dp_size from train_parallel_config + delegating to unit-testable dp_schedule.build_dp_schedule; --load-forge-rollout-data help reworded skip_sglang→skip_vllm / "sglang servers"→"vLLM servers, router, weight_update"
Dropped (faithful to slime)	--disable-rollout-trim-samples + --use-dynamic-global-batch-size (slime #1933 removed them too)
Fixed in this audit (commit 8803d18)	ported the missing tests/utils/test_hf_checkpoint_saver.py (slime #1969) + wired the 6 numeric tests + hf into the cpu CI job — they were present-but-unwired before
Verified	cpu suite green on the rebuilt FCBAGDE stack; hf test ALL_PASS (gb200 job 8918)

[gemini-code-assist]

Code Review

This pull request introduces support for compact/subagent fan-out patterns where a single rollout execution can emit multiple training samples sharing a single group ID. It implements a group-aware DP/microbatch scheduler, precomputes per-group mask sums to ensure correct loss denominators across micro-batches, and updates loss scaling and learning rate scheduling to use the actual step global batch size. Additionally, it adds a raw Megatron-to-HF checkpoint saver and comprehensive unit and end-to-end tests. Feedback from the review suggests interleaving micro-batch indices across virtual stages in VPP mode, replacing __getattribute__ overrides on Sample with standard properties to avoid performance overhead, guarding against division by zero in loss scaling, and supporting tuples in rollout output validation.

[gemini-code-assist]

In virtual pipeline parallel (VPP) training (vpp_size > 1), passing the same flat micro_batch_indices list to all DataIterator instances causes each virtual stage to process the exact same microbatches in the same order. This leads to duplicate data processing and incorrect training behavior.

To fix this, the microbatches should be interleaved across the virtual stages (e.g., micro_batch_indices[i::vpp_size] for stage i).

Suggested change

	def get_data_iterator(rollout_data: RolloutBatch) -> list[DataIterator]:
	"""Build one ``DataIterator`` per VPP stage from the pre-computed schedule in ``rollout_data``."""
	vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1
	micro_batch_indices = rollout_data["micro_batch_indices"]
	return [DataIterator(rollout_data, micro_batch_indices) for _ in range(vpp_size)]
	def get_data_iterator(rollout_data: RolloutBatch) -> list[DataIterator]:
	"""Build one ``DataIterator`` per VPP stage from the pre-computed schedule in ``rollout_data``."""
	vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1
	micro_batch_indices = rollout_data["micro_batch_indices"]
	return [DataIterator(rollout_data, micro_batch_indices[i::vpp_size]) for i in range(vpp_size)]

[gemini-code-assist]

Overriding __getattribute__ and __setattr__ to handle the deprecated rollout_id attribute introduces a significant performance overhead for every attribute access on Sample instances. In RL training, where millions of attributes are accessed, this can noticeably slow down data processing.

Using a standard Python property with a custom getter and setter achieves the exact same deprecation behavior with zero overhead for other attributes, and is much more idiomatic.

    @property
    def rollout_id(self):
        raise AttributeError("Sample.rollout_id is deprecated and write-only; use Sample.group_id instead.")

    @rollout_id.setter
    def rollout_id(self, value):
        if value is None:
            return
        warnings.warn(
            "Sample.rollout_id is deprecated and write-only; set Sample.group_id instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        self.group_id = value

[gemini-code-assist]

If step_global_batch_size is 0 (which can happen in distributed settings with empty partitions or heavy filtering), dividing by it will raise a ZeroDivisionError and crash the training run.

Consider guarding against division by zero by using max(step_global_batch_size, 1).

Suggested change

	if not args.calculate_per_token_loss:
	loss = (
	loss * num_microbatches / global_batch_size * mpu.get_data_parallel_world_size(with_context_parallel=True)
	loss
	* num_microbatches
	/ step_global_batch_size
	* mpu.get_data_parallel_world_size(with_context_parallel=True)
	)
	if not args.calculate_per_token_loss:
	loss = (
	loss
	* num_microbatches
	/ max(step_global_batch_size, 1)
	* mpu.get_data_parallel_world_size(with_context_parallel=True)
	)

[gemini-code-assist]

Strictly asserting isinstance(node, list) can cause unexpected crashes if a custom rollout function returns a tuple of Samples (which is common in Python). Supporting both list and tuple makes the validation much more robust.

def _validate_group_id_annotated(node, depth=0):
    """Walk the rollout function's nested output and validate ``group_id`` only
    when a compact / subagent pattern is detected.

    "Compact" = the rollout function wraps multiple training samples from one
    rollout execution into a ``list[Sample]``. In slime's convention the
    default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout)
    so its leaf ``list[Sample]`` lands at depth 1 and we skip validation,
    preserving backward compatibility. A compact rollout adds a third level:
    ``list[list[list[Sample]]]`` (prompt × rollout × samples-from-one-group),
    so the leaf ``list[Sample]`` lands at depth ≥ 2. At that point we require
    every sibling to carry a non-None ``group_id`` (or legacy ``rollout_id``)
    and to share the same value, so the loss reducer counts the group once
    instead of N times.
    """
    if isinstance(node, Sample):
        return
    assert isinstance(node, (list, tuple)), f"unexpected rollout output node type: {type(node).__name__}"
    if node and isinstance(node[0], Sample):
        if depth >= 2 and len(node) > 1:
            group_ids = [s.group_id for s in node]
            missing = [i for i, group_id in enumerate(group_ids) if group_id is None]
            assert not missing, (
                f"Compact rollout returned {len(node)} samples but group_id is unset on "
                f"positions {missing}. Set Sample.group_id on every sibling so the loss "
                "reducer can aggregate them as one group instead of N."
            )
            assert (
                len(set(group_ids)) == 1
            ), f"Sibling samples from one compact rollout must share group_id; got {group_ids}."
        return
    for item in node:
        _validate_group_id_annotated(item, depth + 1)

[CalvinXKY]

LGTM