From be612444f22d1f422ca9f43faa2b037e71a9d2dd Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 01/15] refactor(diffusion): rename scheduler_rl_mixin.py to flow_rollout_controller.py Pure file rename (git R100); the proof tool has no file-rename primitive. Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../schedulers/scheduling_flow_match_euler_discrete.py | 2 +- .../{scheduler_rl_mixin.py => flow_rollout_controller.py} | 0 .../runtime/post_training/rollout_denoising_mixin.py | 6 +++--- .../multimodal_gen/test/unit/test_scheduler_rollout_unit.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) rename python/sglang/multimodal_gen/runtime/post_training/{scheduler_rl_mixin.py => flow_rollout_controller.py} (100%) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 3e95069e59ae..60d70bef97d8 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -35,7 +35,7 @@ from sglang.multimodal_gen.runtime.models.schedulers.sigma_schedule import ( SigmaScheduleMixin, ) -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( SchedulerRLMixin, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py rename to python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index 674fb1a87933..fd3f43d16eca 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -10,14 +10,14 @@ import torch from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + SchedulerRLMixin, +) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDenoisingEnv, RolloutDitTrajectory, RolloutTrajectoryData, ) -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( - SchedulerRLMixin, -) from sglang.multimodal_gen.runtime.post_training.sp_utils import ( gather_stacked_latents_for_sp, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 21cdb2136468..943ae2d81297 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -4,8 +4,8 @@ import torch -import sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin as rl_mixin_module -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( +import sglang.multimodal_gen.runtime.post_training.flow_rollout_controller as rl_mixin_module +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( SchedulerRLMixin, ) From 7b7a6e5fbae9285a73016988c7d671c05faebd91 Mon Sep 17 00:00:00 2001 From: rockdu Date: Thu, 15 Oct 2026 12:13:52 -0700 Subject: [PATCH 02/15] refactor(move): move SchedulerRLDebugMixin into flow_rollout_controller.py Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 107 ++++++++++++++- .../post_training/scheduler_rl_debug_mixin.py | 122 ------------------ 2 files changed, 104 insertions(+), 125 deletions(-) delete mode 100644 python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index af1c520f10ae..02d901ec5f47 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -9,14 +9,13 @@ import torch from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, get_sp_world_size, ) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutDebugTensors, RolloutSessionData, ) -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_debug_mixin import ( - SchedulerRLDebugMixin, -) if TYPE_CHECKING: from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req @@ -24,6 +23,108 @@ _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) +class SchedulerRLDebugMixin: + @staticmethod + def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: + rollout_session_data.local_variance_noises = [] + rollout_session_data.local_prev_sample_means = [] + rollout_session_data.local_noise_std_devs = [] + rollout_session_data.local_model_outputs = [] + + def append_local_rollout_debug_tensors( + self, + batch, + *, + variance_noise: torch.Tensor, + prev_sample_mean: torch.Tensor, + noise_std_dev: torch.Tensor, + model_output: torch.Tensor, + ) -> None: + rollout_session_data = batch._rollout_session_data + batch_size = variance_noise.shape[0] + # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused. + rollout_session_data.local_variance_noises.append(variance_noise.clone()) + rollout_session_data.local_prev_sample_means.append(prev_sample_mean) + rollout_session_data.local_noise_std_devs.append( + noise_std_dev.expand((batch_size, 1)) + ) + rollout_session_data.local_model_outputs.append(model_output) + + def consume_local_rollout_debug_tensors( + self, + batch, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + rollout_session_data = batch._rollout_session_data + variance_noises = torch.stack(rollout_session_data.local_variance_noises, dim=1) + prev_sample_means = torch.stack( + rollout_session_data.local_prev_sample_means, dim=1 + ) + noise_std_devs = torch.stack(rollout_session_data.local_noise_std_devs, dim=1) + model_outputs = torch.stack(rollout_session_data.local_model_outputs, dim=1) + self._reset_rollout_debug_tensors(rollout_session_data) + return variance_noises, prev_sample_means, noise_std_devs, model_outputs + + def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: + """ + Consume rollout debug tensors and merge for all SP ranks. + + Returns rollout debug tensors with shape [B, T, ...]. + """ + rollout_session_data = batch._rollout_session_data + variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( + self.consume_local_rollout_debug_tensors(batch) + ) + + if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): + variance_noises = variance_noises.to(get_local_torch_device()) + prev_sample_means = prev_sample_means.to(get_local_torch_device()) + noise_std_devs = noise_std_devs.to(get_local_torch_device()) + model_outputs = model_outputs.to(get_local_torch_device()) + pipeline_config = rollout_session_data.pipeline_config + bsz, num_steps = variance_noises.shape[0], variance_noises.shape[1] + + # [B, T, ...] -> [B*T, ...] + variance_noises_packed = variance_noises.contiguous().reshape( + bsz * num_steps, *variance_noises.shape[2:] + ) + prev_sample_means_packed = prev_sample_means.contiguous().reshape( + bsz * num_steps, *prev_sample_means.shape[2:] + ) + model_outputs_packed = model_outputs.contiguous().reshape( + bsz * num_steps, *model_outputs.shape[2:] + ) + + # Gather on packed tensors first. + variance_noises_packed = pipeline_config.gather_latents_for_sp( + variance_noises_packed, batch=batch + ) + prev_sample_means_packed = pipeline_config.gather_latents_for_sp( + prev_sample_means_packed, batch=batch + ) + model_outputs_packed = pipeline_config.gather_latents_for_sp( + model_outputs_packed, batch=batch + ) + + # Unpack back to [B, T, ...]. + variance_noises = variance_noises_packed.reshape( + bsz, num_steps, *variance_noises_packed.shape[1:] + ) + prev_sample_means = prev_sample_means_packed.reshape( + bsz, num_steps, *prev_sample_means_packed.shape[1:] + ) + model_outputs = model_outputs_packed.reshape( + bsz, num_steps, *model_outputs_packed.shape[1:] + ) + # noise_std_devs is same on every device, not a sharded latent tensor. + + return RolloutDebugTensors( + rollout_variance_noises=variance_noises.cpu(), + rollout_prev_sample_means=prev_sample_means.cpu(), + rollout_noise_std_devs=noise_std_devs.cpu(), + rollout_model_outputs=model_outputs.cpu(), + ) + + class SchedulerRLMixin(SchedulerRLDebugMixin): @staticmethod def _get_rollout_session_data(batch) -> RolloutSessionData: diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py deleted file mode 100644 index b6e1222a455c..000000000000 --- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py +++ /dev/null @@ -1,122 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Debug tensor helpers for rollout-enabled schedulers.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -from sglang.multimodal_gen.runtime.distributed import ( - get_local_torch_device, - get_sp_world_size, -) -from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( - RolloutDebugTensors, - RolloutSessionData, -) - -if TYPE_CHECKING: - from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req - - -class SchedulerRLDebugMixin: - @staticmethod - def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: - rollout_session_data.local_variance_noises = [] - rollout_session_data.local_prev_sample_means = [] - rollout_session_data.local_noise_std_devs = [] - rollout_session_data.local_model_outputs = [] - - def append_local_rollout_debug_tensors( - self, - batch, - *, - variance_noise: torch.Tensor, - prev_sample_mean: torch.Tensor, - noise_std_dev: torch.Tensor, - model_output: torch.Tensor, - ) -> None: - rollout_session_data = batch._rollout_session_data - batch_size = variance_noise.shape[0] - # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused. - rollout_session_data.local_variance_noises.append(variance_noise.clone()) - rollout_session_data.local_prev_sample_means.append(prev_sample_mean) - rollout_session_data.local_noise_std_devs.append( - noise_std_dev.expand((batch_size, 1)) - ) - rollout_session_data.local_model_outputs.append(model_output) - - def consume_local_rollout_debug_tensors( - self, - batch, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - rollout_session_data = batch._rollout_session_data - variance_noises = torch.stack(rollout_session_data.local_variance_noises, dim=1) - prev_sample_means = torch.stack( - rollout_session_data.local_prev_sample_means, dim=1 - ) - noise_std_devs = torch.stack(rollout_session_data.local_noise_std_devs, dim=1) - model_outputs = torch.stack(rollout_session_data.local_model_outputs, dim=1) - self._reset_rollout_debug_tensors(rollout_session_data) - return variance_noises, prev_sample_means, noise_std_devs, model_outputs - - def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: - """ - Consume rollout debug tensors and merge for all SP ranks. - - Returns rollout debug tensors with shape [B, T, ...]. - """ - rollout_session_data = batch._rollout_session_data - variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( - self.consume_local_rollout_debug_tensors(batch) - ) - - if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): - variance_noises = variance_noises.to(get_local_torch_device()) - prev_sample_means = prev_sample_means.to(get_local_torch_device()) - noise_std_devs = noise_std_devs.to(get_local_torch_device()) - model_outputs = model_outputs.to(get_local_torch_device()) - pipeline_config = rollout_session_data.pipeline_config - bsz, num_steps = variance_noises.shape[0], variance_noises.shape[1] - - # [B, T, ...] -> [B*T, ...] - variance_noises_packed = variance_noises.contiguous().reshape( - bsz * num_steps, *variance_noises.shape[2:] - ) - prev_sample_means_packed = prev_sample_means.contiguous().reshape( - bsz * num_steps, *prev_sample_means.shape[2:] - ) - model_outputs_packed = model_outputs.contiguous().reshape( - bsz * num_steps, *model_outputs.shape[2:] - ) - - # Gather on packed tensors first. - variance_noises_packed = pipeline_config.gather_latents_for_sp( - variance_noises_packed, batch=batch - ) - prev_sample_means_packed = pipeline_config.gather_latents_for_sp( - prev_sample_means_packed, batch=batch - ) - model_outputs_packed = pipeline_config.gather_latents_for_sp( - model_outputs_packed, batch=batch - ) - - # Unpack back to [B, T, ...]. - variance_noises = variance_noises_packed.reshape( - bsz, num_steps, *variance_noises_packed.shape[1:] - ) - prev_sample_means = prev_sample_means_packed.reshape( - bsz, num_steps, *prev_sample_means_packed.shape[1:] - ) - model_outputs = model_outputs_packed.reshape( - bsz, num_steps, *model_outputs_packed.shape[1:] - ) - # noise_std_devs is same on every device, not a sharded latent tensor. - - return RolloutDebugTensors( - rollout_variance_noises=variance_noises.cpu(), - rollout_prev_sample_means=prev_sample_means.cpu(), - rollout_noise_std_devs=noise_std_devs.cpu(), - rollout_model_outputs=model_outputs.cpu(), - ) From a46a26755b245dd48ab12564706220a235e4f06e Mon Sep 17 00:00:00 2001 From: rockdu Date: Thu, 15 Oct 2026 12:14:10 -0700 Subject: [PATCH 03/15] refactor(move): extract the FlowMatch Euler update into _transition Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 46 ++++++++++++++----- 1 file changed, 35 insertions(+), 11 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 60d70bef97d8..10fde95c12e3 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -542,17 +542,16 @@ def step( batch, model_output, sample, current_sigma, next_sigma, generator ) else: - if self.config.stochastic_sampling: - x0 = sample - current_sigma * model_output - noise = randn_tensor( - sample.shape, - generator=generator, - device=sample.device, - dtype=sample.dtype, - ) - prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise - else: - prev_sample = sample + dt * model_output + prev_sample = self._transition( + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + dt=dt, + generator=generator, + step_index=self.step_index, + timestep=timestep, + ) # upon completion increase step index by one assert self._step_index is not None, "_step_index should not be None" @@ -566,6 +565,31 @@ def step( return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) + def _transition( + self, + *, + model_output: torch.FloatTensor, + sample: torch.FloatTensor, + current_sigma: torch.FloatTensor, + next_sigma: torch.FloatTensor, + dt: torch.FloatTensor, + generator: torch.Generator | list[torch.Generator] | None, + step_index: int, + timestep: int | torch.Tensor, + ) -> torch.FloatTensor: + if self.config.stochastic_sampling: + x0 = sample - current_sigma * model_output + noise = randn_tensor( + sample.shape, + generator=generator, + device=sample.device, + dtype=sample.dtype, + ) + prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise + else: + prev_sample = sample + dt * model_output + return prev_sample + def _time_shift_exponential( self, mu: float, sigma: float, t: torch.Tensor | np.ndarray ) -> torch.Tensor | np.ndarray: From 9ac446cd3cccc3e68f0e355a0352d1c0ac806fa2 Mon Sep 17 00:00:00 2001 From: rockdu Date: Thu, 15 Oct 2026 12:14:24 -0700 Subject: [PATCH 04/15] refactor(move): move RolloutSessionData into flow_rollout_controller.py Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 23 ++++++++++++++++++- .../runtime/post_training/rl_dataclasses.py | 23 +------------------ 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 02d901ec5f47..fe6ed4b35827 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -4,6 +4,7 @@ from __future__ import annotations import math +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Union import torch @@ -14,7 +15,6 @@ ) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDebugTensors, - RolloutSessionData, ) if TYPE_CHECKING: @@ -23,6 +23,27 @@ _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) +@dataclass +class RolloutSessionData: + """Per-batch rollout state created by prepare_rollout(), lives on the batch object. + + Cleared by setting ``batch._rollout_session_data = None``. + """ + + pipeline_config: Any = None + sigma_max: float = 0.0 + latents_shape: tuple | None = None + noise_buffer: torch.Tensor | None = None + + local_log_prob_sum: list[torch.Tensor] = field(default_factory=list) + local_log_prob_count: list[torch.Tensor] = field(default_factory=list) + + local_variance_noises: list[torch.Tensor] = field(default_factory=list) + local_prev_sample_means: list[torch.Tensor] = field(default_factory=list) + local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) + local_model_outputs: list[torch.Tensor] = field(default_factory=list) + + class SchedulerRLDebugMixin: @staticmethod def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: diff --git a/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py b/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py index 012c69c1fc9f..eeb4c39b98dc 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py @@ -1,33 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 """RL-specific dataclasses used by post-training and rollout paths.""" -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any import torch -@dataclass -class RolloutSessionData: - """Per-batch rollout state created by prepare_rollout(), lives on the batch object. - - Cleared by setting ``batch._rollout_session_data = None``. - """ - - pipeline_config: Any = None - sigma_max: float = 0.0 - latents_shape: tuple | None = None - noise_buffer: torch.Tensor | None = None - - local_log_prob_sum: list[torch.Tensor] = field(default_factory=list) - local_log_prob_count: list[torch.Tensor] = field(default_factory=list) - - local_variance_noises: list[torch.Tensor] = field(default_factory=list) - local_prev_sample_means: list[torch.Tensor] = field(default_factory=list) - local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) - local_model_outputs: list[torch.Tensor] = field(default_factory=list) - - @dataclass class RolloutDebugTensors: """Container for rollout debug tensors collected during denoising.""" From 5ae19774c2a44c79f866a7193cbc145948f22996 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 05/15] refactor(diffusion): prepare the rollout session methods to move onto RolloutSessionData Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 11 +- .../pipelines_core/stages/denoising.py | 1 - .../stages/model_specific_stages/cosmos3.py | 1 - .../post_training/flow_rollout_controller.py | 201 ++++++++---------- .../post_training/rollout_denoising_mixin.py | 8 +- .../test/unit/test_scheduler_rollout_unit.py | 113 +++++----- 6 files changed, 161 insertions(+), 174 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 10fde95c12e3..dc7d6a69b3ef 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -538,8 +538,15 @@ def step( if batch is not None and batch.rollout: if not self.already_prepared_rollout(batch): raise RuntimeError("Rollout not prepared before step") - prev_sample = self.flow_sde_sampling( - batch, model_output, sample, current_sigma, next_sigma, generator + prev_sample = SchedulerRLMixin.flow_sde_sampling( + self._get_rollout_session_data(batch), + batch, + model_output, + sample, + current_sigma, + next_sigma, + generator, + step_index=self.step_index, ) else: prev_sample = self._transition( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 205b032dafb2..4ce49db237b8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -2189,7 +2189,6 @@ def _run_denoising_loop( # pre-step value. Gated on batch.rollout to keep the # non-rollout path strictly untouched. if collect_trajectory and batch.rollout: - batch._rollout_loop_step_index = step_index self._maybe_append_dit_trajectory_step( batch=batch, latents=ctx.latents, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 842ba2357cde..f95485515015 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -1778,7 +1778,6 @@ def _denoise_once( if batch.rollout: # Capture the pre-step x_{t_i} before the scheduler advances it. - batch._rollout_loop_step_index = i self._maybe_append_dit_trajectory_step( batch=batch, latents=latents, diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index fe6ed4b35827..856e80b91131 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -10,12 +10,14 @@ import torch from sglang.multimodal_gen.runtime.distributed import ( - get_local_torch_device, get_sp_world_size, ) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDebugTensors, ) +from sglang.multimodal_gen.runtime.post_training.sp_utils import ( + gather_stacked_latents_for_sp, +) if TYPE_CHECKING: from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req @@ -44,99 +46,71 @@ class RolloutSessionData: local_model_outputs: list[torch.Tensor] = field(default_factory=list) -class SchedulerRLDebugMixin: +class SchedulerRLMixin: @staticmethod - def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: - rollout_session_data.local_variance_noises = [] - rollout_session_data.local_prev_sample_means = [] - rollout_session_data.local_noise_std_devs = [] - rollout_session_data.local_model_outputs = [] + def _reset_rollout_debug_tensors(self: RolloutSessionData) -> None: + self.local_variance_noises = [] + self.local_prev_sample_means = [] + self.local_noise_std_devs = [] + self.local_model_outputs = [] + @staticmethod def append_local_rollout_debug_tensors( - self, - batch, + self: RolloutSessionData, *, variance_noise: torch.Tensor, prev_sample_mean: torch.Tensor, noise_std_dev: torch.Tensor, model_output: torch.Tensor, ) -> None: - rollout_session_data = batch._rollout_session_data batch_size = variance_noise.shape[0] - # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused. - rollout_session_data.local_variance_noises.append(variance_noise.clone()) - rollout_session_data.local_prev_sample_means.append(prev_sample_mean) - rollout_session_data.local_noise_std_devs.append( - noise_std_dev.expand((batch_size, 1)) - ) - rollout_session_data.local_model_outputs.append(model_output) + # the underlying noise buffer in self.noise_buffer is reused. + self.local_variance_noises.append(variance_noise.clone()) + self.local_prev_sample_means.append(prev_sample_mean) + self.local_noise_std_devs.append(noise_std_dev.expand((batch_size, 1))) + self.local_model_outputs.append(model_output) + @staticmethod def consume_local_rollout_debug_tensors( - self, - batch, + self: RolloutSessionData, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - rollout_session_data = batch._rollout_session_data - variance_noises = torch.stack(rollout_session_data.local_variance_noises, dim=1) - prev_sample_means = torch.stack( - rollout_session_data.local_prev_sample_means, dim=1 - ) - noise_std_devs = torch.stack(rollout_session_data.local_noise_std_devs, dim=1) - model_outputs = torch.stack(rollout_session_data.local_model_outputs, dim=1) - self._reset_rollout_debug_tensors(rollout_session_data) + variance_noises = torch.stack(self.local_variance_noises, dim=1) + prev_sample_means = torch.stack(self.local_prev_sample_means, dim=1) + noise_std_devs = torch.stack(self.local_noise_std_devs, dim=1) + model_outputs = torch.stack(self.local_model_outputs, dim=1) + SchedulerRLMixin._reset_rollout_debug_tensors(self) return variance_noises, prev_sample_means, noise_std_devs, model_outputs - def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: + @staticmethod + def collect_rollout_debug_tensors( + self: RolloutSessionData, batch: Req + ) -> RolloutDebugTensors: """ Consume rollout debug tensors and merge for all SP ranks. Returns rollout debug tensors with shape [B, T, ...]. """ - rollout_session_data = batch._rollout_session_data variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( - self.consume_local_rollout_debug_tensors(batch) + SchedulerRLMixin.consume_local_rollout_debug_tensors(self) ) - if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): - variance_noises = variance_noises.to(get_local_torch_device()) - prev_sample_means = prev_sample_means.to(get_local_torch_device()) - noise_std_devs = noise_std_devs.to(get_local_torch_device()) - model_outputs = model_outputs.to(get_local_torch_device()) - pipeline_config = rollout_session_data.pipeline_config - bsz, num_steps = variance_noises.shape[0], variance_noises.shape[1] - - # [B, T, ...] -> [B*T, ...] - variance_noises_packed = variance_noises.contiguous().reshape( - bsz * num_steps, *variance_noises.shape[2:] - ) - prev_sample_means_packed = prev_sample_means.contiguous().reshape( - bsz * num_steps, *prev_sample_means.shape[2:] - ) - model_outputs_packed = model_outputs.contiguous().reshape( - bsz * num_steps, *model_outputs.shape[2:] - ) - - # Gather on packed tensors first. - variance_noises_packed = pipeline_config.gather_latents_for_sp( - variance_noises_packed, batch=batch - ) - prev_sample_means_packed = pipeline_config.gather_latents_for_sp( - prev_sample_means_packed, batch=batch - ) - model_outputs_packed = pipeline_config.gather_latents_for_sp( - model_outputs_packed, batch=batch - ) - - # Unpack back to [B, T, ...]. - variance_noises = variance_noises_packed.reshape( - bsz, num_steps, *variance_noises_packed.shape[1:] - ) - prev_sample_means = prev_sample_means_packed.reshape( - bsz, num_steps, *prev_sample_means_packed.shape[1:] - ) - model_outputs = model_outputs_packed.reshape( - bsz, num_steps, *model_outputs_packed.shape[1:] - ) - # noise_std_devs is same on every device, not a sharded latent tensor. + pipeline_config = self.pipeline_config + variance_noises = gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=variance_noises, + ) + prev_sample_means = gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=prev_sample_means, + ) + model_outputs = gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=model_outputs, + ) + # noise_std_devs is same on every device, not a sharded latent tensor. return RolloutDebugTensors( rollout_variance_noises=variance_noises.cpu(), @@ -145,8 +119,6 @@ def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: rollout_model_outputs=model_outputs.cpu(), ) - -class SchedulerRLMixin(SchedulerRLDebugMixin): @staticmethod def _get_rollout_session_data(batch) -> RolloutSessionData: """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" @@ -176,15 +148,15 @@ def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: def already_prepared_rollout(self, batch) -> bool: return getattr(batch, "_rollout_session_data", None) is not None + @staticmethod def _get_or_create_rollout_noise_buffer( - self, - rollout_session_data: RolloutSessionData, + self: RolloutSessionData, full_shape: tuple, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: """Get or create the reusable noise buffer (local or full shape) for rollout.""" - buffer = rollout_session_data.noise_buffer + buffer = self.noise_buffer if ( buffer is None or buffer.shape != full_shape @@ -192,11 +164,12 @@ def _get_or_create_rollout_noise_buffer( or buffer.device != device ): buffer = torch.empty(full_shape, device=device, dtype=dtype) - rollout_session_data.noise_buffer = buffer + self.noise_buffer = buffer return buffer + @staticmethod def _rollout_variance_noise( - self, + self: RolloutSessionData, batch, model_output: torch.FloatTensor, generator: Union[torch.Generator, list[torch.Generator]], @@ -204,7 +177,6 @@ def _rollout_variance_noise( """Generate variance noise for rollout. If generator is a list, use generator[i] for the i-th batch item.""" assert generator is not None, "Generator must be provided" - rollout_session_data = self._get_rollout_session_data(batch) device = model_output.device dtype = model_output.dtype local_shape = tuple(model_output.shape) @@ -218,17 +190,17 @@ def _rollout_variance_noise( "Generator list must have the same length as batch size" ) - buffer = self._get_or_create_rollout_noise_buffer( - rollout_session_data, rollout_session_data.latents_shape, device, dtype + buffer = SchedulerRLMixin._get_or_create_rollout_noise_buffer( + self, self.latents_shape, device, dtype ) for i in range(B): torch.randn( - rollout_session_data.latents_shape, + self.latents_shape, out=buffer[i : i + 1], generator=generator[i], ) - sharded_noise, _ = rollout_session_data.pipeline_config.shard_latents_for_sp( + sharded_noise, _ = self.pipeline_config.shard_latents_for_sp( batch=batch, latents=buffer ) if tuple(sharded_noise.shape) != local_shape: @@ -238,14 +210,16 @@ def _rollout_variance_noise( ) return sharded_noise + @staticmethod def flow_sde_sampling( - self, + self: RolloutSessionData, batch, model_output: torch.FloatTensor, sample: torch.FloatTensor, current_sigma: torch.FloatTensor, next_sigma: torch.FloatTensor, generator: torch.Generator, + step_index: int, ) -> torch.Tensor: """Flow rollout step for log-prob / sampling (see FlowGRPO-style references). @@ -255,7 +229,6 @@ def flow_sde_sampling( 2. ``"cps"``: Coupled Particle Sampling. 3. ``"ode"``: Deterministic ODE step (no diffusion noise). """ - rollout_session_data = self._get_rollout_session_data(batch) sde_type = batch.rollout_sde_type noise_level = float(batch.rollout_noise_level) log_prob_no_const = batch.rollout_log_prob_no_const @@ -268,16 +241,11 @@ def flow_sde_sampling( dt = next_sigma - current_sigma - # step_index comes from the denoising-loop counter stashed by - # DenoisingStage — scheduler._step_index would differ when - # _begin_index != 0 (e.g. partial denoising). sde_step_indices = getattr(batch, "rollout_sde_step_indices", None) - loop_step_index = getattr(batch, "_rollout_loop_step_index", None) if ( sde_type != "ode" and sde_step_indices is not None - and loop_step_index is not None - and loop_step_index not in sde_step_indices + and step_index not in sde_step_indices ): effective_sde_type = "ode" else: @@ -292,10 +260,10 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator + variance_noise = SchedulerRLMixin._rollout_variance_noise( + self, batch, model_output, generator ) - full_variance_noise = rollout_session_data.noise_buffer + full_variance_noise = self.noise_buffer std_dev_t = ( torch.sqrt( current_sigma @@ -303,7 +271,7 @@ def flow_sde_sampling( 1 - torch.where( torch.isclose(current_sigma, current_sigma.new_tensor(1.0)), - rollout_session_data.sigma_max, + self.sigma_max, current_sigma, ) ) @@ -325,10 +293,10 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator + variance_noise = SchedulerRLMixin._rollout_variance_noise( + self, batch, model_output, generator ) - full_variance_noise = rollout_session_data.noise_buffer + full_variance_noise = self.noise_buffer std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2) noise_std_dev = std_dev_t pred_original_sample = sample - current_sigma * model_output @@ -349,7 +317,7 @@ def flow_sde_sampling( (), device=model_output.device, dtype=model_output.dtype ) log_prob_no_const_val = torch.zeros( - rollout_session_data.latents_shape, + self.latents_shape, device=model_output.device, dtype=torch.float32, ) @@ -379,43 +347,48 @@ def flow_sde_sampling( ).sum(dim=list(range(1, len(log_prob_no_const_val.shape)))) if debug_mode: - self.append_local_rollout_debug_tensors( - batch, + SchedulerRLMixin.append_local_rollout_debug_tensors( + self, variance_noise=variance_noise, prev_sample_mean=prev_sample_mean, noise_std_dev=noise_std_dev, model_output=model_output, ) - self.append_local_rollout_log_probs(batch, log_prob_local_sum, local_elem_count) + SchedulerRLMixin.append_local_rollout_log_probs( + self, log_prob_local_sum, local_elem_count + ) return prev_sample + @staticmethod def append_local_rollout_log_probs( - self, batch, log_prob_sum: torch.Tensor, log_prob_count: torch.Tensor + self: RolloutSessionData, + log_prob_sum: torch.Tensor, + log_prob_count: torch.Tensor, ) -> None: - rollout_session_data = self._get_rollout_session_data(batch) - rollout_session_data.local_log_prob_sum.append(log_prob_sum) - rollout_session_data.local_log_prob_count.append(log_prob_count) + self.local_log_prob_sum.append(log_prob_sum) + self.local_log_prob_count.append(log_prob_count) + @staticmethod def consume_local_rollout_log_probs( - self, batch + self: RolloutSessionData, ) -> tuple[torch.Tensor, torch.Tensor]: - rollout_session_data = self._get_rollout_session_data(batch) # [B, T]: batch dim 0, denoising step dim 1 - values_sum = torch.stack(rollout_session_data.local_log_prob_sum, dim=1) - values_count = torch.stack(rollout_session_data.local_log_prob_count, dim=1) - rollout_session_data.local_log_prob_sum = [] - rollout_session_data.local_log_prob_count = [] + values_sum = torch.stack(self.local_log_prob_sum, dim=1) + values_count = torch.stack(self.local_log_prob_count, dim=1) + self.local_log_prob_sum = [] + self.local_log_prob_count = [] return values_sum, values_count - def collect_rollout_log_probs(self, batch: Req) -> torch.Tensor | None: + @staticmethod + def collect_rollout_log_probs(self: RolloutSessionData) -> torch.Tensor | None: """Per-step sums are already computed on the full pre-shard noise buffer inside flow_sde_sampling, so every SP rank holds identical values here and no all-reduce is needed.""" trajectory_log_prob_sum, trajectory_log_prob_count = ( - self.consume_local_rollout_log_probs(batch) + SchedulerRLMixin.consume_local_rollout_log_probs(self) ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index fd3f43d16eca..735ce733ee85 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -67,11 +67,15 @@ def _maybe_collect_rollout_log_probs(self, batch: Req): if batch.rollout_trajectory_data is None: batch.rollout_trajectory_data = RolloutTrajectoryData() batch.rollout_trajectory_data.rollout_log_probs = ( - scheduler.collect_rollout_log_probs(batch) + SchedulerRLMixin.collect_rollout_log_probs( + scheduler._get_rollout_session_data(batch) + ) ) if batch.rollout_debug_mode: batch.rollout_trajectory_data.rollout_debug_tensors = ( - scheduler.collect_rollout_debug_tensors(batch) + SchedulerRLMixin.collect_rollout_debug_tensors( + scheduler._get_rollout_session_data(batch), batch + ) ) scheduler.release_rollout_resources(batch) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 943ae2d81297..667b764ce335 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -38,26 +38,30 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): batch = self._build_batch(debug_mode=False) scheduler.prepare_rollout(batch) - def _raise_if_called(*args, **kwargs): - raise AssertionError("ODE path should not call _rollout_variance_noise") - - scheduler._rollout_variance_noise = _raise_if_called # type: ignore[method-assign] + generator = torch.Generator().manual_seed(1) + generator_state = generator.get_state() sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) model_output = torch.randn_like(sample) current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - prev_sample = scheduler.flow_sde_sampling( + prev_sample = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=torch.Generator(device=sample.device).manual_seed(1), + generator=generator, + step_index=2, ) + # The ODE path must not draw variance noise from the request generator. + self.assertTrue(torch.equal(generator.get_state(), generator_state)) log_prob_local_sum, local_elem_count = ( - scheduler.consume_local_rollout_log_probs(batch) + SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data + ) ) log_prob_local_sum = log_prob_local_sum.squeeze(-1) local_elem_count = local_elem_count.squeeze(-1) @@ -89,13 +93,15 @@ def test_ode_bit_exact_with_non_rollout_path(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) dt = next_sigma - current_sigma - rollout_prev = scheduler.flow_sde_sampling( + rollout_prev = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=torch.Generator(device=sample.device).manual_seed(0), + step_index=2, ) # Exact expression used by the non-rollout branch at # scheduling_flow_match_euler_discrete.py `prev_sample = sample + @@ -132,13 +138,15 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - scheduler.flow_sde_sampling( + SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=torch.Generator(device=sample.device).manual_seed(2), + step_index=2, ) ( @@ -146,7 +154,9 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): prev_sample_means, noise_std_devs, model_outputs, - ) = scheduler.consume_local_rollout_debug_tensors(batch) + ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( + batch._rollout_session_data + ) # [B, T, ...] with one step in this test. self.assertEqual(tuple(variance_noises.shape), (2, 1, 4, 8, 8)) @@ -253,7 +263,7 @@ def test_single_step_matches_flowgrpo_reference(self): shape = (1, 16, 1, 32, 32) atol = 1e-6 pipeline_config = types.SimpleNamespace( - shard_latents_for_sp=lambda _batch, latents: (latents, False) + shard_latents_for_sp=lambda batch, latents: (latents, False) ) for sde_type in ("sde", "cps"): @@ -265,32 +275,29 @@ def test_single_step_matches_flowgrpo_reference(self): g = torch.Generator(device="cpu").manual_seed(seed) model_output = torch.randn(shape, generator=g, dtype=torch.float32) sample = torch.randn(shape, generator=g, dtype=torch.float32) - variance_noise = torch.randn(shape, generator=g, dtype=torch.float32) - - def _mock_rollout_variance_noise(_batch, *_args, **_kwargs): - # flow_sde_sampling reads the full pre-shard noise from - # rollout_session_data.noise_buffer to compute log_prob, so - # the mock must populate it alongside returning the - # (single-GPU trivially-sharded) noise. - scheduler._get_rollout_session_data( # type: ignore[attr-defined] - _batch - ).noise_buffer = variance_noise - return variance_noise - - scheduler._rollout_variance_noise = ( # type: ignore[method-assign] - _mock_rollout_variance_noise + # The controller draws its variance noise from the request + # generator; a twin generator reproduces it for the reference. + noise_generator = torch.Generator(device="cpu").manual_seed(seed + 100) + variance_noise = torch.randn( + shape, + generator=torch.Generator(device="cpu").manual_seed(seed + 100), + dtype=torch.float32, ) - prev_sgl = scheduler.flow_sde_sampling( + prev_sgl = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=g, + generator=noise_generator, + step_index=2, ) - log_prob_sum, elem_count = scheduler.consume_local_rollout_log_probs( - batch + log_prob_sum, elem_count = ( + SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data + ) ) log_prob_sum = log_prob_sum.squeeze(-1) elem_count = elem_count.squeeze(-1) @@ -299,7 +306,9 @@ def _mock_rollout_variance_noise(_batch, *_args, **_kwargs): prev_sample_means, noise_std_devs, _model_outputs, - ) = scheduler.consume_local_rollout_debug_tensors(batch) + ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( + batch._rollout_session_data + ) prev_ref, log_prob_ref, prev_mean_ref, noise_std_ref = ( _flowgrpo_sde_step_with_logprob( @@ -369,21 +378,25 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): # propagates from the (original) model_output.dtype into the # noise buffer. If flow_sde_sampling fails to cast to fp32 at # entry, the buffer is bf16 → log_prob becomes bf16. - scheduler.flow_sde_sampling( + SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=g, + step_index=2, + ) + log_prob_sum, _count = SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data ) - log_prob_sum, _count = scheduler.consume_local_rollout_log_probs(batch) self.assertEqual( log_prob_sum.dtype, torch.float32, msg=f"{sde_type}: log_prob_sum must be fp32 with bf16 model_output", ) - noise_buffer = scheduler._get_rollout_session_data(batch).noise_buffer + noise_buffer = batch._rollout_session_data.noise_buffer self.assertEqual( noise_buffer.dtype, torch.float32, @@ -403,7 +416,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): scheduler = _DummyScheduler() shape = (1, 4, 8, 8) pipeline_config = types.SimpleNamespace( - shard_latents_for_sp=lambda _batch, latents: (latents, False) + shard_latents_for_sp=lambda batch, latents: (latents, False) ) batch = types.SimpleNamespace( rollout_log_prob_no_const=False, @@ -422,48 +435,40 @@ def test_timestep_filters_gate_sde_and_trajectory(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - variance_noise_ref = torch.randn(shape, generator=g, dtype=torch.float32) - variance_noise_call_count = {"n": 0} - - def _mock_variance_noise(_batch, *_args, **_kwargs): - variance_noise_call_count["n"] += 1 - scheduler._get_rollout_session_data( - _batch - ).noise_buffer = variance_noise_ref - return variance_noise_ref - - scheduler._rollout_variance_noise = ( # type: ignore[method-assign] - _mock_variance_noise - ) + generator_state = g.get_state() # Step 0: not in filter → deterministic ODE transition, no noise draw. - batch._rollout_loop_step_index = 0 - prev_0 = scheduler.flow_sde_sampling( + prev_0 = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=g, + step_index=0, ) - self.assertEqual(variance_noise_call_count["n"], 0) + self.assertTrue(torch.equal(g.get_state(), generator_state)) expected_ode = sample + (next_sigma - current_sigma) * model_output self.assertTrue(torch.allclose(prev_0, expected_ode, atol=1e-6)) # Step 1: in filter → real SDE, noise drawn, prev differs from ODE form. - batch._rollout_loop_step_index = 1 - prev_1 = scheduler.flow_sde_sampling( + prev_1 = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=g, + step_index=1, ) - self.assertEqual(variance_noise_call_count["n"], 1) + self.assertFalse(torch.equal(g.get_state(), generator_state)) self.assertFalse(torch.allclose(prev_1, expected_ode, atol=1e-3)) - log_prob_sum, elem_count = scheduler.consume_local_rollout_log_probs(batch) + log_prob_sum, elem_count = SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data + ) self.assertEqual(tuple(log_prob_sum.shape), (shape[0], 2)) # Filtered step contributes zero log-prob; real SDE step does not. self.assertTrue( From 5d612f24c89f3bb6e8d1909bfab7d2b1c13030fa Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 06/15] refactor(move): move the rollout session methods into RolloutSessionData Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 3 +- .../post_training/flow_rollout_controller.py | 119 ++++++++---------- .../post_training/rollout_denoising_mixin.py | 10 +- .../test/unit/test_scheduler_rollout_unit.py | 45 +++---- 4 files changed, 72 insertions(+), 105 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index dc7d6a69b3ef..f81411a77975 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -538,8 +538,7 @@ def step( if batch is not None and batch.rollout: if not self.already_prepared_rollout(batch): raise RuntimeError("Rollout not prepared before step") - prev_sample = SchedulerRLMixin.flow_sde_sampling( - self._get_rollout_session_data(batch), + prev_sample = self._get_rollout_session_data(batch).flow_sde_sampling( batch, model_output, sample, diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 856e80b91131..5735b874c67a 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -45,18 +45,14 @@ class RolloutSessionData: local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) local_model_outputs: list[torch.Tensor] = field(default_factory=list) - -class SchedulerRLMixin: - @staticmethod - def _reset_rollout_debug_tensors(self: RolloutSessionData) -> None: + def _reset_rollout_debug_tensors(self) -> None: self.local_variance_noises = [] self.local_prev_sample_means = [] self.local_noise_std_devs = [] self.local_model_outputs = [] - @staticmethod def append_local_rollout_debug_tensors( - self: RolloutSessionData, + self, *, variance_noise: torch.Tensor, prev_sample_mean: torch.Tensor, @@ -70,28 +66,24 @@ def append_local_rollout_debug_tensors( self.local_noise_std_devs.append(noise_std_dev.expand((batch_size, 1))) self.local_model_outputs.append(model_output) - @staticmethod def consume_local_rollout_debug_tensors( - self: RolloutSessionData, + self, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: variance_noises = torch.stack(self.local_variance_noises, dim=1) prev_sample_means = torch.stack(self.local_prev_sample_means, dim=1) noise_std_devs = torch.stack(self.local_noise_std_devs, dim=1) model_outputs = torch.stack(self.local_model_outputs, dim=1) - SchedulerRLMixin._reset_rollout_debug_tensors(self) + self._reset_rollout_debug_tensors() return variance_noises, prev_sample_means, noise_std_devs, model_outputs - @staticmethod - def collect_rollout_debug_tensors( - self: RolloutSessionData, batch: Req - ) -> RolloutDebugTensors: + def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: """ Consume rollout debug tensors and merge for all SP ranks. Returns rollout debug tensors with shape [B, T, ...]. """ variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( - SchedulerRLMixin.consume_local_rollout_debug_tensors(self) + self.consume_local_rollout_debug_tensors() ) pipeline_config = self.pipeline_config @@ -119,38 +111,8 @@ def collect_rollout_debug_tensors( rollout_model_outputs=model_outputs.cpu(), ) - @staticmethod - def _get_rollout_session_data(batch) -> RolloutSessionData: - """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" - rollout_session_data = getattr(batch, "_rollout_session_data", None) - if rollout_session_data is None: - raise RuntimeError("prepare_rollout() not called before rollout") - return rollout_session_data - - def release_rollout_resources(self, batch) -> None: - """Release rollout-owned resources. Call when denoising ends or before a new rollout.""" - batch._rollout_session_data = None - - def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: - """Enable rollout and set SDE/CPS params. Call once before the denoising loop.""" - if get_sp_world_size() > 1 and pipeline_config is None: - raise RuntimeError( - "SP rollout requires pipeline_config to be passed to prepare_rollout()." - ) - batch._rollout_session_data = RolloutSessionData( - pipeline_config=pipeline_config, - sigma_max=self.sigmas[min(1, len(self.sigmas) - 1)].item(), - latents_shape=( - tuple(batch.latents.shape) if batch.latents is not None else None - ), - ) - - def already_prepared_rollout(self, batch) -> bool: - return getattr(batch, "_rollout_session_data", None) is not None - - @staticmethod def _get_or_create_rollout_noise_buffer( - self: RolloutSessionData, + self, full_shape: tuple, device: torch.device, dtype: torch.dtype, @@ -167,9 +129,8 @@ def _get_or_create_rollout_noise_buffer( self.noise_buffer = buffer return buffer - @staticmethod def _rollout_variance_noise( - self: RolloutSessionData, + self, batch, model_output: torch.FloatTensor, generator: Union[torch.Generator, list[torch.Generator]], @@ -190,8 +151,8 @@ def _rollout_variance_noise( "Generator list must have the same length as batch size" ) - buffer = SchedulerRLMixin._get_or_create_rollout_noise_buffer( - self, self.latents_shape, device, dtype + buffer = self._get_or_create_rollout_noise_buffer( + self.latents_shape, device, dtype ) for i in range(B): torch.randn( @@ -210,9 +171,8 @@ def _rollout_variance_noise( ) return sharded_noise - @staticmethod def flow_sde_sampling( - self: RolloutSessionData, + self, batch, model_output: torch.FloatTensor, sample: torch.FloatTensor, @@ -260,8 +220,8 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = SchedulerRLMixin._rollout_variance_noise( - self, batch, model_output, generator + variance_noise = self._rollout_variance_noise( + batch, model_output, generator ) full_variance_noise = self.noise_buffer std_dev_t = ( @@ -293,8 +253,8 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = SchedulerRLMixin._rollout_variance_noise( - self, batch, model_output, generator + variance_noise = self._rollout_variance_noise( + batch, model_output, generator ) full_variance_noise = self.noise_buffer std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2) @@ -347,32 +307,27 @@ def flow_sde_sampling( ).sum(dim=list(range(1, len(log_prob_no_const_val.shape)))) if debug_mode: - SchedulerRLMixin.append_local_rollout_debug_tensors( - self, + self.append_local_rollout_debug_tensors( variance_noise=variance_noise, prev_sample_mean=prev_sample_mean, noise_std_dev=noise_std_dev, model_output=model_output, ) - SchedulerRLMixin.append_local_rollout_log_probs( - self, log_prob_local_sum, local_elem_count - ) + self.append_local_rollout_log_probs(log_prob_local_sum, local_elem_count) return prev_sample - @staticmethod def append_local_rollout_log_probs( - self: RolloutSessionData, + self, log_prob_sum: torch.Tensor, log_prob_count: torch.Tensor, ) -> None: self.local_log_prob_sum.append(log_prob_sum) self.local_log_prob_count.append(log_prob_count) - @staticmethod def consume_local_rollout_log_probs( - self: RolloutSessionData, + self, ) -> tuple[torch.Tensor, torch.Tensor]: # [B, T]: batch dim 0, denoising step dim 1 values_sum = torch.stack(self.local_log_prob_sum, dim=1) @@ -381,14 +336,44 @@ def consume_local_rollout_log_probs( self.local_log_prob_count = [] return values_sum, values_count - @staticmethod - def collect_rollout_log_probs(self: RolloutSessionData) -> torch.Tensor | None: + def collect_rollout_log_probs(self) -> torch.Tensor | None: """Per-step sums are already computed on the full pre-shard noise buffer inside flow_sde_sampling, so every SP rank holds identical values here and no all-reduce is needed.""" trajectory_log_prob_sum, trajectory_log_prob_count = ( - SchedulerRLMixin.consume_local_rollout_log_probs(self) + self.consume_local_rollout_log_probs() ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() + + +class SchedulerRLMixin: + @staticmethod + def _get_rollout_session_data(batch) -> RolloutSessionData: + """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" + rollout_session_data = getattr(batch, "_rollout_session_data", None) + if rollout_session_data is None: + raise RuntimeError("prepare_rollout() not called before rollout") + return rollout_session_data + + def release_rollout_resources(self, batch) -> None: + """Release rollout-owned resources. Call when denoising ends or before a new rollout.""" + batch._rollout_session_data = None + + def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: + """Enable rollout and set SDE/CPS params. Call once before the denoising loop.""" + if get_sp_world_size() > 1 and pipeline_config is None: + raise RuntimeError( + "SP rollout requires pipeline_config to be passed to prepare_rollout()." + ) + batch._rollout_session_data = RolloutSessionData( + pipeline_config=pipeline_config, + sigma_max=self.sigmas[min(1, len(self.sigmas) - 1)].item(), + latents_shape=( + tuple(batch.latents.shape) if batch.latents is not None else None + ), + ) + + def already_prepared_rollout(self, batch) -> bool: + return getattr(batch, "_rollout_session_data", None) is not None diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index 735ce733ee85..cefacf88033e 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -67,15 +67,13 @@ def _maybe_collect_rollout_log_probs(self, batch: Req): if batch.rollout_trajectory_data is None: batch.rollout_trajectory_data = RolloutTrajectoryData() batch.rollout_trajectory_data.rollout_log_probs = ( - SchedulerRLMixin.collect_rollout_log_probs( - scheduler._get_rollout_session_data(batch) - ) + scheduler._get_rollout_session_data(batch).collect_rollout_log_probs() ) if batch.rollout_debug_mode: batch.rollout_trajectory_data.rollout_debug_tensors = ( - SchedulerRLMixin.collect_rollout_debug_tensors( - scheduler._get_rollout_session_data(batch), batch - ) + scheduler._get_rollout_session_data( + batch + ).collect_rollout_debug_tensors(batch) ) scheduler.release_rollout_resources(batch) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 667b764ce335..5de2b942a0bc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -46,8 +46,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - prev_sample = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_sample = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -59,9 +58,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): # The ODE path must not draw variance noise from the request generator. self.assertTrue(torch.equal(generator.get_state(), generator_state)) log_prob_local_sum, local_elem_count = ( - SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data - ) + batch._rollout_session_data.consume_local_rollout_log_probs() ) log_prob_local_sum = log_prob_local_sum.squeeze(-1) local_elem_count = local_elem_count.squeeze(-1) @@ -93,8 +90,7 @@ def test_ode_bit_exact_with_non_rollout_path(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) dt = next_sigma - current_sigma - rollout_prev = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + rollout_prev = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -138,8 +134,7 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -154,9 +149,7 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): prev_sample_means, noise_std_devs, model_outputs, - ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( - batch._rollout_session_data - ) + ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() # [B, T, ...] with one step in this test. self.assertEqual(tuple(variance_noises.shape), (2, 1, 4, 8, 8)) @@ -284,8 +277,7 @@ def test_single_step_matches_flowgrpo_reference(self): dtype=torch.float32, ) - prev_sgl = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_sgl = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -295,9 +287,7 @@ def test_single_step_matches_flowgrpo_reference(self): step_index=2, ) log_prob_sum, elem_count = ( - SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data - ) + batch._rollout_session_data.consume_local_rollout_log_probs() ) log_prob_sum = log_prob_sum.squeeze(-1) elem_count = elem_count.squeeze(-1) @@ -306,9 +296,7 @@ def test_single_step_matches_flowgrpo_reference(self): prev_sample_means, noise_std_devs, _model_outputs, - ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( - batch._rollout_session_data - ) + ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() prev_ref, log_prob_ref, prev_mean_ref, noise_std_ref = ( _flowgrpo_sde_step_with_logprob( @@ -378,8 +366,7 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): # propagates from the (original) model_output.dtype into the # noise buffer. If flow_sde_sampling fails to cast to fp32 at # entry, the buffer is bf16 → log_prob becomes bf16. - SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -388,8 +375,8 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): generator=g, step_index=2, ) - log_prob_sum, _count = SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data + log_prob_sum, _count = ( + batch._rollout_session_data.consume_local_rollout_log_probs() ) self.assertEqual( log_prob_sum.dtype, @@ -438,8 +425,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): generator_state = g.get_state() # Step 0: not in filter → deterministic ODE transition, no noise draw. - prev_0 = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_0 = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -453,8 +439,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertTrue(torch.allclose(prev_0, expected_ode, atol=1e-6)) # Step 1: in filter → real SDE, noise drawn, prev differs from ODE form. - prev_1 = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_1 = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -466,8 +451,8 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertFalse(torch.equal(g.get_state(), generator_state)) self.assertFalse(torch.allclose(prev_1, expected_ode, atol=1e-3)) - log_prob_sum, elem_count = SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data + log_prob_sum, elem_count = ( + batch._rollout_session_data.consume_local_rollout_log_probs() ) self.assertEqual(tuple(log_prob_sum.shape), (shape[0], 2)) # Filtered step contributes zero log-prob; real SDE step does not. From c17608464bf5786ce9c85276630234cab83781ed Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 07/15] refactor(diffusion): turn RolloutSessionData into a per-stream FlowRolloutController and drop SchedulerRLMixin Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 10 +- .../runtime/pipelines_core/schedule_batch.py | 5 + .../post_training/flow_rollout_controller.py | 174 +++++++----------- .../post_training/rollout_denoising_mixin.py | 76 ++++---- .../test/unit/test_scheduler_rollout_unit.py | 109 +++++------ 5 files changed, 163 insertions(+), 211 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index f81411a77975..b6c5c19cb8d4 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -35,9 +35,6 @@ from sglang.multimodal_gen.runtime.models.schedulers.sigma_schedule import ( SigmaScheduleMixin, ) -from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - SchedulerRLMixin, -) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -58,7 +55,7 @@ class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): class FlowMatchEulerDiscreteScheduler( - SchedulerMixin, ConfigMixin, BaseScheduler, SchedulerRLMixin, SigmaScheduleMixin + SchedulerMixin, ConfigMixin, BaseScheduler, SigmaScheduleMixin ): """ Euler scheduler. @@ -536,10 +533,9 @@ def step( dt = sigma_next - sigma if batch is not None and batch.rollout: - if not self.already_prepared_rollout(batch): + if "visual" not in batch.rollout_controllers: raise RuntimeError("Rollout not prepared before step") - prev_sample = self._get_rollout_session_data(batch).flow_sde_sampling( - batch, + prev_sample = batch.rollout_controllers["visual"].flow_sde_sampling( model_output, sample, current_sigma, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index de10a1666e60..c1e6d15b7c9d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -27,6 +27,9 @@ DataType, SamplingParams, ) +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + FlowRolloutController, +) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutTrajectoryData, ) @@ -198,6 +201,8 @@ class Req: trajectory_timesteps: torch.Tensor | None = None trajectory_latents: torch.Tensor | None = None rollout_trajectory_data: RolloutTrajectoryData | None = None + # Rollout controllers of this request, keyed by latent stream. + rollout_controllers: dict[str, FlowRolloutController] = field(default_factory=dict) trajectory_audio_latents: torch.Tensor | None = None # Extra parameters that might be needed by specific pipeline implementations (e.g., LTX2.3 DenoisingAVStage) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 5735b874c67a..8800f18167b7 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -1,49 +1,63 @@ # SPDX-License-Identifier: Apache-2.0 -"""Flow-matching rollout step utilities for log-prob computation.""" +"""Per-stream flow-matching rollout controller: the rollout transition and its log-probs.""" from __future__ import annotations import math -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Union +from typing import Callable, Union import torch -from sglang.multimodal_gen.runtime.distributed import ( - get_sp_world_size, -) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDebugTensors, ) -from sglang.multimodal_gen.runtime.post_training.sp_utils import ( - gather_stacked_latents_for_sp, -) - -if TYPE_CHECKING: - from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) -@dataclass -class RolloutSessionData: - """Per-batch rollout state created by prepare_rollout(), lives on the batch object. - - Cleared by setting ``batch._rollout_session_data = None``. +class FlowRolloutController: + """Rollout controller of one latent stream of one request, shared by every + flow-matching model: its SDE/CPS/ODE transition and the per-step log-probs. + The denoising stage creates it in ``batch.rollout_controllers`` and clears + that dict when the loop ends. """ - pipeline_config: Any = None - sigma_max: float = 0.0 - latents_shape: tuple | None = None - noise_buffer: torch.Tensor | None = None - - local_log_prob_sum: list[torch.Tensor] = field(default_factory=list) - local_log_prob_count: list[torch.Tensor] = field(default_factory=list) - - local_variance_noises: list[torch.Tensor] = field(default_factory=list) - local_prev_sample_means: list[torch.Tensor] = field(default_factory=list) - local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) - local_model_outputs: list[torch.Tensor] = field(default_factory=list) + def __init__( + self, + *, + sde_type: str, + noise_level: float, + log_prob_no_const: bool, + debug_mode: bool, + sde_step_indices: list[int] | None, + shard_for_sp: Callable[[torch.Tensor], torch.Tensor], + gather_for_sp: Callable[[torch.Tensor], torch.Tensor], + sigmas: torch.Tensor, + latents_shape: tuple | None, + ) -> None: + self.sde_type = sde_type + self.noise_level = noise_level + self.log_prob_no_const = log_prob_no_const + self.debug_mode = debug_mode + self.sde_step_indices = sde_step_indices + # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. + self.shard_for_sp = shard_for_sp + self.gather_for_sp = gather_for_sp + # [T+1] sigma grid of the loop, terminal 0 included. + self.sigmas = sigmas + # FlowGRPO's sigma_max: sigmas[1], used in place of sigma == 1 so the SDE std + # never divides by 1 - sigma = 0. + self.sigma_max = sigmas[min(1, len(sigmas) - 1)].item() + self.latents_shape = latents_shape + self.noise_buffer: torch.Tensor | None = None + + self.local_log_prob_sum: list[torch.Tensor] = [] + self.local_log_prob_count: list[torch.Tensor] = [] + + self.local_variance_noises: list[torch.Tensor] = [] + self.local_prev_sample_means: list[torch.Tensor] = [] + self.local_noise_std_devs: list[torch.Tensor] = [] + self.local_model_outputs: list[torch.Tensor] = [] def _reset_rollout_debug_tensors(self) -> None: self.local_variance_noises = [] @@ -76,7 +90,7 @@ def consume_local_rollout_debug_tensors( self._reset_rollout_debug_tensors() return variance_noises, prev_sample_means, noise_std_devs, model_outputs - def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: + def collect_rollout_debug_tensors(self) -> RolloutDebugTensors: """ Consume rollout debug tensors and merge for all SP ranks. @@ -86,22 +100,9 @@ def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: self.consume_local_rollout_debug_tensors() ) - pipeline_config = self.pipeline_config - variance_noises = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=variance_noises, - ) - prev_sample_means = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=prev_sample_means, - ) - model_outputs = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=model_outputs, - ) + variance_noises = self.gather_for_sp(variance_noises) + prev_sample_means = self.gather_for_sp(prev_sample_means) + model_outputs = self.gather_for_sp(model_outputs) # noise_std_devs is same on every device, not a sharded latent tensor. return RolloutDebugTensors( @@ -131,7 +132,6 @@ def _get_or_create_rollout_noise_buffer( def _rollout_variance_noise( self, - batch, model_output: torch.FloatTensor, generator: Union[torch.Generator, list[torch.Generator]], ) -> torch.FloatTensor: @@ -161,9 +161,7 @@ def _rollout_variance_noise( generator=generator[i], ) - sharded_noise, _ = self.pipeline_config.shard_latents_for_sp( - batch=batch, latents=buffer - ) + sharded_noise = self.shard_for_sp(buffer) if tuple(sharded_noise.shape) != local_shape: raise ValueError( "Rollout SP noise shape mismatch after shard. " @@ -173,7 +171,6 @@ def _rollout_variance_noise( def flow_sde_sampling( self, - batch, model_output: torch.FloatTensor, sample: torch.FloatTensor, current_sigma: torch.FloatTensor, @@ -183,33 +180,27 @@ def flow_sde_sampling( ) -> torch.Tensor: """Flow rollout step for log-prob / sampling (see FlowGRPO-style references). - ``rollout_sde_type`` (from batch SamplingParams): + ``sde_type`` (from the request's ``rollout_sde_type``): 1. ``"sde"``: Standard stochastic differential equation transition (Gaussian). 2. ``"cps"``: Coupled Particle Sampling. 3. ``"ode"``: Deterministic ODE step (no diffusion noise). """ - sde_type = batch.rollout_sde_type - noise_level = float(batch.rollout_noise_level) - log_prob_no_const = batch.rollout_log_prob_no_const - debug_mode = bool(getattr(batch, "rollout_debug_mode", False)) - - if not log_prob_no_const and sde_type != "ode": - assert noise_level > 0, ( + if not self.log_prob_no_const and self.sde_type != "ode": + assert self.noise_level > 0, ( "True log-probability computation requires a non-zero noise level." ) dt = next_sigma - current_sigma - sde_step_indices = getattr(batch, "rollout_sde_step_indices", None) if ( - sde_type != "ode" - and sde_step_indices is not None - and step_index not in sde_step_indices + self.sde_type != "ode" + and self.sde_step_indices is not None + and step_index not in self.sde_step_indices ): effective_sde_type = "ode" else: - effective_sde_type = sde_type + effective_sde_type = self.sde_type # sde/cps: cast to fp32 to match flowGRPO semantics and avoid the # 0-dim-fp32 wrapped-scalar promotion demoting log-prob to bf16. @@ -220,9 +211,7 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator - ) + variance_noise = self._rollout_variance_noise(model_output, generator) full_variance_noise = self.noise_buffer std_dev_t = ( torch.sqrt( @@ -236,7 +225,7 @@ def flow_sde_sampling( ) ) ) - * noise_level + * self.noise_level ) noise_std_dev = std_dev_t * torch.sqrt(-1 * dt) prev_sample_mean = ( @@ -253,11 +242,9 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator - ) + variance_noise = self._rollout_variance_noise(model_output, generator) full_variance_noise = self.noise_buffer - std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2) + std_dev_t = next_sigma * math.sin(self.noise_level * math.pi / 2) noise_std_dev = std_dev_t pred_original_sample = sample - current_sigma * model_output noise_estimate = sample + model_output * (1 - current_sigma) @@ -283,13 +270,13 @@ def flow_sde_sampling( ) # Only enforce the "no full log-prob with ODE" constraint when the # user explicitly chose ODE globally. - if sde_type == "ode": - assert log_prob_no_const, ( + if self.sde_type == "ode": + assert self.log_prob_no_const, ( "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation" ) else: - raise ValueError(f"Unsupported sde_type: {sde_type}") + raise ValueError(f"Unsupported sde_type: {self.sde_type}") reduce_dims = list(range(1, len(log_prob_no_const_val.shape))) local_elem_count = log_prob_no_const_val.new_full( @@ -297,7 +284,7 @@ def flow_sde_sampling( float(math.prod(log_prob_no_const_val.shape[1:])), ) - if log_prob_no_const or effective_sde_type == "ode": + if self.log_prob_no_const or effective_sde_type == "ode": log_prob_local_sum = log_prob_no_const_val.sum(dim=reduce_dims) else: log_prob_local_sum = ( @@ -306,7 +293,7 @@ def flow_sde_sampling( - _LOG_SQRT_2PI ).sum(dim=list(range(1, len(log_prob_no_const_val.shape)))) - if debug_mode: + if self.debug_mode: self.append_local_rollout_debug_tensors( variance_noise=variance_noise, prev_sample_mean=prev_sample_mean, @@ -336,7 +323,7 @@ def consume_local_rollout_log_probs( self.local_log_prob_count = [] return values_sum, values_count - def collect_rollout_log_probs(self) -> torch.Tensor | None: + def collect_rollout_log_probs(self) -> torch.Tensor: """Per-step sums are already computed on the full pre-shard noise buffer inside flow_sde_sampling, so every SP rank holds identical values here and no all-reduce is needed.""" @@ -346,34 +333,3 @@ def collect_rollout_log_probs(self) -> torch.Tensor | None: ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() - - -class SchedulerRLMixin: - @staticmethod - def _get_rollout_session_data(batch) -> RolloutSessionData: - """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" - rollout_session_data = getattr(batch, "_rollout_session_data", None) - if rollout_session_data is None: - raise RuntimeError("prepare_rollout() not called before rollout") - return rollout_session_data - - def release_rollout_resources(self, batch) -> None: - """Release rollout-owned resources. Call when denoising ends or before a new rollout.""" - batch._rollout_session_data = None - - def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: - """Enable rollout and set SDE/CPS params. Call once before the denoising loop.""" - if get_sp_world_size() > 1 and pipeline_config is None: - raise RuntimeError( - "SP rollout requires pipeline_config to be passed to prepare_rollout()." - ) - batch._rollout_session_data = RolloutSessionData( - pipeline_config=pipeline_config, - sigma_max=self.sigmas[min(1, len(self.sigmas) - 1)].item(), - latents_shape=( - tuple(batch.latents.shape) if batch.latents is not None else None - ), - ) - - def already_prepared_rollout(self, batch) -> bool: - return getattr(batch, "_rollout_session_data", None) is not None diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index cefacf88033e..90fc63b2ad4b 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -9,9 +9,12 @@ import torch +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - SchedulerRLMixin, + FlowRolloutController, ) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDenoisingEnv, @@ -38,44 +41,49 @@ def _kwargs_to_cpu(d: Any) -> Any: class RolloutDenoisingMixin: def _maybe_prepare_rollout(self, batch: Req): - """Prepare denoising loop for rollout.""" - scheduler = batch.scheduler - if not isinstance(scheduler, SchedulerRLMixin): - if batch.rollout: - raise ValueError( - f"Scheduler {type(scheduler)} does not support rollout" - ) - return + """Create the rollout controller of the visual stream of this request. - scheduler.release_rollout_resources(batch) - if batch.rollout: - scheduler.prepare_rollout( - batch=batch, - pipeline_config=self.server_args.pipeline_config, + Must run before SP sharding so the controller records the full latent shape. + """ + scheduler = batch.scheduler + if not isinstance(scheduler, FlowMatchEulerDiscreteScheduler): + raise ValueError(f"Scheduler {type(scheduler)} does not support rollout") + + pipeline_config = self.server_args.pipeline_config + batch.rollout_controllers = { + "visual": FlowRolloutController( + sde_type=batch.rollout_sde_type, + noise_level=float(batch.rollout_noise_level), + log_prob_no_const=batch.rollout_log_prob_no_const, + debug_mode=batch.rollout_debug_mode, + sde_step_indices=batch.rollout_sde_step_indices, + shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + batch=batch, latents=latents + )[0], + gather_for_sp=lambda stacked_latents: gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=stacked_latents, + ), + sigmas=scheduler.sigmas, + latents_shape=( + tuple(batch.latents.shape) if batch.latents is not None else None + ), ) + } def _maybe_collect_rollout_log_probs(self, batch: Req): - scheduler = batch.scheduler - if not isinstance(scheduler, SchedulerRLMixin): - if batch.rollout: - raise ValueError( - f"Scheduler {type(scheduler)} does not support rollout" - ) - return - - if batch.rollout: - if batch.rollout_trajectory_data is None: - batch.rollout_trajectory_data = RolloutTrajectoryData() - batch.rollout_trajectory_data.rollout_log_probs = ( - scheduler._get_rollout_session_data(batch).collect_rollout_log_probs() + rollout = batch.rollout_controllers["visual"] + if batch.rollout_trajectory_data is None: + batch.rollout_trajectory_data = RolloutTrajectoryData() + batch.rollout_trajectory_data.rollout_log_probs = ( + rollout.collect_rollout_log_probs() + ) + if batch.rollout_debug_mode: + batch.rollout_trajectory_data.rollout_debug_tensors = ( + rollout.collect_rollout_debug_tensors() ) - if batch.rollout_debug_mode: - batch.rollout_trajectory_data.rollout_debug_tensors = ( - scheduler._get_rollout_session_data( - batch - ).collect_rollout_debug_tensors(batch) - ) - scheduler.release_rollout_resources(batch) + batch.rollout_controllers = {} def _postprocess_rollout_outputs( self, diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 5de2b942a0bc..046f44bc27a5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -4,39 +4,44 @@ import torch -import sglang.multimodal_gen.runtime.post_training.flow_rollout_controller as rl_mixin_module from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - SchedulerRLMixin, + FlowRolloutController, ) +_SIGMAS = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0], dtype=torch.float32) -class _DummyScheduler(SchedulerRLMixin): - def __init__(self): - self.sigmas = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0], dtype=torch.float32) +def _make_rollout_controller(batch, pipeline_config=None) -> FlowRolloutController: + """The controller the denoising stage creates for ``batch`` (single GPU).""" + return FlowRolloutController( + sde_type=batch.rollout_sde_type, + noise_level=float(batch.rollout_noise_level), + log_prob_no_const=batch.rollout_log_prob_no_const, + debug_mode=batch.rollout_debug_mode, + sde_step_indices=batch.rollout_sde_step_indices, + shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + batch=batch, latents=latents + )[0], + gather_for_sp=lambda stacked_latents: stacked_latents, + sigmas=_SIGMAS, + latents_shape=tuple(batch.latents.shape), + ) -class TestSchedulerRolloutOdeUnit(unittest.TestCase): - def setUp(self): - self._orig_get_sp_world_size = rl_mixin_module.get_sp_world_size - rl_mixin_module.get_sp_world_size = lambda: 1 - - def tearDown(self): - rl_mixin_module.get_sp_world_size = self._orig_get_sp_world_size +class TestSchedulerRolloutOdeUnit(unittest.TestCase): def _build_batch(self, *, debug_mode: bool) -> types.SimpleNamespace: return types.SimpleNamespace( rollout_log_prob_no_const=True, rollout_noise_level=0.5, rollout_sde_type="ode", rollout_debug_mode=debug_mode, + rollout_sde_step_indices=None, latents=torch.zeros(2, 4, 8, 8, dtype=torch.float32), - _rollout_session_data=None, ) def test_ode_step_does_not_call_variance_noise_sampler(self): - scheduler = _DummyScheduler() batch = self._build_batch(debug_mode=False) - scheduler.prepare_rollout(batch) + current_rollout_controller = _make_rollout_controller(batch) generator = torch.Generator().manual_seed(1) generator_state = generator.get_state() @@ -46,8 +51,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - prev_sample = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_sample = current_rollout_controller.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -58,7 +62,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): # The ODE path must not draw variance noise from the request generator. self.assertTrue(torch.equal(generator.get_state(), generator_state)) log_prob_local_sum, local_elem_count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() + current_rollout_controller.consume_local_rollout_log_probs() ) log_prob_local_sum = log_prob_local_sum.squeeze(-1) local_elem_count = local_elem_count.squeeze(-1) @@ -80,9 +84,8 @@ def test_ode_bit_exact_with_non_rollout_path(self): wrapped-scalar promotion difference that a spurious ``model_output.float()`` in the ODE branch would introduce is most visible at bf16 precision.""" - scheduler = _DummyScheduler() batch = self._build_batch(debug_mode=False) - scheduler.prepare_rollout(batch) + current_rollout_controller = _make_rollout_controller(batch) sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) model_output = torch.randn_like(sample).to(torch.bfloat16) @@ -90,8 +93,7 @@ def test_ode_bit_exact_with_non_rollout_path(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) dt = next_sigma - current_sigma - rollout_prev = batch._rollout_session_data.flow_sde_sampling( - batch, + rollout_prev = current_rollout_controller.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -125,17 +127,15 @@ def test_ode_bit_exact_with_non_rollout_path(self): self.assertEqual(post_cast_max_abs_diff, 0.0) def test_ode_debug_tensors_have_shape_safe_noise_std(self): - scheduler = _DummyScheduler() batch = self._build_batch(debug_mode=True) - scheduler.prepare_rollout(batch) + current_rollout_controller = _make_rollout_controller(batch) sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) model_output = torch.randn_like(sample) current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - batch._rollout_session_data.flow_sde_sampling( - batch, + current_rollout_controller.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -149,7 +149,7 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): prev_sample_means, noise_std_devs, model_outputs, - ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() + ) = current_rollout_controller.consume_local_rollout_debug_tensors() # [B, T, ...] with one step in this test. self.assertEqual(tuple(variance_noises.shape), (2, 1, 4, 8, 8)) @@ -228,13 +228,6 @@ def _flowgrpo_sde_step_with_logprob( class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase): - def setUp(self): - self._orig_get_sp_world_size = rl_mixin_module.get_sp_world_size - rl_mixin_module.get_sp_world_size = lambda: 1 - - def tearDown(self): - rl_mixin_module.get_sp_world_size = self._orig_get_sp_world_size - def _build_batch( self, *, sde_type: str, shape: tuple[int, ...] ) -> types.SimpleNamespace: @@ -243,14 +236,13 @@ def _build_batch( rollout_noise_level=0.5, rollout_sde_type=sde_type, rollout_debug_mode=True, + rollout_sde_step_indices=None, latents=torch.empty(shape, dtype=torch.float32), - _rollout_session_data=None, ) def test_single_step_matches_flowgrpo_reference(self): """Verify prev_sample, prev_sample_mean, noise_std_dev, and log_prob all match FlowGRPO's ``sde_step_with_logprob`` for SDE and CPS.""" - scheduler = _DummyScheduler() current_sigma = torch.tensor(0.5, dtype=torch.float32) next_sigma = torch.tensor(0.3, dtype=torch.float32) shape = (1, 16, 1, 32, 32) @@ -262,8 +254,9 @@ def test_single_step_matches_flowgrpo_reference(self): for sde_type in ("sde", "cps"): for seed in (0, 1, 2, 3): batch = self._build_batch(sde_type=sde_type, shape=shape) - scheduler.release_rollout_resources(batch) - scheduler.prepare_rollout(batch=batch, pipeline_config=pipeline_config) + current_rollout_controller = _make_rollout_controller( + batch, pipeline_config + ) g = torch.Generator(device="cpu").manual_seed(seed) model_output = torch.randn(shape, generator=g, dtype=torch.float32) @@ -277,8 +270,7 @@ def test_single_step_matches_flowgrpo_reference(self): dtype=torch.float32, ) - prev_sgl = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_sgl = current_rollout_controller.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -287,7 +279,7 @@ def test_single_step_matches_flowgrpo_reference(self): step_index=2, ) log_prob_sum, elem_count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() + current_rollout_controller.consume_local_rollout_log_probs() ) log_prob_sum = log_prob_sum.squeeze(-1) elem_count = elem_count.squeeze(-1) @@ -296,7 +288,7 @@ def test_single_step_matches_flowgrpo_reference(self): prev_sample_means, noise_std_devs, _model_outputs, - ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() + ) = current_rollout_controller.consume_local_rollout_debug_tensors() prev_ref, log_prob_ref, prev_mean_ref, noise_std_ref = ( _flowgrpo_sde_step_with_logprob( @@ -305,7 +297,7 @@ def test_single_step_matches_flowgrpo_reference(self): variance_noise=variance_noise, sigma=current_sigma, sigma_prev=next_sigma, - sigma_max=scheduler.sigmas[1].item(), + sigma_max=_SIGMAS[1].item(), noise_level=0.5, sde_type=sde_type, ) @@ -343,7 +335,6 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): branches therefore cast ``model_output.float()`` at entry. Passing bf16 ``model_output`` must still yield an fp32 noise buffer and an fp32 log-prob sum.""" - scheduler = _DummyScheduler() current_sigma = torch.tensor(0.5, dtype=torch.float32) next_sigma = torch.tensor(0.3, dtype=torch.float32) shape = (1, 16, 1, 32, 32) @@ -353,8 +344,9 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): for sde_type in ("sde", "cps"): batch = self._build_batch(sde_type=sde_type, shape=shape) - scheduler.release_rollout_resources(batch) - scheduler.prepare_rollout(batch=batch, pipeline_config=pipeline_config) + current_rollout_controller = _make_rollout_controller( + batch, pipeline_config + ) g = torch.Generator(device="cpu").manual_seed(0) model_output = torch.randn(shape, generator=g, dtype=torch.float32).to( @@ -366,8 +358,7 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): # propagates from the (original) model_output.dtype into the # noise buffer. If flow_sde_sampling fails to cast to fp32 at # entry, the buffer is bf16 → log_prob becomes bf16. - batch._rollout_session_data.flow_sde_sampling( - batch, + current_rollout_controller.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -376,14 +367,14 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): step_index=2, ) log_prob_sum, _count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() + current_rollout_controller.consume_local_rollout_log_probs() ) self.assertEqual( log_prob_sum.dtype, torch.float32, msg=f"{sde_type}: log_prob_sum must be fp32 with bf16 model_output", ) - noise_buffer = batch._rollout_session_data.noise_buffer + noise_buffer = current_rollout_controller.noise_buffer self.assertEqual( noise_buffer.dtype, torch.float32, @@ -400,7 +391,6 @@ def test_timestep_filters_gate_sde_and_trajectory(self): ) # --- Part 1: rollout_sde_step_indices gates SDE noise injection --- - scheduler = _DummyScheduler() shape = (1, 4, 8, 8) pipeline_config = types.SimpleNamespace( shard_latents_for_sp=lambda batch, latents: (latents, False) @@ -412,9 +402,8 @@ def test_timestep_filters_gate_sde_and_trajectory(self): rollout_debug_mode=False, rollout_sde_step_indices=[1], # only step 1 is stochastic latents=torch.empty(shape, dtype=torch.float32), - _rollout_session_data=None, ) - scheduler.prepare_rollout(batch=batch, pipeline_config=pipeline_config) + current_rollout_controller = _make_rollout_controller(batch, pipeline_config) g = torch.Generator(device="cpu").manual_seed(0) sample = torch.randn(shape, generator=g, dtype=torch.float32) @@ -425,8 +414,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): generator_state = g.get_state() # Step 0: not in filter → deterministic ODE transition, no noise draw. - prev_0 = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_0 = current_rollout_controller.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -439,8 +427,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertTrue(torch.allclose(prev_0, expected_ode, atol=1e-6)) # Step 1: in filter → real SDE, noise drawn, prev differs from ODE form. - prev_1 = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_1 = current_rollout_controller.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -452,7 +439,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertFalse(torch.allclose(prev_1, expected_ode, atol=1e-3)) log_prob_sum, elem_count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() + current_rollout_controller.consume_local_rollout_log_probs() ) self.assertEqual(tuple(log_prob_sum.shape), (shape[0], 2)) # Filtered step contributes zero log-prob; real SDE step does not. @@ -492,7 +479,7 @@ class _DummyDit(RolloutDenoisingMixin): len(traj_filtered._rollout_denoising_env_state["step_latents"]), 2 ) self.assertEqual( - len(traj_filtered._rollout_denoising_env_state["step_timesteps"]), 2 + len(traj_filtered._rollout_denoising_env_state["step_timesteps"]), 3 ) # None (default) → all steps appended (back-compat). @@ -537,7 +524,7 @@ class _DummyDit(RolloutDenoisingMixin): len(traj_exclude_final._rollout_denoising_env_state["step_latents"]), 3 ) self.assertEqual( - len(traj_exclude_final._rollout_denoising_env_state["step_timesteps"]), 3 + len(traj_exclude_final._rollout_denoising_env_state["step_timesteps"]), 4 ) # Filter includes only step_index=T → only the final latent survives. @@ -564,7 +551,7 @@ class _DummyDit(RolloutDenoisingMixin): len(traj_only_final._rollout_denoising_env_state["step_latents"]), 1 ) self.assertEqual( - len(traj_only_final._rollout_denoising_env_state["step_timesteps"]), 1 + len(traj_only_final._rollout_denoising_env_state["step_timesteps"]), 4 ) From 0db14f717173f54646d03601c9a455e042d6da7b Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 08/15] refactor(diffusion): pass the rollout transition to scheduler step instead of branching on batch Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 44 +++++++++---------- .../pipelines_core/stages/denoising.py | 7 ++- .../stages/model_specific_stages/cosmos3.py | 26 +++++------ .../model_specific_stages/ltx_2/denoising.py | 29 +++++------- .../post_training/flow_rollout_controller.py | 33 +++++++++++++- .../post_training/rollout_denoising_mixin.py | 10 ++++- .../test/unit/test_scheduler_rollout_unit.py | 38 ++++++++++++++++ 7 files changed, 125 insertions(+), 62 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index b6c5c19cb8d4..301fcb43f193 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -22,7 +22,7 @@ # ============================================================================== import math from dataclasses import dataclass -from typing import Any +from typing import Any, Callable import numpy as np import torch @@ -462,7 +462,7 @@ def step( s_noise: float = 1.0, generator: torch.Generator | list[torch.Generator] | None = None, per_token_timesteps: torch.Tensor | None = None, - batch=None, + transition: Callable[..., torch.FloatTensor] | None = None, return_dict: bool = True, ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]: """ @@ -485,6 +485,10 @@ def step( A random number generator, or one generator per batch item. per_token_timesteps (`torch.Tensor`, *optional*): The timesteps for each token in the sample. + transition (`Callable`, *optional*): + Replaces this scheduler's own `_transition` for this call, with the + same keyword arguments; rollout passes + `FlowRolloutController.transition` here. return_dict (`bool`): Whether or not to return a [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. @@ -532,28 +536,18 @@ def step( next_sigma = sigma_next dt = sigma_next - sigma - if batch is not None and batch.rollout: - if "visual" not in batch.rollout_controllers: - raise RuntimeError("Rollout not prepared before step") - prev_sample = batch.rollout_controllers["visual"].flow_sde_sampling( - model_output, - sample, - current_sigma, - next_sigma, - generator, - step_index=self.step_index, - ) - else: - prev_sample = self._transition( - model_output=model_output, - sample=sample, - current_sigma=current_sigma, - next_sigma=next_sigma, - dt=dt, - generator=generator, - step_index=self.step_index, - timestep=timestep, - ) + if transition is None: + transition = self._transition + prev_sample = transition( + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + dt=dt, + generator=generator, + step_index=self.step_index, + timestep=timestep, + ) # upon completion increase step index by one assert self._step_index is not None, "_step_index should not be None" @@ -567,6 +561,8 @@ def step( return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) + # Old ODE transition: the scheduler's own Euler update (or its stochastic_sampling step), + # used when step() is not given a transition. def _transition( self, *, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 4ce49db237b8..4e3541832c73 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -156,6 +156,7 @@ ) from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( RolloutDenoisingMixin, + rollout_transition, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -1336,7 +1337,11 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( scheduler.step, - {"generator": batch.generator, "eta": batch.eta, "batch": batch}, + { + "generator": batch.generator, + "eta": batch.eta, + "transition": rollout_transition(batch), + }, ) # Setup precision and autocast settings diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index f95485515015..e60c7018f2fa 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -61,6 +61,7 @@ from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( RolloutDenoisingMixin, + rollout_transition, ) from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import ( prepare_rollout_request_scheduler, @@ -1516,6 +1517,7 @@ def _denoise_once( }, guidance=None, ) + transition = rollout_transition(batch) do_cfg = guidance_scale > 1.0 # Control-CFG runs even when text guidance is off (its own extra @@ -1784,22 +1786,14 @@ def _denoise_once( timestep_value=t, step_index=i, ) - latents = scheduler.step( - noise_pred, - t, - latents, - generator=batch.generator, - batch=batch, - return_dict=False, - )[0] - else: - latents = scheduler.step( - noise_pred, - t, - latents, - generator=generator, - return_dict=False, - )[0] + latents = scheduler.step( + noise_pred, + t, + latents, + generator=generator, + transition=transition, + return_dict=False, + )[0] if action_noise_pred is not None: # Zero the velocity at conditioned (clean) action tokens and at diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index c0dd7d93a843..29f7fe72ef12 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -40,6 +40,9 @@ from sglang.multimodal_gen.runtime.platforms import ( current_platform, ) +from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( + rollout_transition, +) from sglang.multimodal_gen.runtime.server_args import ( ServerArgs, is_ltx2_two_stage_pipeline_name, @@ -186,12 +189,6 @@ def __init__( ] = OrderedDict() self._ltx2_coords_cache_max_entries = 4 - def _scheduler_step_kwargs(self, batch: Req, scheduler) -> dict: - return self.prepare_extra_func_kwargs( - scheduler.step, - {"generator": batch.generator, "eta": batch.eta, "batch": batch}, - ) - @staticmethod def _randn_like_with_batch_generators( reference_tensor: torch.Tensor, batch: Req @@ -1901,19 +1898,13 @@ def _stage2_midpoint_model_call( midpoint_model_call=_stage2_midpoint_model_call, ) else: - if batch.rollout: - ctx.scheduler._step_index = step.step_index - ctx.latents = ctx.scheduler.step( - model_video, - step.t_device, - ctx.latents, - return_dict=False, - **self._scheduler_step_kwargs(batch, ctx.scheduler), - )[0] - else: - ctx.latents = ctx.scheduler.step( - model_video, step.t_device, ctx.latents, return_dict=False - )[0] + ctx.latents = ctx.scheduler.step( + model_video, + step.t_device, + ctx.latents, + transition=rollout_transition(batch), + return_dict=False, + )[0] ctx.audio_latents = ctx.audio_scheduler.step( model_audio, step.t_device, ctx.audio_latents, return_dict=False )[0] diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 8800f18167b7..47f5c6b46461 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -4,7 +4,7 @@ from __future__ import annotations import math -from typing import Callable, Union +from typing import Any, Callable, Union import torch @@ -30,6 +30,7 @@ def __init__( log_prob_no_const: bool, debug_mode: bool, sde_step_indices: list[int] | None, + generator: Union[torch.Generator, list[torch.Generator]], shard_for_sp: Callable[[torch.Tensor], torch.Tensor], gather_for_sp: Callable[[torch.Tensor], torch.Tensor], sigmas: torch.Tensor, @@ -40,6 +41,7 @@ def __init__( self.log_prob_no_const = log_prob_no_const self.debug_mode = debug_mode self.sde_step_indices = sde_step_indices + self.generator = generator # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. self.shard_for_sp = shard_for_sp self.gather_for_sp = gather_for_sp @@ -112,6 +114,35 @@ def collect_rollout_debug_tensors(self) -> RolloutDebugTensors: rollout_model_outputs=model_outputs.cpu(), ) + def transition( + self, + *, + model_output: torch.FloatTensor, + sample: torch.FloatTensor, + current_sigma: torch.FloatTensor, + next_sigma: torch.FloatTensor, + dt: torch.FloatTensor, + generator: Any, + step_index: int, + timestep: Any, + ) -> torch.FloatTensor: + """Stand-in for ``scheduler._transition``, passed as ``step(transition=...)``; + ``sample`` arrives already upcast to fp32 by ``step()``. + + Draws variance noise from the request generator captured at creation, so + the caller's ``generator`` argument is not used. + """ + if current_sigma.dim() != 0: + raise ValueError("Rollout does not support per-token timesteps") + return self.flow_sde_sampling( + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + generator=self.generator, + step_index=step_index, + ) + def _get_or_create_rollout_noise_buffer( self, full_shape: tuple, diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index 90fc63b2ad4b..5fd749ca1469 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, Callable import torch @@ -39,6 +39,13 @@ def _kwargs_to_cpu(d: Any) -> Any: return d +def rollout_transition(batch: Req) -> Callable[..., torch.FloatTensor] | None: + """The ``transition`` to pass to ``scheduler.step`` for this request: the + visual-stream rollout controller, or None to keep the scheduler's own update.""" + rollout = batch.rollout_controllers.get("visual") + return None if rollout is None else rollout.transition + + class RolloutDenoisingMixin: def _maybe_prepare_rollout(self, batch: Req): """Create the rollout controller of the visual stream of this request. @@ -57,6 +64,7 @@ def _maybe_prepare_rollout(self, batch: Req): log_prob_no_const=batch.rollout_log_prob_no_const, debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, + generator=batch.generator, shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 046f44bc27a5..e3f0ed4c9184 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -4,6 +4,9 @@ import torch +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( FlowRolloutController, ) @@ -19,6 +22,7 @@ def _make_rollout_controller(batch, pipeline_config=None) -> FlowRolloutControll log_prob_no_const=batch.rollout_log_prob_no_const, debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, + generator=torch.Generator().manual_seed(0), shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], @@ -126,6 +130,40 @@ def test_ode_bit_exact_with_non_rollout_path(self): self.assertEqual(pre_cast_max_abs_diff, 0.0) self.assertEqual(post_cast_max_abs_diff, 0.0) + def test_ode_transition_matches_serving_step(self): + """Passing the rollout transition to a real FlowMatchEulerDiscreteScheduler + must reproduce rollout=False bit for bit in ODE mode: step() keeps its + index, sigma lookup and dtype casts around either transition.""" + serving, rollout_scheduler = ( + FlowMatchEulerDiscreteScheduler(shift=3.0), + FlowMatchEulerDiscreteScheduler(shift=3.0), + ) + for scheduler in (serving, rollout_scheduler): + scheduler.set_timesteps(5) + scheduler.set_begin_index(0) + current_rollout_controller = _make_rollout_controller( + self._build_batch(debug_mode=False) + ) + + g = torch.Generator().manual_seed(0) + sample = torch.randn(2, 4, 8, 8) + expected, actual = sample.clone(), sample.clone() + for t in serving.timesteps: + model_output = torch.randn(sample.shape, generator=g).to(torch.bfloat16) + expected = serving.step(model_output, t, expected, return_dict=False)[0] + actual = rollout_scheduler.step( + model_output, + t, + actual, + transition=current_rollout_controller.transition, + return_dict=False, + )[0] + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(rollout_scheduler.step_index, serving.step_index) + self.assertEqual( + len(current_rollout_controller.local_log_prob_sum), len(serving.timesteps) + ) + def test_ode_debug_tensors_have_shape_safe_noise_std(self): batch = self._build_batch(debug_mode=True) current_rollout_controller = _make_rollout_controller(batch) From eade7c48cb70d2ee30667ee379efa078b4cd538c Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 09/15] refactor(diffusion): record rollout latents in the transition instead of the denoise loops Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu --- .../pipelines_core/stages/denoising.py | 12 -- .../stages/model_specific_stages/cosmos3.py | 8 - .../post_training/flow_rollout_controller.py | 31 ++- .../post_training/rollout_denoising_mixin.py | 43 +--- .../test/unit/test_scheduler_rollout_unit.py | 184 +++++++++--------- 5 files changed, 127 insertions(+), 151 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 4e3541832c73..738d98bbe5ab 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -2188,18 +2188,6 @@ def _run_denoising_loop( t_host, timesteps_cpu, ) - # Capture the raw (pre-scale, pre-I2V-concat) noisy latent - # x_{t_i} for rollout trajectory collection. Must run - # BEFORE _run_denoising_step so ctx.latents is still the - # pre-step value. Gated on batch.rollout to keep the - # non-rollout path strictly untouched. - if collect_trajectory and batch.rollout: - self._maybe_append_dit_trajectory_step( - batch=batch, - latents=ctx.latents, - timestep_value=step.t_host, - step_index=step_index, - ) self._run_denoising_step(ctx, step, batch, server_args) if collect_trajectory: self._record_trajectory(ctx, step, batch, server_args) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index e60c7018f2fa..defb43b3a252 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -1778,14 +1778,6 @@ def _denoise_once( if velocity_mask is not None: noise_pred = noise_pred * velocity_mask - if batch.rollout: - # Capture the pre-step x_{t_i} before the scheduler advances it. - self._maybe_append_dit_trajectory_step( - batch=batch, - latents=latents, - timestep_value=t, - step_index=i, - ) latents = scheduler.step( noise_pred, t, diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 47f5c6b46461..2a80fb1d6a74 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -31,6 +31,8 @@ def __init__( debug_mode: bool, sde_step_indices: list[int] | None, generator: Union[torch.Generator, list[torch.Generator]], + return_dit_trajectory: bool, + return_step_indices: list[int] | None, shard_for_sp: Callable[[torch.Tensor], torch.Tensor], gather_for_sp: Callable[[torch.Tensor], torch.Tensor], sigmas: torch.Tensor, @@ -42,6 +44,8 @@ def __init__( self.debug_mode = debug_mode self.sde_step_indices = sde_step_indices self.generator = generator + self.return_dit_trajectory = return_dit_trajectory + self.return_step_indices = return_step_indices # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. self.shard_for_sp = shard_for_sp self.gather_for_sp = gather_for_sp @@ -61,6 +65,9 @@ def __init__( self.local_noise_std_devs: list[torch.Tensor] = [] self.local_model_outputs: list[torch.Tensor] = [] + self.step_latents: list[torch.Tensor] = [] + self.step_timesteps: list[torch.Tensor] = [] + def _reset_rollout_debug_tensors(self) -> None: self.local_variance_noises = [] self.local_prev_sample_means = [] @@ -130,10 +137,14 @@ def transition( ``sample`` arrives already upcast to fp32 by ``step()``. Draws variance noise from the request generator captured at creation, so - the caller's ``generator`` argument is not used. + the caller's ``generator`` argument is not used. Records ``sample`` as the + step's x_t before advancing it. """ if current_sigma.dim() != 0: raise ValueError("Rollout does not support per-token timesteps") + self._maybe_append_dit_trajectory_step( + latents=sample, timestep_value=timestep, step_index=step_index + ) return self.flow_sde_sampling( model_output=model_output, sample=sample, @@ -364,3 +375,21 @@ def collect_rollout_log_probs(self) -> torch.Tensor: ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() + + def _maybe_append_dit_trajectory_step( + self, + latents: torch.Tensor, + timestep_value: torch.Tensor, + step_index: int, + ) -> None: + if not self.return_dit_trajectory: + return + + if ( + self.return_step_indices is not None + and step_index not in self.return_step_indices + ): + return + + self.step_latents.append(latents.detach()) + self.step_timesteps.append(timestep_value.detach().cpu()) diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index 5fd749ca1469..8b545f6f55b4 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -65,6 +65,8 @@ def _maybe_prepare_rollout(self, batch: Req): debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, generator=batch.generator, + return_dit_trajectory=batch.rollout_return_dit_trajectory, + return_step_indices=batch.rollout_return_step_indices, shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], @@ -91,7 +93,6 @@ def _maybe_collect_rollout_log_probs(self, batch: Req): batch.rollout_trajectory_data.rollout_debug_tensors = ( rollout.collect_rollout_debug_tensors() ) - batch.rollout_controllers = {} def _postprocess_rollout_outputs( self, @@ -111,8 +112,7 @@ def _postprocess_rollout_outputs( # Append final denoised latent as the (T+1)-th entry (step_index=T), # routed through the same filter so rollout_return_step_indices can # include/exclude it. - self._maybe_append_dit_trajectory_step( - batch=batch, + batch.rollout_controllers["visual"]._maybe_append_dit_trajectory_step( latents=latents, timestep_value=final_timestep, step_index=num_inference_steps, @@ -121,6 +121,7 @@ def _postprocess_rollout_outputs( batch=batch, pipeline_config=server_args.pipeline_config, ) + batch.rollout_controllers = {} def _maybe_init_denoising_env_collection( self, @@ -155,55 +156,31 @@ def _maybe_init_denoising_env_collection( batch._rollout_denoising_env_state = { "env": env, - "step_latents": [], - "step_timesteps": [], "pos_cond_kwargs_src": pos_src, "neg_cond_kwargs_src": neg_src, } - def _maybe_append_dit_trajectory_step( - self, - batch, - latents: torch.Tensor, - timestep_value: torch.Tensor, - step_index: int, - ) -> None: - if not batch.rollout or not batch.rollout_return_dit_trajectory: - return - state = getattr(batch, "_rollout_denoising_env_state", None) - if state is None: - return - - return_step_indices = getattr(batch, "rollout_return_step_indices", None) - if return_step_indices is not None and step_index not in return_step_indices: - return - - state["step_latents"].append(latents.detach()) - state["step_timesteps"].append(timestep_value.detach().cpu()) - def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> None: state = getattr(batch, "_rollout_denoising_env_state", None) if state is None: return env: RolloutDenoisingEnv | None = state["env"] - step_latents: list[torch.Tensor] = state["step_latents"] - step_timesteps: list[torch.Tensor] = state["step_timesteps"] + rollout = batch.rollout_controllers["visual"] + step_latents: list[torch.Tensor] = rollout.step_latents + step_timesteps: list[torch.Tensor] = rollout.step_timesteps if batch.rollout_trajectory_data is None: batch.rollout_trajectory_data = RolloutTrajectoryData() if step_latents and batch.rollout_return_dit_trajectory: - step_latents_tensor = torch.stack(step_latents, dim=1) - step_latents_tensor = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=step_latents_tensor, + step_latents_tensor = rollout.gather_for_sp( + torch.stack(step_latents, dim=1) ) batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( latents=step_latents_tensor.cpu(), timesteps=torch.stack(step_timesteps, dim=0).cpu(), - sigmas=batch.scheduler.sigmas.detach().cpu().clone(), + sigmas=rollout.sigmas.detach().cpu().clone(), ) if env is not None and batch.rollout_return_denoising_env: diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index e3f0ed4c9184..d12cfc86d360 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -1,3 +1,28 @@ +"""Unit tests for FlowRolloutController, whose ``transition`` the rollout passes +to ``scheduler.step()``. + + scheduler.step(model_output, t, x_t, transition=current_rollout_controller.transition) + | upcast x_t to fp32, look up sigma_t / sigma_next, dt = sigma_next - sigma_t + v + current_rollout_controller.transition(...) + |-- record x_t into the trajectory (filtered by return_step_indices) + +-- flow_sde_sampling(...) + |-- ode : x_next = x_t + dt * v (no noise, log_prob 0) + |-- sde / cps: x_next = mean + std * noise (fp32, log_prob summed) + +-- step outside sde_step_indices -> ode form + v + x_next (cast back by step()) + +TestSchedulerRolloutOdeUnit + ODE rollout == serving step, bit for bit (bare update and through step()); + ODE draws no noise; the trajectory records x_t, not x_next; + debug tensors keep [B, T, ...] shapes. +TestSchedulerFlowGRPOStepAlignmentUnit + SDE / CPS match the verbatim FlowGRPO reference; bf16 model_output still + yields fp32 noise and log-probs; sde_step_indices gates the noise draw; + return_step_indices filters latents while timesteps cover every step. +""" + import math import types import unittest @@ -23,6 +48,8 @@ def _make_rollout_controller(batch, pipeline_config=None) -> FlowRolloutControll debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, generator=torch.Generator().manual_seed(0), + return_dit_trajectory=batch.rollout_return_dit_trajectory, + return_step_indices=batch.rollout_return_step_indices, shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], @@ -40,6 +67,8 @@ def _build_batch(self, *, debug_mode: bool) -> types.SimpleNamespace: rollout_sde_type="ode", rollout_debug_mode=debug_mode, rollout_sde_step_indices=None, + rollout_return_dit_trajectory=False, + rollout_return_step_indices=None, latents=torch.zeros(2, 4, 8, 8, dtype=torch.float32), ) @@ -164,6 +193,29 @@ def test_ode_transition_matches_serving_step(self): len(current_rollout_controller.local_log_prob_sum), len(serving.timesteps) ) + def test_transition_records_pre_step_latents(self): + """The rollout transition records x_t (the latents entering the step), + not the updated x_next it returns.""" + scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0) + scheduler.set_timesteps(5) + scheduler.set_begin_index(0) + batch = self._build_batch(debug_mode=False) + batch.rollout_return_dit_trajectory = True + current_rollout_controller = _make_rollout_controller(batch) + sample = torch.randn(2, 4, 8, 8) + model_output = torch.randn(2, 4, 8, 8) + prev_sample = scheduler.step( + model_output, + scheduler.timesteps[0], + sample, + transition=current_rollout_controller.transition, + return_dict=False, + )[0] + self.assertTrue(torch.equal(current_rollout_controller.step_latents[0], sample)) + self.assertFalse( + torch.equal(current_rollout_controller.step_latents[0], prev_sample) + ) + def test_ode_debug_tensors_have_shape_safe_noise_std(self): batch = self._build_batch(debug_mode=True) current_rollout_controller = _make_rollout_controller(batch) @@ -275,6 +327,8 @@ def _build_batch( rollout_sde_type=sde_type, rollout_debug_mode=True, rollout_sde_step_indices=None, + rollout_return_dit_trajectory=False, + rollout_return_step_indices=None, latents=torch.empty(shape, dtype=torch.float32), ) @@ -424,9 +478,6 @@ def test_timestep_filters_gate_sde_and_trajectory(self): injection (excluded steps = ODE transition + zero log-prob); independently, rollout_return_step_indices gates the dit_trajectory append. Both features are exercised here because they share the same step_index predicate.""" - from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( - RolloutDenoisingMixin, - ) # --- Part 1: rollout_sde_step_indices gates SDE noise injection --- shape = (1, 4, 8, 8) @@ -439,6 +490,8 @@ def test_timestep_filters_gate_sde_and_trajectory(self): rollout_sde_type="sde", rollout_debug_mode=False, rollout_sde_step_indices=[1], # only step 1 is stochastic + rollout_return_dit_trajectory=False, + rollout_return_step_indices=None, latents=torch.empty(shape, dtype=torch.float32), ) current_rollout_controller = _make_rollout_controller(batch, pipeline_config) @@ -492,105 +545,42 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertTrue(torch.all(elem_count > 0)) # --- Part 2: rollout_return_step_indices gates dit trajectory append --- - class _DummyDit(RolloutDenoisingMixin): - pass - - dit = _DummyDit() lat = torch.zeros(1, 4, 8, 8) ts = torch.tensor(0.5) - # Filter [0, 2] over steps 0,1,2 → steps 0 and 2 appended, step 1 skipped. - traj_filtered = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=[0, 2], - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_filtered, - latents=lat, - timestep_value=ts, - step_index=i, - ) - self.assertEqual( - len(traj_filtered._rollout_denoising_env_state["step_latents"]), 2 - ) - self.assertEqual( - len(traj_filtered._rollout_denoising_env_state["step_timesteps"]), 3 - ) - - # None (default) → all steps appended (back-compat). - traj_all = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=None, - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_all, - latents=lat, - timestep_value=ts, - step_index=i, - ) - self.assertEqual(len(traj_all._rollout_denoising_env_state["step_latents"]), 3) - - # Filter excludes step_index=T (the final/(T+1)-th latent appended by - # _postprocess_rollout_outputs). Simulate T=3 loop steps + final append. - traj_exclude_final = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=[0, 1, 2], # excludes T=3 - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_exclude_final, - latents=lat, - timestep_value=ts, - step_index=i, + def _recorded_steps(return_step_indices: list[int] | None) -> list[int]: + # Three loop steps, then the final latent appended at step_index=T=3. + current_rollout_controller = _make_rollout_controller( + types.SimpleNamespace( + rollout_log_prob_no_const=False, + rollout_noise_level=0.5, + rollout_sde_type="sde", + rollout_debug_mode=False, + rollout_sde_step_indices=None, + rollout_return_dit_trajectory=True, + rollout_return_step_indices=return_step_indices, + latents=lat, + ) ) - # Mimic the final append routed through the same filter. - dit._maybe_append_dit_trajectory_step( - batch=traj_exclude_final, - latents=lat, - timestep_value=torch.zeros(()), - step_index=3, - ) - self.assertEqual( - len(traj_exclude_final._rollout_denoising_env_state["step_latents"]), 3 - ) - self.assertEqual( - len(traj_exclude_final._rollout_denoising_env_state["step_timesteps"]), 4 - ) - - # Filter includes only step_index=T → only the final latent survives. - traj_only_final = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=[3], - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_only_final, - latents=lat, - timestep_value=ts, - step_index=i, + # Step i records the timestep value i, so the timesteps name the steps kept. + for i in range(4): + current_rollout_controller._maybe_append_dit_trajectory_step( + latents=lat, timestep_value=torch.tensor(float(i)), step_index=i + ) + self.assertEqual( + len(current_rollout_controller.step_latents), + len(current_rollout_controller.step_timesteps), ) - dit._maybe_append_dit_trajectory_step( - batch=traj_only_final, - latents=lat, - timestep_value=torch.zeros(()), - step_index=3, - ) - self.assertEqual( - len(traj_only_final._rollout_denoising_env_state["step_latents"]), 1 - ) - self.assertEqual( - len(traj_only_final._rollout_denoising_env_state["step_timesteps"]), 4 - ) + return [int(t) for t in current_rollout_controller.step_timesteps] + + # Filter [0, 2] over steps 0,1,2 -> steps 0 and 2 appended, step 1 skipped. + self.assertEqual(_recorded_steps([0, 2]), [0, 2]) + # None (default) -> all steps appended, including the final latent. + self.assertEqual(_recorded_steps(None), [0, 1, 2, 3]) + # Filter excludes step_index=T (the final/(T+1)-th latent). + self.assertEqual(_recorded_steps([0, 1, 2]), [0, 1, 2]) + # Filter includes only step_index=T -> only the final latent survives. + self.assertEqual(_recorded_steps([3]), [3]) if __name__ == "__main__": From b42c596a102543c40e4b12a9935fb9b4ab7405e3 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 10/15] refactor(diffusion): replace the rollout denoising mixin with rollout hook functions Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu --- .../pipelines_core/stages/denoising.py | 44 ++-- .../stages/model_specific_stages/cosmos3.py | 58 ++--- .../model_specific_stages/ltx_2/denoising.py | 2 +- .../post_training/flow_rollout_controller.py | 2 +- .../post_training/rollout_denoising_mixin.py | 203 ---------------- .../runtime/post_training/rollout_hooks.py | 230 ++++++++++++++++++ 6 files changed, 284 insertions(+), 255 deletions(-) delete mode 100644 python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py create mode 100644 python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 738d98bbe5ab..6017cd55aef7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -154,8 +154,10 @@ AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( - RolloutDenoisingMixin, +from sglang.multimodal_gen.runtime.post_training.rollout_hooks import ( + finalize_rollout_if_enabled, + init_rollout_env_if_enabled, + prepare_rollout_if_enabled, rollout_transition, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -362,7 +364,7 @@ class DualTransformerExecutionMode(str, Enum): PAIRED_PER_STEP = "paired_per_step" -class DenoisingStage(PipelineStage, RolloutDenoisingMixin): +class DenoisingStage(PipelineStage): """ Stage for running the denoising loop in diffusion pipelines. @@ -1331,8 +1333,9 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): if freshly_loaded: register_loaded_transformer(self, server_args, pipeline) - if batch.rollout: - self._maybe_prepare_rollout(batch) + prepare_rollout_if_enabled( + batch=batch, pipeline_config=server_args.pipeline_config + ) # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( @@ -2107,26 +2110,23 @@ def _denoise( Run the denoising loop. """ ctx = self._prepare_denoising_loop(batch, server_args) - if batch.rollout: - self._maybe_init_denoising_env_collection( - batch=batch, - pipeline_config=server_args.pipeline_config, - image_kwargs=ctx.image_kwargs, - pos_cond_kwargs=ctx.pos_cond_kwargs, - neg_cond_kwargs=ctx.neg_cond_kwargs, - guidance=ctx.guidance, - ) + init_rollout_env_if_enabled( + batch=batch, + image_kwargs=ctx.image_kwargs, + pos_cond_kwargs=ctx.pos_cond_kwargs, + neg_cond_kwargs=ctx.neg_cond_kwargs, + guidance=ctx.guidance, + ) timesteps_cpu = self._run_denoising_loop(ctx, batch, server_args) # collect rollout outputs before finalization gathers or replaces latents - if batch.rollout: - self._postprocess_rollout_outputs( - batch=batch, - latents=ctx.latents, - num_inference_steps=len(timesteps_cpu), - final_timestep=timesteps_cpu.new_zeros(()), - server_args=server_args, - ) + finalize_rollout_if_enabled( + batch=batch, + latents=ctx.latents, + num_inference_steps=len(timesteps_cpu), + final_timestep=timesteps_cpu.new_zeros(()), + pipeline_config=server_args.pipeline_config, + ) self._finalize_denoising_loop(ctx, batch, server_args) return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index defb43b3a252..0bb0073e0ba0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -59,8 +59,10 @@ VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform -from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( - RolloutDenoisingMixin, +from sglang.multimodal_gen.runtime.post_training.rollout_hooks import ( + finalize_rollout_if_enabled, + init_rollout_env_if_enabled, + prepare_rollout_if_enabled, rollout_transition, ) from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import ( @@ -1060,7 +1062,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: return batch -class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): +class Cosmos3DenoisingStage(PipelineStage): def default_workload_iterations( self, batch: Req, num_inference_steps: int ) -> int | None: @@ -1500,23 +1502,24 @@ def _denoise_once( raise ValueError( "Cosmos3 rollout does not support action/sound modalities." ) - self._maybe_prepare_rollout(batch) - self._maybe_init_denoising_env_collection( - batch=batch, - pipeline_config=server_args.pipeline_config, - image_kwargs={}, - pos_cond_kwargs={ - "text_ids": cond_text_ids, - "text_mask": cond_text_mask, - "fps": fps, - }, - neg_cond_kwargs={ - "text_ids": uncond_text_ids, - "text_mask": uncond_text_mask, - "fps": fps, - }, - guidance=None, - ) + prepare_rollout_if_enabled( + batch=batch, pipeline_config=server_args.pipeline_config + ) + init_rollout_env_if_enabled( + batch=batch, + image_kwargs={}, + pos_cond_kwargs={ + "text_ids": cond_text_ids, + "text_mask": cond_text_mask, + "fps": fps, + }, + neg_cond_kwargs={ + "text_ids": uncond_text_ids, + "text_mask": uncond_text_mask, + "fps": fps, + }, + guidance=None, + ) transition = rollout_transition(batch) do_cfg = guidance_scale > 1.0 @@ -1834,14 +1837,13 @@ def _denoise_once( # next request's steps. self.transformer.reset_denoising_step() - if batch.rollout: - self._postprocess_rollout_outputs( - batch=batch, - latents=latents, - num_inference_steps=len(timesteps), - final_timestep=timesteps.new_zeros(()).cpu(), - server_args=server_args, - ) + finalize_rollout_if_enabled( + batch=batch, + latents=latents, + num_inference_steps=len(timesteps), + final_timestep=timesteps.new_zeros(()).cpu(), + pipeline_config=server_args.pipeline_config, + ) batch.latents = latents if action_latents is not None: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index 29f7fe72ef12..f348070628e6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -40,7 +40,7 @@ from sglang.multimodal_gen.runtime.platforms import ( current_platform, ) -from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( +from sglang.multimodal_gen.runtime.post_training.rollout_hooks import ( rollout_transition, ) from sglang.multimodal_gen.runtime.server_args import ( diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 2a80fb1d6a74..bd170e152983 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -36,7 +36,7 @@ def __init__( shard_for_sp: Callable[[torch.Tensor], torch.Tensor], gather_for_sp: Callable[[torch.Tensor], torch.Tensor], sigmas: torch.Tensor, - latents_shape: tuple | None, + latents_shape: tuple, ) -> None: self.sde_type = sde_type self.noise_level = noise_level diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py deleted file mode 100644 index 8b545f6f55b4..000000000000 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ /dev/null @@ -1,203 +0,0 @@ -"""Mixin for rollout-related denoising hooks. - -Moved out of DenoisingStage to keep the core stage lean. -""" - -from __future__ import annotations - -from typing import Any, Callable - -import torch - -from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler, -) -from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - FlowRolloutController, -) -from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( - RolloutDenoisingEnv, - RolloutDitTrajectory, - RolloutTrajectoryData, -) -from sglang.multimodal_gen.runtime.post_training.sp_utils import ( - gather_stacked_latents_for_sp, -) -from sglang.multimodal_gen.runtime.server_args import ServerArgs - - -def _kwargs_to_cpu(d: Any) -> Any: - if isinstance(d, torch.Tensor): - return d.detach().cpu() - if isinstance(d, dict): - return {k: _kwargs_to_cpu(v) for k, v in d.items()} - if isinstance(d, list): - return [_kwargs_to_cpu(v) for v in d] - if isinstance(d, tuple): - return tuple(_kwargs_to_cpu(v) for v in d) - return d - - -def rollout_transition(batch: Req) -> Callable[..., torch.FloatTensor] | None: - """The ``transition`` to pass to ``scheduler.step`` for this request: the - visual-stream rollout controller, or None to keep the scheduler's own update.""" - rollout = batch.rollout_controllers.get("visual") - return None if rollout is None else rollout.transition - - -class RolloutDenoisingMixin: - def _maybe_prepare_rollout(self, batch: Req): - """Create the rollout controller of the visual stream of this request. - - Must run before SP sharding so the controller records the full latent shape. - """ - scheduler = batch.scheduler - if not isinstance(scheduler, FlowMatchEulerDiscreteScheduler): - raise ValueError(f"Scheduler {type(scheduler)} does not support rollout") - - pipeline_config = self.server_args.pipeline_config - batch.rollout_controllers = { - "visual": FlowRolloutController( - sde_type=batch.rollout_sde_type, - noise_level=float(batch.rollout_noise_level), - log_prob_no_const=batch.rollout_log_prob_no_const, - debug_mode=batch.rollout_debug_mode, - sde_step_indices=batch.rollout_sde_step_indices, - generator=batch.generator, - return_dit_trajectory=batch.rollout_return_dit_trajectory, - return_step_indices=batch.rollout_return_step_indices, - shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( - batch=batch, latents=latents - )[0], - gather_for_sp=lambda stacked_latents: gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=stacked_latents, - ), - sigmas=scheduler.sigmas, - latents_shape=( - tuple(batch.latents.shape) if batch.latents is not None else None - ), - ) - } - - def _maybe_collect_rollout_log_probs(self, batch: Req): - rollout = batch.rollout_controllers["visual"] - if batch.rollout_trajectory_data is None: - batch.rollout_trajectory_data = RolloutTrajectoryData() - batch.rollout_trajectory_data.rollout_log_probs = ( - rollout.collect_rollout_log_probs() - ) - if batch.rollout_debug_mode: - batch.rollout_trajectory_data.rollout_debug_tensors = ( - rollout.collect_rollout_debug_tensors() - ) - - def _postprocess_rollout_outputs( - self, - batch: Req, - latents: torch.Tensor, - num_inference_steps: int, - final_timestep: torch.Tensor, - server_args: ServerArgs, - ) -> None: - """Finalize rollout-only outputs. - - Must be called before ``_post_denoising_loop`` so that ``latents`` (the - last ``scheduler.step`` output) is still SP-sharded and can be gathered - uniformly with the per-step trajectory latents. - """ - self._maybe_collect_rollout_log_probs(batch) - # Append final denoised latent as the (T+1)-th entry (step_index=T), - # routed through the same filter so rollout_return_step_indices can - # include/exclude it. - batch.rollout_controllers["visual"]._maybe_append_dit_trajectory_step( - latents=latents, - timestep_value=final_timestep, - step_index=num_inference_steps, - ) - self._maybe_finalize_denoising_env_collection( - batch=batch, - pipeline_config=server_args.pipeline_config, - ) - batch.rollout_controllers = {} - - def _maybe_init_denoising_env_collection( - self, - batch, - pipeline_config, - image_kwargs: dict[str, Any], - pos_cond_kwargs: dict[str, Any], - neg_cond_kwargs: dict[str, Any], - guidance: torch.Tensor | None, - ) -> None: - collect_env = batch.rollout_return_denoising_env - collect_traj = batch.rollout_return_dit_trajectory - if not (collect_env or collect_traj): - batch._rollout_denoising_env_state = None - return - - if collect_env: - env = RolloutDenoisingEnv( - image_kwargs=_kwargs_to_cpu(image_kwargs), - pos_cond_kwargs=_kwargs_to_cpu(pos_cond_kwargs), - neg_cond_kwargs=( - _kwargs_to_cpu(neg_cond_kwargs) if neg_cond_kwargs else None - ), - guidance=guidance.detach().cpu() if guidance is not None else None, - ) - pos_src = pos_cond_kwargs - neg_src = neg_cond_kwargs - else: - env = None - pos_src = None - neg_src = None - - batch._rollout_denoising_env_state = { - "env": env, - "pos_cond_kwargs_src": pos_src, - "neg_cond_kwargs_src": neg_src, - } - - def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> None: - state = getattr(batch, "_rollout_denoising_env_state", None) - if state is None: - return - - env: RolloutDenoisingEnv | None = state["env"] - rollout = batch.rollout_controllers["visual"] - step_latents: list[torch.Tensor] = rollout.step_latents - step_timesteps: list[torch.Tensor] = rollout.step_timesteps - - if batch.rollout_trajectory_data is None: - batch.rollout_trajectory_data = RolloutTrajectoryData() - - if step_latents and batch.rollout_return_dit_trajectory: - step_latents_tensor = rollout.gather_for_sp( - torch.stack(step_latents, dim=1) - ) - batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( - latents=step_latents_tensor.cpu(), - timesteps=torch.stack(step_timesteps, dim=0).cpu(), - sigmas=rollout.sigmas.detach().cpu().clone(), - ) - - if env is not None and batch.rollout_return_denoising_env: - gather_fn = getattr( - pipeline_config, "gather_denoising_env_static_for_sp", None - ) - - pos_src = state.get("pos_cond_kwargs_src") - if pos_src is not None and env.pos_cond_kwargs is not None: - gathered_pos = gather_fn(batch, pos_src) if gather_fn else pos_src - env.pos_cond_kwargs = _kwargs_to_cpu(gathered_pos) - - neg_src = state.get("neg_cond_kwargs_src") - if neg_src is not None and env.neg_cond_kwargs is not None: - gathered_neg = gather_fn(batch, neg_src) if gather_fn else neg_src - env.neg_cond_kwargs = _kwargs_to_cpu(gathered_neg) - - batch.rollout_trajectory_data.denoising_env = env - - batch._rollout_denoising_env_state = None diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py new file mode 100644 index 000000000000..9cf41b82d850 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rollout hooks for denoise loops. + +A loop calls ``prepare_rollout_if_enabled`` before SP sharding, +``init_rollout_env_if_enabled`` at the loop start and +``finalize_rollout_if_enabled`` at the loop end, and passes +``rollout_transition(batch)`` to ``scheduler.step``. Every hook is a no-op for +serving requests. +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch + +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + FlowRolloutController, +) +from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutDenoisingEnv, + RolloutDitTrajectory, + RolloutTrajectoryData, +) +from sglang.multimodal_gen.runtime.post_training.sp_utils import ( + gather_stacked_latents_for_sp, +) + + +# Key of the image/video latent stream in ``Req.rollout_controllers``. +VISUAL_STREAM = "visual" + + +def _kwargs_to_cpu(d: Any) -> Any: + if isinstance(d, torch.Tensor): + return d.detach().cpu() + if isinstance(d, dict): + return {k: _kwargs_to_cpu(v) for k, v in d.items()} + if isinstance(d, list): + return [_kwargs_to_cpu(v) for v in d] + if isinstance(d, tuple): + return tuple(_kwargs_to_cpu(v) for v in d) + return d + + +def rollout_transition(batch: Req) -> Callable[..., torch.FloatTensor] | None: + """The ``transition`` to pass to ``scheduler.step`` for this request: the + visual-stream rollout controller, or None to keep the scheduler's own update.""" + current_rollout_controller = batch.rollout_controllers.get(VISUAL_STREAM) + return ( + None + if current_rollout_controller is None + else current_rollout_controller.transition + ) + + +def _build_scheduler_rollout_controller( + *, batch: Req, pipeline_config: Any +) -> FlowRolloutController: + """The visual-stream controller of a loop that steps ``batch.scheduler``.""" + scheduler = batch.scheduler + if not isinstance(scheduler, FlowMatchEulerDiscreteScheduler): + raise ValueError(f"Scheduler {type(scheduler)} does not support rollout") + + current_rollout_controller = FlowRolloutController( + sde_type=batch.rollout_sde_type, + noise_level=float(batch.rollout_noise_level), + log_prob_no_const=batch.rollout_log_prob_no_const, + debug_mode=batch.rollout_debug_mode, + sde_step_indices=batch.rollout_sde_step_indices, + generator=batch.generator, + return_dit_trajectory=batch.rollout_return_dit_trajectory, + return_step_indices=batch.rollout_return_step_indices, + shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + batch=batch, latents=latents + )[0], + gather_for_sp=lambda stacked_latents: gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=stacked_latents, + ), + sigmas=scheduler.sigmas, + latents_shape=tuple(batch.latents.shape), + ) + return current_rollout_controller + + +def prepare_rollout_if_enabled( + *, + batch: Req, + pipeline_config: Any, + build_controller: Callable[..., FlowRolloutController] = ( + _build_scheduler_rollout_controller + ), +) -> None: + """Create the rollout controller of the visual stream of this request. + + Must run before SP sharding so the controller records the full latent shape. + ``build_controller(batch=, pipeline_config=)`` returns it; a loop without a + scheduler passes its own. + """ + if not batch.rollout: + return + current_rollout_controller = build_controller( + batch=batch, pipeline_config=pipeline_config + ) + batch.rollout_controllers = {VISUAL_STREAM: current_rollout_controller} + + +def init_rollout_env_if_enabled( + *, + batch: Req, + image_kwargs: dict[str, Any], + pos_cond_kwargs: dict[str, Any], + neg_cond_kwargs: dict[str, Any], + guidance: torch.Tensor | None, +) -> None: + """Snapshot the DiT conditioning that the trainer replays.""" + if not batch.rollout: + return + collect_env = batch.rollout_return_denoising_env + collect_traj = batch.rollout_return_dit_trajectory + if not (collect_env or collect_traj): + batch._rollout_denoising_env_state = None + return + + if collect_env: + env = RolloutDenoisingEnv( + image_kwargs=_kwargs_to_cpu(image_kwargs), + pos_cond_kwargs=_kwargs_to_cpu(pos_cond_kwargs), + neg_cond_kwargs=( + _kwargs_to_cpu(neg_cond_kwargs) if neg_cond_kwargs else None + ), + guidance=guidance.detach().cpu() if guidance is not None else None, + ) + pos_src = pos_cond_kwargs + neg_src = neg_cond_kwargs + else: + env = None + pos_src = None + neg_src = None + + batch._rollout_denoising_env_state = { + "env": env, + "pos_cond_kwargs_src": pos_src, + "neg_cond_kwargs_src": neg_src, + } + + +def finalize_rollout_if_enabled( + *, + batch: Req, + latents: torch.Tensor, + num_inference_steps: int, + final_timestep: torch.Tensor, + pipeline_config: Any, +) -> None: + """Finalize rollout-only outputs. + + Must be called before ``_post_denoising_loop`` so that ``latents`` (the + last ``scheduler.step`` output) is still SP-sharded and can be gathered + uniformly with the per-step trajectory latents. + """ + if not batch.rollout: + return + _collect_rollout_log_probs(batch) + # Append final denoised latent as the (T+1)-th entry (step_index=T), + # routed through the same filter so rollout_return_step_indices can + # include/exclude it. + batch.rollout_controllers[VISUAL_STREAM]._maybe_append_dit_trajectory_step( + latents=latents, + timestep_value=final_timestep, + step_index=num_inference_steps, + ) + _finalize_rollout_env(batch=batch, pipeline_config=pipeline_config) + batch.rollout_controllers = {} + + +def _collect_rollout_log_probs(batch: Req) -> None: + current_rollout_controller = batch.rollout_controllers[VISUAL_STREAM] + if batch.rollout_trajectory_data is None: + batch.rollout_trajectory_data = RolloutTrajectoryData() + batch.rollout_trajectory_data.rollout_log_probs = ( + current_rollout_controller.collect_rollout_log_probs() + ) + if batch.rollout_debug_mode: + batch.rollout_trajectory_data.rollout_debug_tensors = ( + current_rollout_controller.collect_rollout_debug_tensors() + ) + + +def _finalize_rollout_env(*, batch: Req, pipeline_config: Any) -> None: + state = batch._rollout_denoising_env_state + if state is None: + return + + env: RolloutDenoisingEnv | None = state["env"] + current_rollout_controller = batch.rollout_controllers[VISUAL_STREAM] + step_latents: list[torch.Tensor] = current_rollout_controller.step_latents + step_timesteps: list[torch.Tensor] = current_rollout_controller.step_timesteps + + if step_latents: + step_latents_tensor = current_rollout_controller.gather_for_sp( + torch.stack(step_latents, dim=1) + ) + batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( + latents=step_latents_tensor.cpu(), + timesteps=torch.stack(step_timesteps, dim=0).cpu(), + sigmas=current_rollout_controller.sigmas.detach().cpu().clone(), + ) + + if env is not None: + env.pos_cond_kwargs = _kwargs_to_cpu( + pipeline_config.gather_denoising_env_static_for_sp( + batch, state["pos_cond_kwargs_src"] + ) + ) + if env.neg_cond_kwargs is not None: + env.neg_cond_kwargs = _kwargs_to_cpu( + pipeline_config.gather_denoising_env_static_for_sp( + batch, state["neg_cond_kwargs_src"] + ) + ) + batch.rollout_trajectory_data.denoising_env = env + + batch._rollout_denoising_env_state = None From dc188a7b92d15503f5d9be2c075e511e1052ae2f Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 15:48:20 -0700 Subject: [PATCH 11/15] feat(diffusion): let a model pass its own ODE update to the rollout controller Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 12 ++++- .../test/unit/test_scheduler_rollout_unit.py | 51 ++++++++++++++++++- 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index bd170e152983..3dd571305141 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -37,6 +37,7 @@ def __init__( gather_for_sp: Callable[[torch.Tensor], torch.Tensor], sigmas: torch.Tensor, latents_shape: tuple, + ode_transition: Callable[..., torch.Tensor] | None = None, ) -> None: self.sde_type = sde_type self.noise_level = noise_level @@ -55,6 +56,10 @@ def __init__( # never divides by 1 - sigma = 0. self.sigma_max = sigmas[min(1, len(sigmas) - 1)].item() self.latents_shape = latents_shape + # The model's own deterministic update, called as + # ode_transition(model_output=, sample=, step_index=) on ODE steps; + # None means the flow-matching Euler update sample + dt * model_output. + self.ode_transition = ode_transition self.noise_buffer: torch.Tensor | None = None self.local_log_prob_sum: list[torch.Tensor] = [] @@ -299,7 +304,12 @@ def flow_sde_sampling( log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2) elif effective_sde_type == "ode": - prev_sample = sample + dt * model_output + if self.ode_transition is None: + prev_sample = sample + dt * model_output + else: + prev_sample = self.ode_transition( + model_output=model_output, sample=sample, step_index=step_index + ) prev_sample_mean = prev_sample variance_noise = torch.zeros_like(model_output) noise_std_dev = torch.zeros( diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index d12cfc86d360..502714ce2049 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -8,6 +8,7 @@ |-- record x_t into the trajectory (filtered by return_step_indices) +-- flow_sde_sampling(...) |-- ode : x_next = x_t + dt * v (no noise, log_prob 0) + | or the model's own ode_transition |-- sde / cps: x_next = mean + std * noise (fp32, log_prob summed) +-- step outside sde_step_indices -> ode form v @@ -16,6 +17,7 @@ TestSchedulerRolloutOdeUnit ODE rollout == serving step, bit for bit (bare update and through step()); ODE draws no noise; the trajectory records x_t, not x_next; + a model's own ode_transition replaces the Euler update on ODE steps only; debug tensors keep [B, T, ...] shapes. TestSchedulerFlowGRPOStepAlignmentUnit SDE / CPS match the verbatim FlowGRPO reference; bf16 model_output still @@ -39,7 +41,9 @@ _SIGMAS = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0], dtype=torch.float32) -def _make_rollout_controller(batch, pipeline_config=None) -> FlowRolloutController: +def _make_rollout_controller( + batch, pipeline_config=None, ode_transition=None +) -> FlowRolloutController: """The controller the denoising stage creates for ``batch`` (single GPU).""" return FlowRolloutController( sde_type=batch.rollout_sde_type, @@ -56,6 +60,7 @@ def _make_rollout_controller(batch, pipeline_config=None) -> FlowRolloutControll gather_for_sp=lambda stacked_latents: stacked_latents, sigmas=_SIGMAS, latents_shape=tuple(batch.latents.shape), + ode_transition=ode_transition, ) @@ -216,6 +221,50 @@ def test_transition_records_pre_step_latents(self): torch.equal(current_rollout_controller.step_latents[0], prev_sample) ) + def test_ode_transition_replaces_euler_on_ode_steps_only(self): + """A model's own ode_transition runs on the steps outside the SDE window + and nowhere else; SDE steps keep the shared stochastic update.""" + batch = self._build_batch(debug_mode=False) + batch.rollout_sde_type = "sde" + batch.rollout_sde_step_indices = [1] + pipeline_config = types.SimpleNamespace( + shard_latents_for_sp=lambda batch, latents: (latents, False) + ) + ode_calls = [] + + def ode_transition(*, model_output, sample, step_index): + ode_calls.append(step_index) + return sample * 2 + + current_rollout_controller = _make_rollout_controller( + batch, pipeline_config, ode_transition=ode_transition + ) + sample = torch.randn(2, 4, 8, 8) + model_output = torch.randn_like(sample) + + # step 0 is outside the window: the model's update, zero log-prob + prev_0 = current_rollout_controller.flow_sde_sampling( + model_output=model_output, + sample=sample, + current_sigma=torch.tensor(0.6), + next_sigma=torch.tensor(0.4), + generator=None, + step_index=0, + ) + self.assertTrue(torch.equal(prev_0, sample * 2)) + # step 1 is inside the window: stochastic update, ode_transition not called + current_rollout_controller.flow_sde_sampling( + model_output=model_output, + sample=sample, + current_sigma=torch.tensor(0.6), + next_sigma=torch.tensor(0.4), + generator=[torch.Generator().manual_seed(i) for i in range(2)], + step_index=1, + ) + self.assertEqual(ode_calls, [0]) + log_prob_sum, _ = current_rollout_controller.consume_local_rollout_log_probs() + self.assertTrue(torch.equal(log_prob_sum[:, 0], torch.zeros(2))) + def test_ode_debug_tensors_have_shape_safe_noise_std(self): batch = self._build_batch(debug_mode=True) current_rollout_controller = _make_rollout_controller(batch) From 8dc3ba55ee6880db94b73378b92be5992e964ead Mon Sep 17 00:00:00 2001 From: rockdu Date: Thu, 15 Oct 2026 21:23:05 -0700 Subject: [PATCH 12/15] docs(diffusion): mark the MiniMax H3 Diffusers converter as cloned by miles-diffusion Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 01002aefd6ee..644f4419f5e4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -132,6 +132,9 @@ def _reject_non_lora_delta_tensors(adapter: dict[str, torch.Tensor]) -> None: ) +# DO NOT CHANGE, MILES-D IS USING: miles_diffusion keeps a clone in +# miles/backends/fsdp_utils/models/diffusers/h3/weight_bridge.py; change both together. +# TODO(rockdu): load weights and update weights in one format so the clone can go. def _diffusers_h3_checkpoint( iterator: Iterable[tuple[str, torch.Tensor]], ) -> Iterator[tuple[str, torch.Tensor]]: From 0e57e33c11f43f1e10f2765adcbf2274fe5fb739 Mon Sep 17 00:00:00 2001 From: rockdu Date: Fri, 16 Oct 2026 21:49:20 -0700 Subject: [PATCH 13/15] fix(diffusion): pass rollout step kwargs only to rollout requests and keep trajectory timesteps on device Refactor-Kind: non_mechanical_provable Co-authored-by: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Co-authored-by: JingwenGu0829 <75733630+JingwenGu0829@users.noreply.github.com> Signed-off-by: rockdu --- .../pipelines_core/stages/denoising.py | 6 +-- .../stages/model_specific_stages/cosmos3.py | 7 ++- .../model_specific_stages/ltx_2/denoising.py | 4 +- .../post_training/flow_rollout_controller.py | 2 +- .../runtime/post_training/rollout_hooks.py | 18 ++++---- .../test/unit/test_cosmos3_rollout.py | 43 +++++++++++++++++++ .../test/unit/test_scheduler_rollout_unit.py | 15 +++++++ 7 files changed, 75 insertions(+), 20 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 6017cd55aef7..fddc52999704 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -158,7 +158,7 @@ finalize_rollout_if_enabled, init_rollout_env_if_enabled, prepare_rollout_if_enabled, - rollout_transition, + rollout_step_kwargs, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -1343,7 +1343,7 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): { "generator": batch.generator, "eta": batch.eta, - "transition": rollout_transition(batch), + **rollout_step_kwargs(batch), }, ) @@ -2124,7 +2124,7 @@ def _denoise( batch=batch, latents=ctx.latents, num_inference_steps=len(timesteps_cpu), - final_timestep=timesteps_cpu.new_zeros(()), + final_timestep=ctx.timesteps.new_zeros(()), pipeline_config=server_args.pipeline_config, ) self._finalize_denoising_loop(ctx, batch, server_args) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 0bb0073e0ba0..33c62f3f62c6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -63,7 +63,7 @@ finalize_rollout_if_enabled, init_rollout_env_if_enabled, prepare_rollout_if_enabled, - rollout_transition, + rollout_step_kwargs, ) from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import ( prepare_rollout_request_scheduler, @@ -1520,7 +1520,6 @@ def _denoise_once( }, guidance=None, ) - transition = rollout_transition(batch) do_cfg = guidance_scale > 1.0 # Control-CFG runs even when text guidance is off (its own extra @@ -1786,7 +1785,7 @@ def _denoise_once( t, latents, generator=generator, - transition=transition, + **rollout_step_kwargs(batch), return_dict=False, )[0] @@ -1841,7 +1840,7 @@ def _denoise_once( batch=batch, latents=latents, num_inference_steps=len(timesteps), - final_timestep=timesteps.new_zeros(()).cpu(), + final_timestep=timesteps.new_zeros(()), pipeline_config=server_args.pipeline_config, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index f348070628e6..358d035ed63c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -41,7 +41,7 @@ current_platform, ) from sglang.multimodal_gen.runtime.post_training.rollout_hooks import ( - rollout_transition, + rollout_step_kwargs, ) from sglang.multimodal_gen.runtime.server_args import ( ServerArgs, @@ -1902,7 +1902,7 @@ def _stage2_midpoint_model_call( model_video, step.t_device, ctx.latents, - transition=rollout_transition(batch), + **rollout_step_kwargs(batch), return_dict=False, )[0] ctx.audio_latents = ctx.audio_scheduler.step( diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 3dd571305141..1cf115b1d0c6 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -402,4 +402,4 @@ def _maybe_append_dit_trajectory_step( return self.step_latents.append(latents.detach()) - self.step_timesteps.append(timestep_value.detach().cpu()) + self.step_timesteps.append(timestep_value.detach()) diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py index 9cf41b82d850..245c6f4606d8 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py @@ -4,7 +4,7 @@ A loop calls ``prepare_rollout_if_enabled`` before SP sharding, ``init_rollout_env_if_enabled`` at the loop start and ``finalize_rollout_if_enabled`` at the loop end, and passes -``rollout_transition(batch)`` to ``scheduler.step``. Every hook is a no-op for +``**rollout_step_kwargs(batch)`` to ``scheduler.step``. Every hook is a no-op for serving requests. """ @@ -30,7 +30,6 @@ gather_stacked_latents_for_sp, ) - # Key of the image/video latent stream in ``Req.rollout_controllers``. VISUAL_STREAM = "visual" @@ -47,15 +46,14 @@ def _kwargs_to_cpu(d: Any) -> Any: return d -def rollout_transition(batch: Req) -> Callable[..., torch.FloatTensor] | None: - """The ``transition`` to pass to ``scheduler.step`` for this request: the - visual-stream rollout controller, or None to keep the scheduler's own update.""" +def rollout_step_kwargs(batch: Req) -> dict[str, Any]: + """The rollout-only keyword arguments of ``scheduler.step`` for this request: + the visual-stream controller's ``transition``. Empty for a serving request, so + its step call is unchanged and works with any scheduler.""" current_rollout_controller = batch.rollout_controllers.get(VISUAL_STREAM) - return ( - None - if current_rollout_controller is None - else current_rollout_controller.transition - ) + if current_rollout_controller is None: + return {} + return {"transition": current_rollout_controller.transition} def _build_scheduler_rollout_controller( diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py index ffd1788b732f..8793dc98a102 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py @@ -4,15 +4,22 @@ import types import unittest +from contextlib import nullcontext +from unittest import mock import torch +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, ) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages import ( + cosmos3, +) from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import ( prepare_rollout_request_scheduler, ) @@ -105,6 +112,42 @@ def test_rl_capable_serving_scheduler_passes_through(self): torch.testing.assert_close(serving.sigmas, expected) +class TestCosmos3ServingStep(unittest.TestCase): + def test_serving_request_steps_unipc_without_rollout_kwargs(self): + """A serving request runs Cosmos3's denoise loop on the UniPC serving + scheduler, whose step() takes no ``transition``: the loop must not pass + rollout-only keyword arguments to it.""" + stage = cosmos3.Cosmos3DenoisingStage.__new__(cosmos3.Cosmos3DenoisingStage) + stage.scheduler = _serving_scheduler() + stage.transformer = mock.Mock() + stage._logged_parallel_config = True + stage.log_info = mock.Mock() + stage.progress_bar = lambda iterable, **kwargs: iterable + stage.use_declared_component = lambda **kwargs: nullcontext() + # Zero velocity: every UniPC step leaves the latents unchanged. + stage._run_transformer = lambda **kwargs: torch.zeros_like(kwargs["latents"]) + text = torch.ones(1, 2, dtype=torch.long) + batch = Req( + sampling_params=SamplingParams(prompt="test", guidance_scale=1.0), + latents=torch.ones(1, 1, 1, 1, 1), + timesteps=stage.scheduler.timesteps, + extra={ + "cond_text_ids": text, + "cond_text_mask": text, + "uncond_text_ids": text, + "uncond_text_mask": text, + "cond_text_seq_len": 2, + "video_shape": (1, 1, 1), + }, + ) + server_args = types.SimpleNamespace( + pipeline_config=None, enable_cfg_parallel=False + ) + with mock.patch.object(cosmos3, "get_sp_world_size", return_value=1): + stage.forward(batch, server_args) + torch.testing.assert_close(batch.latents, torch.ones_like(batch.latents)) + + class _FusedParamModule(torch.nn.Module): """Diffusers-style q/k/v weights that map into one fused to_qkv param.""" diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 502714ce2049..67d443df7499 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -77,6 +77,21 @@ def _build_batch(self, *, debug_mode: bool) -> types.SimpleNamespace: latents=torch.zeros(2, 4, 8, 8, dtype=torch.float32), ) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_trajectory_timestep_stays_on_device(self): + """The transition records the step's timestep without a per-step copy to + the host; finalization stacks them and moves them to the CPU once.""" + batch = self._build_batch(debug_mode=False) + batch.rollout_return_dit_trajectory = True + current_rollout_controller = _make_rollout_controller(batch) + timestep = torch.tensor(500.0, device="cuda") + current_rollout_controller._maybe_append_dit_trajectory_step( + latents=batch.latents, timestep_value=timestep, step_index=0 + ) + self.assertEqual( + current_rollout_controller.step_timesteps[0].device, timestep.device + ) + def test_ode_step_does_not_call_variance_noise_sampler(self): batch = self._build_batch(debug_mode=False) current_rollout_controller = _make_rollout_controller(batch) From 62a2a17d83bf84fca75c3b7304d7a729f76d6202 Mon Sep 17 00:00:00 2001 From: rockdu Date: Fri, 16 Oct 2026 22:06:02 -0700 Subject: [PATCH 14/15] fix(diffusion): pass the request to the rollout controller per call instead of capturing it Refactor-Kind: non_mechanical_provable Co-authored-by: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Co-authored-by: JingwenGu0829 <75733630+JingwenGu0829@users.noreply.github.com> Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 32 ++++++---- .../runtime/post_training/rollout_hooks.py | 15 +++-- .../test/unit/test_scheduler_rollout_unit.py | 62 ++++++++++++++++--- 3 files changed, 85 insertions(+), 24 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 1cf115b1d0c6..025d15d90acf 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -33,8 +33,8 @@ def __init__( generator: Union[torch.Generator, list[torch.Generator]], return_dit_trajectory: bool, return_step_indices: list[int] | None, - shard_for_sp: Callable[[torch.Tensor], torch.Tensor], - gather_for_sp: Callable[[torch.Tensor], torch.Tensor], + shard_for_sp: Callable[[Any, torch.Tensor], torch.Tensor], + gather_for_sp: Callable[[Any, torch.Tensor], torch.Tensor], sigmas: torch.Tensor, latents_shape: tuple, ode_transition: Callable[..., torch.Tensor] | None = None, @@ -47,7 +47,8 @@ def __init__( self.generator = generator self.return_dit_trajectory = return_dit_trajectory self.return_step_indices = return_step_indices - # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. + # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. Both + # read the request's live SP state, so the request is passed in per call. self.shard_for_sp = shard_for_sp self.gather_for_sp = gather_for_sp # [T+1] sigma grid of the loop, terminal 0 included. @@ -104,7 +105,7 @@ def consume_local_rollout_debug_tensors( self._reset_rollout_debug_tensors() return variance_noises, prev_sample_means, noise_std_devs, model_outputs - def collect_rollout_debug_tensors(self) -> RolloutDebugTensors: + def collect_rollout_debug_tensors(self, batch: Any) -> RolloutDebugTensors: """ Consume rollout debug tensors and merge for all SP ranks. @@ -114,9 +115,9 @@ def collect_rollout_debug_tensors(self) -> RolloutDebugTensors: self.consume_local_rollout_debug_tensors() ) - variance_noises = self.gather_for_sp(variance_noises) - prev_sample_means = self.gather_for_sp(prev_sample_means) - model_outputs = self.gather_for_sp(model_outputs) + variance_noises = self.gather_for_sp(batch, variance_noises) + prev_sample_means = self.gather_for_sp(batch, prev_sample_means) + model_outputs = self.gather_for_sp(batch, model_outputs) # noise_std_devs is same on every device, not a sharded latent tensor. return RolloutDebugTensors( @@ -129,6 +130,7 @@ def collect_rollout_debug_tensors(self) -> RolloutDebugTensors: def transition( self, *, + batch: Any, model_output: torch.FloatTensor, sample: torch.FloatTensor, current_sigma: torch.FloatTensor, @@ -138,7 +140,8 @@ def transition( step_index: int, timestep: Any, ) -> torch.FloatTensor: - """Stand-in for ``scheduler._transition``, passed as ``step(transition=...)``; + """Stand-in for ``scheduler._transition``, passed as ``step(transition=...)`` + with ``batch`` bound by ``rollout_step_kwargs``; ``sample`` arrives already upcast to fp32 by ``step()``. Draws variance noise from the request generator captured at creation, so @@ -151,6 +154,7 @@ def transition( latents=sample, timestep_value=timestep, step_index=step_index ) return self.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -179,6 +183,7 @@ def _get_or_create_rollout_noise_buffer( def _rollout_variance_noise( self, + batch: Any, model_output: torch.FloatTensor, generator: Union[torch.Generator, list[torch.Generator]], ) -> torch.FloatTensor: @@ -208,7 +213,7 @@ def _rollout_variance_noise( generator=generator[i], ) - sharded_noise = self.shard_for_sp(buffer) + sharded_noise = self.shard_for_sp(batch, buffer) if tuple(sharded_noise.shape) != local_shape: raise ValueError( "Rollout SP noise shape mismatch after shard. " @@ -218,6 +223,7 @@ def _rollout_variance_noise( def flow_sde_sampling( self, + batch: Any, model_output: torch.FloatTensor, sample: torch.FloatTensor, current_sigma: torch.FloatTensor, @@ -258,7 +264,9 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise(model_output, generator) + variance_noise = self._rollout_variance_noise( + batch, model_output, generator + ) full_variance_noise = self.noise_buffer std_dev_t = ( torch.sqrt( @@ -289,7 +297,9 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise(model_output, generator) + variance_noise = self._rollout_variance_noise( + batch, model_output, generator + ) full_variance_noise = self.noise_buffer std_dev_t = next_sigma * math.sin(self.noise_level * math.pi / 2) noise_std_dev = std_dev_t diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py index 245c6f4606d8..c7cf99ea8574 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py @@ -10,6 +10,7 @@ from __future__ import annotations +import functools from typing import Any, Callable import torch @@ -53,7 +54,11 @@ def rollout_step_kwargs(batch: Req) -> dict[str, Any]: current_rollout_controller = batch.rollout_controllers.get(VISUAL_STREAM) if current_rollout_controller is None: return {} - return {"transition": current_rollout_controller.transition} + return { + "transition": functools.partial( + current_rollout_controller.transition, batch=batch + ) + } def _build_scheduler_rollout_controller( @@ -73,10 +78,10 @@ def _build_scheduler_rollout_controller( generator=batch.generator, return_dit_trajectory=batch.rollout_return_dit_trajectory, return_step_indices=batch.rollout_return_step_indices, - shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + shard_for_sp=lambda batch, latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], - gather_for_sp=lambda stacked_latents: gather_stacked_latents_for_sp( + gather_for_sp=lambda batch, stacked_latents: gather_stacked_latents_for_sp( pipeline_config=pipeline_config, batch=batch, stacked_latents=stacked_latents, @@ -187,7 +192,7 @@ def _collect_rollout_log_probs(batch: Req) -> None: ) if batch.rollout_debug_mode: batch.rollout_trajectory_data.rollout_debug_tensors = ( - current_rollout_controller.collect_rollout_debug_tensors() + current_rollout_controller.collect_rollout_debug_tensors(batch) ) @@ -203,7 +208,7 @@ def _finalize_rollout_env(*, batch: Req, pipeline_config: Any) -> None: if step_latents: step_latents_tensor = current_rollout_controller.gather_for_sp( - torch.stack(step_latents, dim=1) + batch, torch.stack(step_latents, dim=1) ) batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( latents=step_latents_tensor.cpu(), diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 67d443df7499..b04907cf235e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -1,7 +1,7 @@ """Unit tests for FlowRolloutController, whose ``transition`` the rollout passes to ``scheduler.step()``. - scheduler.step(model_output, t, x_t, transition=current_rollout_controller.transition) + scheduler.step(model_output, t, x_t, **rollout_step_kwargs(batch)) | upcast x_t to fp32, look up sigma_t / sigma_next, dt = sigma_next - sigma_t v current_rollout_controller.transition(...) @@ -25,18 +25,26 @@ return_step_indices filters latents while timesteps cover every step. """ +import functools +import gc import math import types import unittest +import weakref import torch +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( FlowRolloutController, ) +from sglang.multimodal_gen.runtime.post_training.rollout_hooks import ( + prepare_rollout_if_enabled, +) _SIGMAS = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0], dtype=torch.float32) @@ -54,16 +62,42 @@ def _make_rollout_controller( generator=torch.Generator().manual_seed(0), return_dit_trajectory=batch.rollout_return_dit_trajectory, return_step_indices=batch.rollout_return_step_indices, - shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + shard_for_sp=lambda batch, latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], - gather_for_sp=lambda stacked_latents: stacked_latents, + gather_for_sp=lambda batch, stacked_latents: stacked_latents, sigmas=_SIGMAS, latents_shape=tuple(batch.latents.shape), ode_transition=ode_transition, ) +class TestRolloutControllerOwnership(unittest.TestCase): + def test_controller_does_not_keep_the_request_alive(self): + """batch -> rollout controller must not point back to the batch: a request + dropped mid-denoise (e.g. on an exception) is freed by refcounting, without + waiting for a gc pass. + + batch --rollout_controllers--> controller --x--> batch + """ + batch = Req(sampling_params=SamplingParams(prompt="test", rollout=True)) + batch.scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0) + batch.scheduler.set_timesteps(5) + batch.latents = torch.zeros(1, 4, 8, 8) + prepare_rollout_if_enabled(batch=batch, pipeline_config=types.SimpleNamespace()) + self.assertTrue(batch.rollout_controllers) + + batch_ref = weakref.ref(batch) + gc.disable() + try: + del batch + self.assertTrue( + batch_ref() is None, "the rollout controller keeps the request alive" + ) + finally: + gc.enable() + + class TestSchedulerRolloutOdeUnit(unittest.TestCase): def _build_batch(self, *, debug_mode: bool) -> types.SimpleNamespace: return types.SimpleNamespace( @@ -105,6 +139,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) prev_sample = current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -147,6 +182,7 @@ def test_ode_bit_exact_with_non_rollout_path(self): dt = next_sigma - current_sigma rollout_prev = current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -190,9 +226,8 @@ def test_ode_transition_matches_serving_step(self): for scheduler in (serving, rollout_scheduler): scheduler.set_timesteps(5) scheduler.set_begin_index(0) - current_rollout_controller = _make_rollout_controller( - self._build_batch(debug_mode=False) - ) + batch = self._build_batch(debug_mode=False) + current_rollout_controller = _make_rollout_controller(batch) g = torch.Generator().manual_seed(0) sample = torch.randn(2, 4, 8, 8) @@ -204,7 +239,9 @@ def test_ode_transition_matches_serving_step(self): model_output, t, actual, - transition=current_rollout_controller.transition, + transition=functools.partial( + current_rollout_controller.transition, batch=batch + ), return_dict=False, )[0] self.assertTrue(torch.equal(actual, expected)) @@ -228,7 +265,9 @@ def test_transition_records_pre_step_latents(self): model_output, scheduler.timesteps[0], sample, - transition=current_rollout_controller.transition, + transition=functools.partial( + current_rollout_controller.transition, batch=batch + ), return_dict=False, )[0] self.assertTrue(torch.equal(current_rollout_controller.step_latents[0], sample)) @@ -259,6 +298,7 @@ def ode_transition(*, model_output, sample, step_index): # step 0 is outside the window: the model's update, zero log-prob prev_0 = current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=torch.tensor(0.6), @@ -269,6 +309,7 @@ def ode_transition(*, model_output, sample, step_index): self.assertTrue(torch.equal(prev_0, sample * 2)) # step 1 is inside the window: stochastic update, ode_transition not called current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=torch.tensor(0.6), @@ -290,6 +331,7 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -427,6 +469,7 @@ def test_single_step_matches_flowgrpo_reference(self): ) prev_sgl = current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -515,6 +558,7 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): # noise buffer. If flow_sde_sampling fails to cast to fp32 at # entry, the buffer is bf16 → log_prob becomes bf16. current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -570,6 +614,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): # Step 0: not in filter → deterministic ODE transition, no noise draw. prev_0 = current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -583,6 +628,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): # Step 1: in filter → real SDE, noise drawn, prev differs from ODE form. prev_1 = current_rollout_controller.flow_sde_sampling( + batch=batch, model_output=model_output, sample=sample, current_sigma=current_sigma, From b19c15b427fe87ab989425db042f9fe51ba16755 Mon Sep 17 00:00:00 2001 From: rockdu Date: Fri, 16 Oct 2026 22:19:18 -0700 Subject: [PATCH 15/15] refactor(diffusion): merge the rollout finalize helpers, drop the redundant generator argument and make the trajectory recorder public Refactor-Kind: non_mechanical_provable Co-authored-by: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Co-authored-by: JingwenGu0829 <75733630+JingwenGu0829@users.noreply.github.com> Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 29 +++++------ .../runtime/post_training/rollout_hooks.py | 50 +++++++------------ .../test/unit/test_scheduler_rollout_unit.py | 41 +++++++-------- 3 files changed, 50 insertions(+), 70 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 025d15d90acf..2aaf12527381 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -150,7 +150,7 @@ def transition( """ if current_sigma.dim() != 0: raise ValueError("Rollout does not support per-token timesteps") - self._maybe_append_dit_trajectory_step( + self.record_trajectory_step( latents=sample, timestep_value=timestep, step_index=step_index ) return self.flow_sde_sampling( @@ -159,7 +159,6 @@ def transition( sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=self.generator, step_index=step_index, ) @@ -185,23 +184,24 @@ def _rollout_variance_noise( self, batch: Any, model_output: torch.FloatTensor, - generator: Union[torch.Generator, list[torch.Generator]], ) -> torch.FloatTensor: - """Generate variance noise for rollout. If generator is a list, use generator[i] for the i-th batch item.""" - assert generator is not None, "Generator must be provided" + """Generate variance noise for rollout from the request generator; a list + of generators gives one per batch item.""" + assert self.generator is not None, "Generator must be provided" device = model_output.device dtype = model_output.dtype local_shape = tuple(model_output.shape) B = local_shape[0] - if isinstance(generator, torch.Generator): + if isinstance(self.generator, torch.Generator): assert B == 1, "Generator must be a list if batch size is not 1" - generator = [generator] + generators = [self.generator] else: - assert len(generator) == B, ( + assert len(self.generator) == B, ( "Generator list must have the same length as batch size" ) + generators = self.generator buffer = self._get_or_create_rollout_noise_buffer( self.latents_shape, device, dtype @@ -210,7 +210,7 @@ def _rollout_variance_noise( torch.randn( self.latents_shape, out=buffer[i : i + 1], - generator=generator[i], + generator=generators[i], ) sharded_noise = self.shard_for_sp(batch, buffer) @@ -228,7 +228,6 @@ def flow_sde_sampling( sample: torch.FloatTensor, current_sigma: torch.FloatTensor, next_sigma: torch.FloatTensor, - generator: torch.Generator, step_index: int, ) -> torch.Tensor: """Flow rollout step for log-prob / sampling (see FlowGRPO-style references). @@ -264,9 +263,7 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator - ) + variance_noise = self._rollout_variance_noise(batch, model_output) full_variance_noise = self.noise_buffer std_dev_t = ( torch.sqrt( @@ -297,9 +294,7 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator - ) + variance_noise = self._rollout_variance_noise(batch, model_output) full_variance_noise = self.noise_buffer std_dev_t = next_sigma * math.sin(self.noise_level * math.pi / 2) noise_std_dev = std_dev_t @@ -396,7 +391,7 @@ def collect_rollout_log_probs(self) -> torch.Tensor: rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() - def _maybe_append_dit_trajectory_step( + def record_trajectory_step( self, latents: torch.Tensor, timestep_value: torch.Tensor, diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py index c7cf99ea8574..ca3f9f17b570 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py @@ -170,20 +170,6 @@ def finalize_rollout_if_enabled( """ if not batch.rollout: return - _collect_rollout_log_probs(batch) - # Append final denoised latent as the (T+1)-th entry (step_index=T), - # routed through the same filter so rollout_return_step_indices can - # include/exclude it. - batch.rollout_controllers[VISUAL_STREAM]._maybe_append_dit_trajectory_step( - latents=latents, - timestep_value=final_timestep, - step_index=num_inference_steps, - ) - _finalize_rollout_env(batch=batch, pipeline_config=pipeline_config) - batch.rollout_controllers = {} - - -def _collect_rollout_log_probs(batch: Req) -> None: current_rollout_controller = batch.rollout_controllers[VISUAL_STREAM] if batch.rollout_trajectory_data is None: batch.rollout_trajectory_data = RolloutTrajectoryData() @@ -195,28 +181,29 @@ def _collect_rollout_log_probs(batch: Req) -> None: current_rollout_controller.collect_rollout_debug_tensors(batch) ) - -def _finalize_rollout_env(*, batch: Req, pipeline_config: Any) -> None: - state = batch._rollout_denoising_env_state - if state is None: - return - - env: RolloutDenoisingEnv | None = state["env"] - current_rollout_controller = batch.rollout_controllers[VISUAL_STREAM] - step_latents: list[torch.Tensor] = current_rollout_controller.step_latents - step_timesteps: list[torch.Tensor] = current_rollout_controller.step_timesteps - + # Append final denoised latent as the (T+1)-th entry (step_index=T), + # routed through the same filter so rollout_return_step_indices can + # include/exclude it. + current_rollout_controller.record_trajectory_step( + latents=latents, + timestep_value=final_timestep, + step_index=num_inference_steps, + ) + step_latents = current_rollout_controller.step_latents if step_latents: - step_latents_tensor = current_rollout_controller.gather_for_sp( - batch, torch.stack(step_latents, dim=1) - ) batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( - latents=step_latents_tensor.cpu(), - timesteps=torch.stack(step_timesteps, dim=0).cpu(), + latents=current_rollout_controller.gather_for_sp( + batch, torch.stack(step_latents, dim=1) + ).cpu(), + timesteps=torch.stack( + current_rollout_controller.step_timesteps, dim=0 + ).cpu(), sigmas=current_rollout_controller.sigmas.detach().cpu().clone(), ) - if env is not None: + state = batch._rollout_denoising_env_state + if state is not None and state["env"] is not None: + env: RolloutDenoisingEnv = state["env"] env.pos_cond_kwargs = _kwargs_to_cpu( pipeline_config.gather_denoising_env_static_for_sp( batch, state["pos_cond_kwargs_src"] @@ -231,3 +218,4 @@ def _finalize_rollout_env(*, batch: Req, pipeline_config: Any) -> None: batch.rollout_trajectory_data.denoising_env = env batch._rollout_denoising_env_state = None + batch.rollout_controllers = {} diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index b04907cf235e..c264790b4f70 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -50,7 +50,7 @@ def _make_rollout_controller( - batch, pipeline_config=None, ode_transition=None + batch, pipeline_config=None, ode_transition=None, generator=None ) -> FlowRolloutController: """The controller the denoising stage creates for ``batch`` (single GPU).""" return FlowRolloutController( @@ -59,7 +59,7 @@ def _make_rollout_controller( log_prob_no_const=batch.rollout_log_prob_no_const, debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, - generator=torch.Generator().manual_seed(0), + generator=torch.Generator().manual_seed(0) if generator is None else generator, return_dit_trajectory=batch.rollout_return_dit_trajectory, return_step_indices=batch.rollout_return_step_indices, shard_for_sp=lambda batch, latents: pipeline_config.shard_latents_for_sp( @@ -119,7 +119,7 @@ def test_trajectory_timestep_stays_on_device(self): batch.rollout_return_dit_trajectory = True current_rollout_controller = _make_rollout_controller(batch) timestep = torch.tensor(500.0, device="cuda") - current_rollout_controller._maybe_append_dit_trajectory_step( + current_rollout_controller.record_trajectory_step( latents=batch.latents, timestep_value=timestep, step_index=0 ) self.assertEqual( @@ -128,10 +128,11 @@ def test_trajectory_timestep_stays_on_device(self): def test_ode_step_does_not_call_variance_noise_sampler(self): batch = self._build_batch(debug_mode=False) - current_rollout_controller = _make_rollout_controller(batch) - generator = torch.Generator().manual_seed(1) generator_state = generator.get_state() + current_rollout_controller = _make_rollout_controller( + batch, generator=generator + ) sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) model_output = torch.randn_like(sample) @@ -144,7 +145,6 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=generator, step_index=2, ) # The ODE path must not draw variance noise from the request generator. @@ -187,7 +187,6 @@ def test_ode_bit_exact_with_non_rollout_path(self): sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=torch.Generator(device=sample.device).manual_seed(0), step_index=2, ) # Exact expression used by the non-rollout branch at @@ -291,7 +290,10 @@ def ode_transition(*, model_output, sample, step_index): return sample * 2 current_rollout_controller = _make_rollout_controller( - batch, pipeline_config, ode_transition=ode_transition + batch, + pipeline_config, + ode_transition=ode_transition, + generator=[torch.Generator().manual_seed(i) for i in range(2)], ) sample = torch.randn(2, 4, 8, 8) model_output = torch.randn_like(sample) @@ -303,7 +305,6 @@ def ode_transition(*, model_output, sample, step_index): sample=sample, current_sigma=torch.tensor(0.6), next_sigma=torch.tensor(0.4), - generator=None, step_index=0, ) self.assertTrue(torch.equal(prev_0, sample * 2)) @@ -314,7 +315,6 @@ def ode_transition(*, model_output, sample, step_index): sample=sample, current_sigma=torch.tensor(0.6), next_sigma=torch.tensor(0.4), - generator=[torch.Generator().manual_seed(i) for i in range(2)], step_index=1, ) self.assertEqual(ode_calls, [0]) @@ -336,7 +336,6 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=torch.Generator(device=sample.device).manual_seed(2), step_index=2, ) @@ -452,16 +451,17 @@ def test_single_step_matches_flowgrpo_reference(self): for sde_type in ("sde", "cps"): for seed in (0, 1, 2, 3): batch = self._build_batch(sde_type=sde_type, shape=shape) + # The controller draws its variance noise from the request + # generator; a twin generator reproduces it for the reference. current_rollout_controller = _make_rollout_controller( - batch, pipeline_config + batch, + pipeline_config, + generator=torch.Generator(device="cpu").manual_seed(seed + 100), ) g = torch.Generator(device="cpu").manual_seed(seed) model_output = torch.randn(shape, generator=g, dtype=torch.float32) sample = torch.randn(shape, generator=g, dtype=torch.float32) - # The controller draws its variance noise from the request - # generator; a twin generator reproduces it for the reference. - noise_generator = torch.Generator(device="cpu").manual_seed(seed + 100) variance_noise = torch.randn( shape, generator=torch.Generator(device="cpu").manual_seed(seed + 100), @@ -474,7 +474,6 @@ def test_single_step_matches_flowgrpo_reference(self): sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=noise_generator, step_index=2, ) log_prob_sum, elem_count = ( @@ -563,7 +562,6 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=g, step_index=2, ) log_prob_sum, _count = ( @@ -602,9 +600,10 @@ def test_timestep_filters_gate_sde_and_trajectory(self): rollout_return_step_indices=None, latents=torch.empty(shape, dtype=torch.float32), ) - current_rollout_controller = _make_rollout_controller(batch, pipeline_config) - g = torch.Generator(device="cpu").manual_seed(0) + current_rollout_controller = _make_rollout_controller( + batch, pipeline_config, generator=g + ) sample = torch.randn(shape, generator=g, dtype=torch.float32) model_output = torch.randn(shape, generator=g, dtype=torch.float32) current_sigma = torch.tensor(0.6, dtype=torch.float32) @@ -619,7 +618,6 @@ def test_timestep_filters_gate_sde_and_trajectory(self): sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=g, step_index=0, ) self.assertTrue(torch.equal(g.get_state(), generator_state)) @@ -633,7 +631,6 @@ def test_timestep_filters_gate_sde_and_trajectory(self): sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=g, step_index=1, ) self.assertFalse(torch.equal(g.get_state(), generator_state)) @@ -674,7 +671,7 @@ def _recorded_steps(return_step_indices: list[int] | None) -> list[int]: ) # Step i records the timestep value i, so the timesteps name the steps kept. for i in range(4): - current_rollout_controller._maybe_append_dit_trajectory_step( + current_rollout_controller.record_trajectory_step( latents=lat, timestep_value=torch.tensor(float(i)), step_index=i ) self.assertEqual(