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

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` | |

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

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
 

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]:

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(

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))