fix(dp_schedule): drop trailing rollouts when the aligned micro-batch target exceeds the sample count

[EazyReal]

What changed

In build_dp_schedule (slime/utils/dp_schedule.py), when use_dynamic_batch_size and align_to > 1, drop the minimum number of whole trailing rollouts before aligning the micro-batch count, so the up-rounded target (ceil(K/align_to)*align_to) stays reachable. The kept rollout count is recorded in global_batch_sizes[s] (previously always the constant global_batch_size) so per-step loss scaling and LR increments track the samples that actually train. If rollout atomicity makes alignment impossible (no whole-rollout prefix has a sample count divisible by align_to), the bare AssertionError is replaced with a ValueError naming the step, kept sample count, and align_to.

The micro-batch target computation is factored into _aligned_target, and per-step sample collection / packing into _collect_step_samples / _pack, reused across the drop loop.

Why

In the long-trajectory dynamic-batch regime, first-fit packing puts roughly one sample per micro-batch, so the micro-batch count equals the sample count. When align_to (= dp_size, or dp_size * mb_group under VPP) rounds that count up past the sample count, expand_bins_by_splitting cannot reach the target — it stops early once every bin is a singleton — and build_dp_schedule died on a bare AssertionError. Concretely at DP4, a step of 514 single-sample bins needs 516 micro-batches, which is unreachable.

Dropping whole trailing rollouts keeps GRPO groups atomic (a multi-sample rollout is never split across the drop) and keeps the kept rollouts a contiguous 0..k-1 prefix. The drop floor is sample-count based: keep dropping while the target is unreachable as long as a full aligned block (>= align_to samples) remains. The loop is a no-op on the common one-sample-per-rollout and static-batch paths, where the aligned target never exceeds the sample count and global_batch_sizes stays equal to global_batch_size.

Validation

tests/test_dp_schedule.py is CPU-only / torch-free and already registered in the cpu-unittest job. Three new @pytest.mark.unit cases:

• 512 rollouts with 2 fanning out to 2 samples (514 samples) at DP4: drops the 2 trailing rollouts to a 510-rollout / 512-sample prefix that tiles by 4; asserts global_batch_sizes == [510], equal micro-batch count per rank, no rollout split.
• Ragged sizes [3,2,3,1,2] at DP4: drops past the naive dp_size floor down to the 3-rollout / 8-sample prefix that aligns.
• Ragged sizes [1,1,4,1,3,1] at DP4 (prefix sample counts 1,2,6,7,10,11, none divisible by 4): asserts the actionable ValueError.

Run: python tests/test_dp_schedule.py (or pytest tests/test_dp_schedule.py). The existing suite, including the FLOPs-balancing cases, continues to pass.

[EazyReal]

@zhuzilin could you review this one? The DP schedule can align the micro-batch target above the available sample count; this drops only the trailing rollouts that cannot fit, avoiding invalid empty/misaligned groups.

[EazyReal]

@zhuzilin refreshed on latest main and checks are green. This drops trailing rollout groups when the aligned micro-batch target exceeds the sample count, preventing invalid DP scheduling. Could you review when you have bandwidth?