Add correctness-contract tests for train-data-pair-dispatch (+ extract schedule helpers)

[zhihengy]

Why

PR radixark#10 (train-data-pair-dispatch) reshaped diffusion-GRPO training data from a 2D (sample x SDE-step) grid + tiling into a flat list of train-pairs: rollout flattens samples into pairs -> contiguous equal split across DP ranks -> the trainer feeds contiguous micro-batches to the forward, normalizing gradients as loss_sum / num_local_pairs. Correctness rests on a set of invariants (flatten completeness, contiguous DP split + drop-tail, micro-batch schedule coverage, conditioning assembly equivalence, ...), none of which were guarded by tests. A regression would either silently change the gradient scale or make DP ranks run a different number of FSDP collectives and deadlock.

This PR adds tests for those invariants, plus a small enabling refactor. No behavior change, and CI stays CPU-only.

1. Refactor (commit 1)

Move build_microbatch_schedule and validate_same_microbatch_counts_across_dp out of actor.py (which imports ray / diffusers / sglang at module load) into miles/utils/train_data_utils.py (torch + types only), so they can be unit-tested on a CPU runner without the heavy import chain. Also factor the DP count comparison into a pure microbatch_counts_agree helper, decoupled from the all_gather_object collective so it is testable without a process group. Functions are moved verbatim; actor.py imports the same names.

2. CPU tests (commit 2) — tests/fast/utils/test_train_data_utils.py

register_cpu_ci(suite="stage-a-cpu", labels=[]) -> runs on every PR. 23 tests:

• flatten: one pair per selected step, sample-major order, latent/next_latent offset, advantage broadcast (vs separate raw_reward), scheduler meta from the first sample; required-step-indices and empty-output guards.
• DP split: contiguous equal shards, drop-tail + warning, scheduler-meta broadcast, sigmas-optional, raises on num_pairs < dp_size / dp_size <= 0.
• microbatch schedule: golden, perfect-cover (no gap/overlap), short trailing batch, absolute offsets; equal shards -> equal per-rank micro-batch counts (the deadlock-prevention invariant).
• scheduler-meta reconstruction: fallback sigmas, stored sigmas, missing timesteps.

3. GPU tests (commit 3) — tests/e2e/test_train_pair_cond_equivalence.py

register_cuda_ci — not wired into any PR stage (no GPU runner; CI stays CPU-only). Run locally: pytest tests/e2e/test_train_pair_cond_equivalence.py. 10 tests on a tiny random-weight QwenImageTransformer2DModel:

• conditioning assembly: expand vs collate build identical encoder_hidden_states for a same-sample micro-batch (exact, torch.equal); collate pads + masks variable lengths.
• a same-sample forward is identical whichever path assembled the cond (exact).
• batching different-length samples together (pad + mask) leaves each sample's forward unchanged vs running it alone (tolerance — cross-batch fp rounding).
• classifier-free guidance: joint positive+negative collate order/mask, and the cfg_combine formula including the true-cfg rescale branch.
• the per-pair forward is invariant to micro_batch_size (tolerance).
• driving the real _forward_train_pair_batch: the accumulated gradient is invariant to micro_batch_size (tolerance).

torch.equal is used where the comparison keeps batch size + inputs constant; allclose only where it spans different batch sizes (on GPU, cuBLAS picks different kernels/reduction orders per shape, so cross-batch results differ at fp rounding — exact equality is unattainable there).

Local test results

• CPU: 23 passed (~1s); run_suite.py --hw cpu --suite stage-a-cpu discovers the file.
• GPU: 10 passed (B200).

CI impact

No changes to .github/workflows — CI remains CPU-only. The CPU tests auto-run on every PR via stage-a-cpu; the GPU tests are register_cuda_ci and run manually until a GPU runner exists.

🤖 Generated with Claude Code