[DMD2 3/7] feat(rollout): stop the rollout at a random step for DMD2 backward simulation

[zhihengy]

Third of the DMD2 stacked series, on top of #330. A step strategy can now stop the engine at step i, so one rollout returns the student's input x_{σ_i} and its x0 prediction x̂0_i: the data that DMD2's backward simulation trains on.

What

• Step strategies return (sde, ret, stop). stop is sent as rollout_stop_step_index. The three built-ins return stop=None, so their requests are unchanged. A custom strategy must now return the third element; None keeps the old behavior.
• New built-in stop_at_random_step. It draws i ~ U{0..N-1} from the request seed and returns ret=[i, i+1], stop=i and no SDE filter. The engine runs steps 0..i and sends step i straight to σ=0. The trajectory then holds x_{σ_i} followed by x̂0_i as its final latent. With --diffusion-sde-type cps --diffusion-noise-level 1.0, every earlier step re-noises its x0 prediction with fresh noise, which is the DMD2 multi-step sampler.
• Docs: the strategy signature in customization.md and cli-reference.md, and the new built-in.

Nothing consumes this yet. The critic and actor PRs (4/7, 5/7) read i = latent_step_indices[0] and σ_i = sigmas[i] from the trajectory.

Why

• DMD2 trains the student at a random step i, on an input the student produced itself (backward simulation). The first i steps are a no-grad rollout, so they belong in the engine. This is the same split as Flow-GRPO, whose selected SDE steps are picked miles-side by --diffusion-step-strategy-path.
• The critic trains on x̂0_i, and the generator recomputes it from x_{σ_i}, so the rollout has to return both. Truncating the schedule makes x̂0_i the final latent. That needs no new response field, and the engine runs i+1 DiT forwards instead of N.
• Re-noising with fresh noise is the existing cps transition at noise_level=1: x_next = (1-σ') x̂0 + σ' ε. At σ'=0 it returns x̂0 exactly; so does ode.

sglang side

sgl-project/sglang#45230 is one commit on sglang-miles bff3f6ae7 (+94 lines, 34 of them non-test). Merge it and bump the image before this PR: an engine without it drops the unknown field and silently runs every step.

• RolloutRequest / SamplingParams get a new rollout_stop_step_index field.
• After set_timesteps, TimestepPreparationStage truncates the schedule to timesteps[:i+1] and sigmas[:i+1] + terminal. It rejects i outside the schedule.
• The field is part of the stage's dedup fingerprint, so grouped requests with different stops don't share a schedule.

Validation

Engine A/B on FLUX.2-klein-base-4B (1×B200, 512², 4 steps, cps / noise 1.0, guidance 1, debug tensors on). SD3.5-M was not run because its weights are gated and unavailable on this devbox. FLUX.2 also tests a resolution-dependent mu shift. Requests come from stop_at_random_step and are parsed by apply_rollout_image_response. Each one is compared with a full 4-step rollout that uses the same seed:

stop i	returned σ	latent_step_indices	x_{σ_i} vs full run	final latent vs x_i - σ_i·v_i (bf16)	DiT forwards
0	1.0, 0	[0, 1]	bitwise equal	bitwise equal	1
1	1.0, 0.9581, 0	[1, 2]	bitwise equal	bitwise equal	2
2	1.0, 0.9581, 0.884, 0	[2, 3]	bitwise equal	bitwise equal	3
3	1.0, 0.9581, 0.884, 0.7175, 0	[3, 4]	bitwise equal	bitwise equal	4

• Without the field, the sglang branch is bitwise equal to sglang-miles bff3f6ae7 on a full rollout: latents, sigmas, model outputs and the decoded image.
• sglang unit tests: 53 passed in test_scheduler_rollout_unit.py and test_rollout_api.py. The new case checks that a stopped rollout replays the full run's prefix bit for bit and ends on the x0 prediction. It fails if the terminal sigma isn't re-appended.

Tests

• tests/fast/rollout/test_step_strategy_hub.py: the new strategy returns [i, i+1] and stop=i, and it draws every step over seeds. The built-ins return stop=None.

Checklist

• pre-commit hooks pass on the changed files (both repos)
• Added/updated tests for new behaviour
• pytest tests/fast is green (518 passed, 6 skipped)
• If launch flags changed, python3 train_diffusion.py --help still parses (no flag added; the --diffusion-step-strategy-path contract changed)
• If a public flag was added, it appears in the CLI reference docs (signature updated in cli-reference.md and customization.md)
• If an example was added, it has a real walkthrough (no examples added)

Stack

DMD2: distill a flow-matching teacher into a few-step student with distribution matching. Merge in order:

1. [DMD2 1/7] [DMD2 1/7] feat(ray): train several colocated trainer groups in the sync driver #329: several colocated trainer groups in the sync driver
2. [DMD2 2/7] [DMD2 2/7] feat(ray): hand earlier trainer groups' outputs to later groups as external_data #330: hand earlier trainer groups' outputs to later groups as external_data
3. [DMD2 3/7] this PR: few-step rollout and the DMD2 step strategy
4. [DMD2 4/7] [DMD2 4/7] feat(dmd2): train a fake-score critic group from the teacher and score the student's samples #333: DMD2 critic group: critic training and teacher/critic scoring
5. [DMD2 5/7] [DMD2 5/7] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores #334: --loss-type dmd2 on the actor
6. [DMD2 6/7] SD3.5-Medium 4-step recipe and eval (not opened yet)
7. [DMD2 7/7] GAN head (phase 2, not opened yet)