diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py index 9b3e63a805e7..99af7b33760d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py @@ -8,6 +8,9 @@ using the modular pipeline architecture. """ +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, ) @@ -17,6 +20,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, + TimestepPreparationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.wan import ( WanProgressiveDenoisingStage, @@ -49,10 +53,19 @@ def initialize_pipeline(self, server_args: ServerArgs): ) def create_pipeline_stages(self, server_args: ServerArgs) -> None: + shift = server_args.pipeline_config.flow_shift self.add_stage(InputValidationStage()) self.add_standard_text_encoding_stage() self.add_standard_latent_preparation_stage() - self.add_standard_timestep_preparation_stage() + # rollout=True requests lazily bind a flow-match Euler scheduler (RL SDE path). + self.add_stage( + TimestepPreparationStage( + scheduler=self.get_module("scheduler"), + rollout_scheduler_factory=lambda: FlowMatchEulerDiscreteScheduler( + shift=1.0 if shift is None else shift + ), + ) + ) self.add_progressive_denoising_stage(WanProgressiveDenoisingStage) self.add_standard_decoding_stage() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py index ea581f9b841f..24ebc65c543a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py @@ -28,6 +28,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) +from sglang.multimodal_gen.runtime.post_training.rollout_timestep_mixin import ( + RolloutTimestepPreparationMixin, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -45,7 +48,7 @@ class TimestepPreparationFingerprint: num_frames: int | None -class TimestepPreparationStage(PipelineStage): +class TimestepPreparationStage(PipelineStage, RolloutTimestepPreparationMixin): """ Stage for preparing timesteps for the diffusion process. @@ -63,9 +66,13 @@ def __init__( prepare_extra_set_timesteps_kwargs: ( list[Callable[[Req, ServerArgs], Tuple[str, Any]]] | None ) = None, + rollout_scheduler_factory: Callable[[], Any] | None = None, ) -> None: super().__init__() self.scheduler = scheduler + # See RolloutTimestepPreparationMixin. + self.rollout_scheduler_factory = rollout_scheduler_factory + self._rollout_scheduler = None self.prepare_extra_set_timesteps_kwargs = list( prepare_extra_set_timesteps_kwargs or [] ) @@ -90,7 +97,11 @@ def forward( if batch.scheduler is not None and batch.timesteps is not None: return batch - scheduler = get_or_create_request_scheduler(batch, self.scheduler) + rollout_template = self._resolve_rollout_scheduler(batch) + scheduler = get_or_create_request_scheduler( + batch, + rollout_template if rollout_template is not None else self.scheduler, + ) device = get_local_torch_device() num_inference_steps = batch.num_inference_steps timesteps = batch.timesteps @@ -153,6 +164,9 @@ def forward( ) timesteps = scheduler.timesteps + if rollout_template is not None: + self._check_rollout_timesteps(scheduler) + # Update batch with prepared timesteps batch.timesteps = timesteps batch.scheduler = scheduler 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 7b6e222ca6d2..5547c40292a2 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 @@ -38,27 +38,38 @@ def _kwargs_to_cpu(d: Any) -> Any: class RolloutDenoisingMixin: + def _request_scheduler(self, batch: Req): + """Scheduler in effect for this request. + + The timestep preparation stage binds per-request schedulers (e.g. the + RL rollout scheduler) to ``batch.scheduler``; the stage module is only + a fallback for pipelines that never bind one. + """ + return batch.scheduler if batch.scheduler is not None else self.scheduler + def _maybe_prepare_rollout(self, batch: Req): """Prepare denoising loop for rollout.""" - if not isinstance(self.scheduler, SchedulerRLMixin): + scheduler = self._request_scheduler(batch) + if not isinstance(scheduler, SchedulerRLMixin): if batch.rollout: raise ValueError( - f"Scheduler {type(self.scheduler)} does not support rollout" + f"Scheduler {type(scheduler)} does not support rollout" ) return - self.scheduler.release_rollout_resources(batch) + scheduler.release_rollout_resources(batch) if batch.rollout: - self.scheduler.prepare_rollout( + scheduler.prepare_rollout( batch=batch, pipeline_config=self.server_args.pipeline_config, ) def _maybe_collect_rollout_log_probs(self, batch: Req): - if not isinstance(self.scheduler, SchedulerRLMixin): + scheduler = self._request_scheduler(batch) + if not isinstance(scheduler, SchedulerRLMixin): if batch.rollout: raise ValueError( - f"Scheduler {type(self.scheduler)} does not support rollout" + f"Scheduler {type(scheduler)} does not support rollout" ) return @@ -66,13 +77,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 = ( - self.scheduler.collect_rollout_log_probs(batch) + scheduler.collect_rollout_log_probs(batch) ) if batch.rollout_debug_mode: batch.rollout_trajectory_data.rollout_debug_tensors = ( - self.scheduler.collect_rollout_debug_tensors(batch) + scheduler.collect_rollout_debug_tensors(batch) ) - self.scheduler.release_rollout_resources(batch) + scheduler.release_rollout_resources(batch) def _postprocess_rollout_outputs( self, diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_timestep_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_timestep_mixin.py new file mode 100644 index 000000000000..ae281bc84905 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_timestep_mixin.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Mixin for per-request rollout scheduler binding in TimestepPreparationStage. + +Kept under post_training to keep the core stage lean; mirrors +RolloutDenoisingMixin on DenoisingStage. +""" + +from __future__ import annotations + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class RolloutTimestepPreparationMixin: + """Bind an alternate scheduler to rollout=True requests. + + The rollout SDE/log-prob path needs a first-order flow-match Euler + scheduler, which not every pipeline serves (e.g. Wan serves UniPC). The + host stage sets ``self.rollout_scheduler_factory``; the scheduler is + created on the first rollout=True request and cached, so an engine that + never sees rollout requests never initializes it. None keeps the serving + scheduler for rollout requests. Downstream stages read the scheduler + from ``batch.scheduler``, so the host stage is the single switch point. + """ + + # Class-level so the rollout info log prints once per process, not once + # per stage instance. + _logged_rollout_scheduler_check = False + + def _resolve_rollout_scheduler(self, batch: Req): + """Return the rollout scheduler template for this request, or None.""" + if not batch.rollout or self.rollout_scheduler_factory is None: + return None + if self._rollout_scheduler is None: + self._rollout_scheduler = self.rollout_scheduler_factory() + return self._rollout_scheduler + + def _check_rollout_timesteps(self, scheduler) -> None: + # The rollout SDE/log-prob math assumes the flow-match Euler + # convention timesteps == sigmas[:-1] * num_train_timesteps. + sigmas = scheduler.sigmas + timesteps = scheduler.timesteps + if sigmas is None or timesteps is None or sigmas.numel() < 2: + return + reconstructed = sigmas[:-1].to(device=timesteps.device) * float( + scheduler.config.num_train_timesteps + ) + max_abs_diff = (timesteps.float() - reconstructed.float()).abs().max().item() + if max_abs_diff > 1e-3: + raise ValueError( + f"rollout timestep/sigma mismatch: max_abs_diff={max_abs_diff:.6g}" + ) + if not RolloutTimestepPreparationMixin._logged_rollout_scheduler_check: + logger.info( + "RL rollout using %s (timesteps dtype=%s, sigmas dtype=%s, " + "max_abs_diff=%.6g)", + type(scheduler).__name__, + timesteps.dtype, + sigmas.dtype, + max_abs_diff, + ) + RolloutTimestepPreparationMixin._logged_rollout_scheduler_check = True