From acfa478fa28400f1ad0afcac955105fc795b7566 Mon Sep 17 00:00:00 2001 From: Andy Ye <89891424+zhihengy@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:12:21 +0000 Subject: [PATCH] feat(rollout): stop the rollout at a random step for DMD2 backward simulation --- docs/user-guide/cli-reference.md | 2 +- docs/user-guide/customization.md | 6 ++-- miles/rollout/sglang_diffusion_rollout.py | 3 +- miles/rollout/step_strategy_hub.py | 32 +++++++++++++------- miles/utils/arguments.py | 2 +- tests/fast/rollout/test_step_strategy_hub.py | 25 ++++++++++++--- 6 files changed, 50 insertions(+), 20 deletions(-) diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index e438c99e8..9c9b8e0ac 100644 --- a/docs/user-guide/cli-reference.md +++ b/docs/user-guide/cli-reference.md @@ -200,7 +200,7 @@ See [Dtype Control](../advanced/dtype-control.md). | `--diffusion-num-sde-steps` | int | `0` | Read by the step strategy. | | `--diffusion-sde-window-range` | `"lo,hi"` | – | For `sde_window`. Defaults to `[0, num_inference_steps)`. | | `--diffusion-sde-candidate-steps` | `"1,2,3"` | – | Required by `epoch_global_random_choice`. | -| `--diffusion-step-strategy-path` | str | – | `(args, sample, num_steps, seed) -> (sde, ret)` choosing the SDE and returned steps, e.g. `miles.rollout.step_strategy_hub.sde_window`. Flow-GRPO needs one. | +| `--diffusion-step-strategy-path` | str | – | `(args, sample, num_steps, seed) -> (sde, ret, stop)` choosing the SDE, returned and last steps, e.g. `miles.rollout.step_strategy_hub.sde_window`. Flow-GRPO needs one. | | `--rollout-return-full-trajectory` | flag | off | Ask the engine for the whole denoising trajectory instead of only the steps the step strategy requests. The train tensors are identical; kept for debugging and A/B runs. | | `--diffusion-log-prob-no-const` | flag | off | Drop log-prob constants on the engine (pairs with the CPS backend). Forced on for `cps` / `ode`. | | `--diffusion-generator-device` | str | `cuda` | | diff --git a/docs/user-guide/customization.md b/docs/user-guide/customization.md index 88fa86b64..aaf1bfc7e 100644 --- a/docs/user-guide/customization.md +++ b/docs/user-guide/customization.md @@ -98,10 +98,11 @@ latents the engine ships back for the trainer. Stock implementations live in `miles/rollout/step_strategy_hub.py`. ```python -def strategy(args, sample, num_steps, seed) -> tuple[list[int] | None, list[int] | None]: - # returns (sde_step_indices, return_step_indices) +def strategy(args, sample, num_steps, seed) -> tuple[list[int] | None, list[int] | None, int | None]: + # returns (sde_step_indices, return_step_indices, stop_step_index) # sde None = no per-step filter: every step uses --diffusion-sde-type # return None = the engine returns every step's latent + # stop None = the engine runs every step; otherwise its last step goes straight to sigma 0 ... ``` @@ -113,6 +114,7 @@ steps for debugging / A-B runs. | `sde_window` | Random contiguous window (`--diffusion-num-sde-steps`, `--diffusion-sde-window-range`); returns all latents its steps touch | | `epoch_global_random_choice` | Subset of `--diffusion-sde-candidate-steps`, shared by each block of `--rollout-batch-size` consecutive prompt groups (one rollout, unless a dynamic-sampling filter drops groups); returns all latents its steps touch | | `ode_and_return_last` | No step filter (every step uses `--diffusion-sde-type`, which the DiffusionNFT recipes set to `ode`); return only the final clean `x0` | +| `stop_at_random_step` | DMD2 backward simulation: no step filter; stop at a random step `i` and return `x_i` plus its `x0` prediction | DiffusionNFT picks the timesteps it trains on with a separate hook, `--diffusion-nft-timestep-strategy-path`: `(args, num_steps, generator) -> list[int]`, built in diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 2b7ef5662..e952cc19a 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -176,7 +176,7 @@ async def generate_microgroup( # SGL-D TODO: support seed list for multiple samples in one request # currently only support assigning the first seed, SGL-D generates samples with seed, seed+1, seed+2, ... if not evaluation and state.step_strategy_fn is not None: - sde_indices, return_indices = state.step_strategy_fn( + sde_indices, return_indices, stop_index = state.step_strategy_fn( args, microgroup[0], int(sampling_params["num_inference_steps"]), @@ -186,6 +186,7 @@ async def generate_microgroup( return_indices = None # None asks the engine for every step sampling_params["rollout_sde_step_indices"] = sde_indices sampling_params["rollout_return_step_indices"] = return_indices + sampling_params["rollout_stop_step_index"] = stop_index else: sde_indices = None diff --git a/miles/rollout/step_strategy_hub.py b/miles/rollout/step_strategy_hub.py index 3ecc893a0..a59e3b3fe 100644 --- a/miles/rollout/step_strategy_hub.py +++ b/miles/rollout/step_strategy_hub.py @@ -1,11 +1,13 @@ """Library of functions that fill ``rollout_sde_step_indices`` / -``rollout_return_step_indices`` for one sglang-diffusion rollout request. +``rollout_return_step_indices`` / ``rollout_stop_step_index`` for one +sglang-diffusion rollout request. -Each function has signature ``(args, sample, num_steps, seed) -> (sde, ret)``, -both ``list[int] | None``. A strategy owns both halves of the request: +Each function has signature ``(args, sample, num_steps, seed) -> (sde, ret, stop)``, +``sde`` and ``ret`` are ``list[int] | None``, ``stop`` is ``int | None``. A strategy owns the whole request: ``sde`` picks the steps sampled as SDE (``None`` = no SDE steps, a pure ODE -rollout), and ``ret`` picks the trajectory latents the engine ships back for -the trainer (``None`` = all steps). +rollout), ``ret`` picks the trajectory latents the engine ships back for +the trainer (``None`` = all steps), and ``stop`` is the last step the engine +runs, sent straight to sigma 0 (``None`` = all steps). Point ``--diffusion-step-strategy-path`` at any such function. """ @@ -27,7 +29,7 @@ def _sde_steps_to_latent_indices(indices: list[int]) -> list[int]: def sde_window( args: Namespace, sample: Sample, num_steps: int, seed: int -) -> tuple[list[int] | None, list[int] | None]: +) -> tuple[list[int] | None, list[int] | None, int | None]: """flow_grpo-style random contiguous SDE window; the engine returns all latents the window's steps touch (x_i and x_{i+1} for each step i).""" window_size = int(args.diffusion_num_sde_steps) @@ -48,12 +50,12 @@ def sde_window( rng = np.random.default_rng(seed) start = int(rng.integers(lo, hi - window_size + 1)) indices = list(range(start, start + window_size)) - return indices, _sde_steps_to_latent_indices(indices) + return indices, _sde_steps_to_latent_indices(indices), None def epoch_global_random_choice( args: Namespace, sample: Sample, num_steps: int, seed: int -) -> tuple[list[int] | None, list[int] | None]: +) -> tuple[list[int] | None, list[int] | None, int | None]: """Per-epoch global random SDE subset: draw ``--diffusion-num-sde-steps`` candidate steps at random once per epoch (every sample in the epoch shares them), returned in ascending step order (matches the wan2.2 recipe).""" @@ -68,15 +70,23 @@ def epoch_global_random_choice( generator = torch.Generator().manual_seed(epoch + int(args.rollout_seed)) selected = torch.randperm(len(candidates), generator=generator)[:num_sde_steps] indices = sorted(candidates[i] for i in selected.tolist()) - return indices, _sde_steps_to_latent_indices(indices) + return indices, _sde_steps_to_latent_indices(indices), None def ode_and_return_last( args: Namespace, sample: Sample, num_steps: int, seed: int -) -> tuple[list[int] | None, list[int] | None]: +) -> tuple[list[int] | None, list[int] | None, int | None]: """DiffusionNFT-style request: no SDE steps (pure ODE rollout) and only the final clean x0 shipped back.""" - return None, [num_steps] + return None, [num_steps], None + + +def stop_at_random_step( + args: Namespace, sample: Sample, num_steps: int, seed: int +) -> tuple[list[int] | None, list[int] | None, int | None]: + """DMD2 backward simulation: stop at a random step i and ship x_i plus its x0 prediction.""" + stop = int(np.random.default_rng(seed).integers(num_steps)) + return None, [stop, stop + 1], stop def _sde_candidate_steps(args: Namespace, num_steps: int) -> list[int]: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index f4d759811..e7fc8ee93 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -437,7 +437,7 @@ def add_rollout_arguments(parser): type=str, default=None, help="Dotted path to `fn(args, sample, num_steps, seed) -> (sde_step_indices, " - "returned_latent_indices)`, called per engine request. The built-ins in " + "returned_latent_indices, stop_step_index)`, called per engine request. The built-ins in " "miles.rollout.step_strategy_hub read --diffusion-num-sde-steps.", ) parser.add_argument( diff --git a/tests/fast/rollout/test_step_strategy_hub.py b/tests/fast/rollout/test_step_strategy_hub.py index 312f5985e..29a256b70 100644 --- a/tests/fast/rollout/test_step_strategy_hub.py +++ b/tests/fast/rollout/test_step_strategy_hub.py @@ -7,16 +7,22 @@ from argparse import Namespace -from miles.rollout.step_strategy_hub import epoch_global_random_choice, ode_and_return_last, sde_window +from miles.rollout.step_strategy_hub import ( + epoch_global_random_choice, + ode_and_return_last, + sde_window, + stop_at_random_step, +) from miles.utils.types import Sample def test_sde_window_returns_window_plus_boundary(): args = Namespace(diffusion_num_sde_steps=3, diffusion_sde_window_range="0,10") - sde, ret = sde_window(args, Sample(), num_steps=10, seed=7) + sde, ret, stop = sde_window(args, Sample(), num_steps=10, seed=7) assert sde == list(range(sde[0], sde[0] + 3)) assert ret == sorted(set(sde) | {i + 1 for i in sde}) + assert stop is None def test_epoch_global_random_choice_returns_subset_plus_boundary(): @@ -26,11 +32,22 @@ def test_epoch_global_random_choice_returns_subset_plus_boundary(): rollout_batch_size=4, rollout_seed=0, ) - sde, ret = epoch_global_random_choice(args, Sample(group_index=0), num_steps=10, seed=7) + sde, ret, stop = epoch_global_random_choice(args, Sample(group_index=0), num_steps=10, seed=7) assert len(sde) == 2 and set(sde) <= {1, 3, 5, 7} assert ret == sorted(set(sde) | {i + 1 for i in sde}) + assert stop is None def test_ode_and_return_last_requests_only_x0(): - assert ode_and_return_last(Namespace(), Sample(), num_steps=10, seed=7) == (None, [10]) + assert ode_and_return_last(Namespace(), Sample(), num_steps=10, seed=7) == (None, [10], None) + + +def test_stop_at_random_step_ships_x_i_and_its_x0_prediction(): + stops = set() + for seed in range(32): + sde, ret, stop = stop_at_random_step(Namespace(), Sample(), num_steps=4, seed=seed) + assert sde is None + assert ret == [stop, stop + 1] + stops.add(stop) + assert stops == set(range(4))