diff --git a/examples/diffusion/bagel/bagel_trainside_lora.yaml b/examples/diffusion/bagel/bagel_trainside_lora.yaml index 7757b6ad1..b2e7fd21a 100644 --- a/examples/diffusion/bagel/bagel_trainside_lora.yaml +++ b/examples/diffusion/bagel/bagel_trainside_lora.yaml @@ -168,7 +168,7 @@ sampling: seed: 42 init_same_noise: false # per-sample x_T (NoiseRecipe r{rollout}:{sample}); diverse GRPO groups # SDE steps: 2 random steps drawn from the early half [0, 7) of the 14-step schedule, - # resolved PER ROLLOUT via AllSDEScheduler → resolve_sde_indices (shared across the + # resolved PER ROLLOUT via AllSDEScheduler → get_sde_indices (shared across the # GRPO group, seeded by rollout_id). Uses UniRL's NATIVE timestep_fraction (early-half # restriction) + num_sde_steps (random-k) — zero framework change. Differs from # flow_grpo only in that the 2 SDE steps are SCATTERED, not a contiguous window; diff --git a/examples/diffusion/minimax_h3/minimax_h3_t2va_nft.yaml b/examples/diffusion/minimax_h3/minimax_h3_t2va_nft.yaml index 081315468..11f15a2b3 100644 --- a/examples/diffusion/minimax_h3/minimax_h3_t2va_nft.yaml +++ b/examples/diffusion/minimax_h3/minimax_h3_t2va_nft.yaml @@ -223,7 +223,7 @@ sampling: logprob_precision: fp32 scheduler: # num_sde_steps: 0 IS the forward-process switch, with eta: 0.0 above — it - # makes resolve_sde_indices return [], read as "pure ODE, record no + # makes get_sde_indices return [], read as "pure ODE, record no # log-probs", and makes the stage store only the terminal clean latent. _target_: unirl.sde.index_schedule.AllSDEScheduler num_timesteps: ${..num_inference_steps} diff --git a/examples/pe/pe_trainside_pickscore.yaml b/examples/pe/pe_trainside_pickscore.yaml index d5917e679..3cb88e6e8 100644 --- a/examples/pe/pe_trainside_pickscore.yaml +++ b/examples/pe/pe_trainside_pickscore.yaml @@ -243,7 +243,7 @@ sampling: # Sparse SDE-index schedule (FlowDPPO-style, mirrors sd3_flowdppo.yaml): # num_sde_steps random steps drawn from the first half of the denoising # schedule, resolved PER ROLLOUT (seeded by rollout_id) by - # PETrainer._build_req via resolve_sde_indices. Replaces the static all-SDE + # PETrainer._build_req via get_sde_indices. Replaces the static all-SDE # list with selective stochastic steps. scheduler: _target_: unirl.sde.index_schedule.AllSDEScheduler diff --git a/unirl/sde/README.md b/unirl/sde/README.md index 0f1278cf8..3ec5ef00c 100644 --- a/unirl/sde/README.md +++ b/unirl/sde/README.md @@ -45,7 +45,7 @@ adapter's private choice. match SGLang). - **SDE indices own the policy-gradient density.** The selection is a `TimestepScheduler` (`index_schedule.py`) wired under `sampling.scheduler`, which - `DiffusionSamplingParams.resolve_sde_indices` asks once per rollout id. Selected + `DiffusionSamplingParams.get_sde_indices` asks once per rollout id. Selected `sde_indices` get a stochastic transition with a real per-step Gaussian log-prob (→ `LatentSegment.sde_logp`). Trainside loops collapse the same SDE kernel to Euler with `eta=0`; the FastVideo adapter instead implements the model @@ -123,3 +123,9 @@ MixGRPO keeps `FlowSDEStrategy` and adds a `WindowScheduler` under *disabled*, not "stretch to zero" — normalize falsy values to `None`. - **`compute_mu` is the single per-model μ override point.** FLUX.2-klein's μ depends on **both** `image_seq_len` and `num_inference_steps`, unlike the base formula. +- **`WindowScheduler` indexes the denoising loop.** `num_timesteps` is the full step count + and only full windows are visited. MixGRPO's `max_timesteps = sampling_steps - 2` indexes + that repo's sample dict, a different coordinate space. +- **`exp_decay_threshold` defaults to 13.** Decay starts only once a window start passes it. + Shipped mixgrpo schedules are shorter, so `strategy=exp_decay` stays at `iters_per_window` + until the threshold is lowered. diff --git a/unirl/sde/index_schedule.py b/unirl/sde/index_schedule.py index 3b4c7fb54..4d5592eb2 100644 --- a/unirl/sde/index_schedule.py +++ b/unirl/sde/index_schedule.py @@ -1,8 +1,11 @@ """Index schedulers used by GRPO-style algorithms.""" +import bisect +import math from abc import ABC, abstractmethod from collections.abc import Sequence from dataclasses import dataclass +from itertools import accumulate from typing import List, Literal, Optional, Set, Tuple, Union import numpy as np @@ -26,11 +29,29 @@ class WindowConfig: exp_decay_k: float = 0.1 def __post_init__(self) -> None: + if self.strategy != "all" and self.window_size < 1: + raise ValueError(f"WindowConfig requires window_size >= 1, got {self.window_size}") + if self.strategy in ("progressive", "decay", "exp_decay"): + if not 0 <= self.overlap_size < self.window_size: + raise ValueError( + f"WindowConfig({self.strategy}) requires 0 <= overlap_size < window_size, " + f"got overlap_size={self.overlap_size}, window_size={self.window_size}" + ) + if not math.isfinite(self.iters_per_window) or self.iters_per_window < 1: + raise ValueError(f"WindowConfig requires a finite iters_per_window >= 1, got {self.iters_per_window}") if self.strategy == "decay": if self.max_iters_per_window is None: self.max_iters_per_window = self.iters_per_window if self.min_iters_per_window is None: self.min_iters_per_window = max(1, self.iters_per_window // 4) + lo, hi = self.min_iters_per_window, self.max_iters_per_window + if not math.isfinite(lo) or not math.isfinite(hi) or lo < 1 or hi < lo: + raise ValueError( + "WindowConfig(decay) requires finite bounds with " + f"1 <= min_iters_per_window <= max_iters_per_window, got ({lo}, {hi})" + ) + if self.strategy == "exp_decay" and not math.isfinite(self.exp_decay_k): + raise ValueError(f"WindowConfig(exp_decay) requires a finite exp_decay_k, got {self.exp_decay_k}") class TimestepScheduler(ABC): @@ -106,6 +127,8 @@ class WindowScheduler(TimestepScheduler): "all": None, "progressive": "_resolve_progressive", "random": "_resolve_random", + "decay": "_resolve_sliding", + "exp_decay": "_resolve_sliding", } def __init__(self, num_timesteps: int, config: WindowConfig): @@ -127,20 +150,48 @@ def get_sde_indices(self, step: Optional[int] = None) -> Set[int]: return resolve_method(0 if step is None else int(step)) def _resolve_progressive(self, step: int) -> Set[int]: + starts = self._window_starts() window_step = step // self.config.iters_per_window - stride = self.config.window_size - self.config.overlap_size - remaining = self.num_timesteps - self.config.init_timestep - self.config.window_size - num_one_round_window_steps = max(1, remaining // stride + 1) - if window_step >= num_one_round_window_steps and not self.config.roll_back: - window_step = num_one_round_window_steps - 1 - return self._resolve_progressive(window_step * self.config.iters_per_window) - - window_step = window_step % num_one_round_window_steps - cur_timestep = self.config.init_timestep + window_step * stride - return set(range(cur_timestep, cur_timestep + self.config.window_size)) + if window_step >= len(starts) and not self.config.roll_back: + window_step = len(starts) - 1 + else: + window_step = window_step % len(starts) + cur = starts[window_step] + return set(range(cur, cur + self.config.window_size)) def _resolve_random(self, step: int) -> Set[int]: rng = np.random.default_rng(step) max_start = max(0, self.num_timesteps - self.config.window_size) cur_timestep = int(rng.integers(0, max_start + 1)) return set(range(cur_timestep, cur_timestep + self.config.window_size)) + + def _window_starts(self) -> List[int]: + """Start index of every full window in one sweep, in visiting order.""" + stride = self.config.window_size - self.config.overlap_size + remaining = self.num_timesteps - self.config.init_timestep - self.config.window_size + count = max(1, remaining // stride + 1) + return [self.config.init_timestep + i * stride for i in range(count)] + + def _resolve_sliding(self, step: int) -> Set[int]: + """Walk full windows whose dwell is the decay or exp_decay value at each start.""" + starts = self._window_starts() + if self.config.strategy == "decay": + lo = self.config.min_iters_per_window + hi = self.config.max_iters_per_window + dwell = [] + for start in starts: + progress = start / self.num_timesteps + dwell.append(max(lo, int(hi * (1.0 - progress) + lo * progress))) + else: + base = self.config.iters_per_window + k = self.config.exp_decay_k + threshold = self.config.exp_decay_threshold + dwell = [int(math.ceil(base * math.exp(-k * max(0, start - threshold)))) for start in starts] + total = sum(dwell) + if step >= total: + if not self.config.roll_back: + cur = starts[-1] + return set(range(cur, cur + self.config.window_size)) + step = step % total + cur = starts[bisect.bisect_right(list(accumulate(dwell)), step)] + return set(range(cur, cur + self.config.window_size)) diff --git a/unirl/trainer/diffusion.py b/unirl/trainer/diffusion.py index 6611e9add..7af067499 100644 --- a/unirl/trainer/diffusion.py +++ b/unirl/trainer/diffusion.py @@ -692,7 +692,7 @@ def _build_request_sample( sp = sampling if sampling is not None else self.sampling_params noise_latent_shape = self._eval_noise_latent_shape if sampling is not None else self._noise_latent_shape diffusion = sp.get("diffusion") - sde_indices = diffusion.resolve_sde_indices(rollout_id) + sde_indices = diffusion.get_sde_indices(rollout_id) diffusion = dataclasses.replace( diffusion, sde_indices=sde_indices, scheduler=None, init_noise_latent_shape=noise_latent_shape ) diff --git a/unirl/trainer/pe.py b/unirl/trainer/pe.py index d6ecca9a5..a12f7638c 100644 --- a/unirl/trainer/pe.py +++ b/unirl/trainer/pe.py @@ -156,7 +156,7 @@ def _build_request_sample( base = sampling if sampling is not None else self.sampling_params diff_params = base.get("diffusion") ar_params = base.get("ar") - sde_indices = diff_params.resolve_sde_indices(rollout_id) + sde_indices = diff_params.get_sde_indices(rollout_id) diffusion = dataclasses.replace(diff_params, sde_indices=sde_indices, scheduler=None) request = prepare_input_sample( inputs, diff --git a/unirl/trainer/unified_model.py b/unirl/trainer/unified_model.py index 3a772f61e..98eabd520 100644 --- a/unirl/trainer/unified_model.py +++ b/unirl/trainer/unified_model.py @@ -200,7 +200,7 @@ def _build_request_sample( base = sampling if sampling is not None else self.sampling_params diff_params = base.get("diffusion") ar_params = base.get("ar") - sde_indices = diff_params.resolve_sde_indices(rollout_id) + sde_indices = diff_params.get_sde_indices(rollout_id) disable_xt = bool(os.environ.get("DISABLE_DRIVER_XT")) or bool(getattr(diff_params, "disable_driver_xt", False)) diffusion = dataclasses.replace( diff_params, sde_indices=sde_indices, scheduler=None, disable_driver_xt=disable_xt diff --git a/unirl/types/sampling.py b/unirl/types/sampling.py index 53fee4fef..72430bfab 100644 --- a/unirl/types/sampling.py +++ b/unirl/types/sampling.py @@ -182,8 +182,8 @@ def __post_init__(self) -> None: f"{cls}.sampler_kwargs cannot contain reserved keys {sorted(shadowed)}; set them as fields instead", ) - def resolve_sde_indices(self, rollout_id: int) -> List[int]: - """Resolve which denoising steps record SDE log-probs for ``rollout_id``.""" + def get_sde_indices(self, rollout_id: int) -> List[int]: + """SDE steps for ``rollout_id``: a preset list, else ``scheduler.get_sde_indices``, else every step.""" if self.sde_indices is not None: return [int(i) for i in self.sde_indices] scheduler: Optional[TimestepScheduler] = self.scheduler