From 395a2201e40352eb65b87d3b7186e53c00d800a4 Mon Sep 17 00:00:00 2001 From: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Date: Thu, 9 Jul 2026 10:24:06 +0800 Subject: [PATCH] [diffusion] rl: enable RL rollout path for LTX-2.3 post-training (#28926) --- .../multimodal_gen/configs/sample/ltx_2.py | 3 +++ .../entrypoints/post_training/rollout_api.py | 4 +++- .../entrypoints/post_training/utils.py | 3 +++ .../runtime/pipelines/ltx_2_pipeline.py | 4 +++- .../ltx_2/decoding_av.py | 1 + .../model_specific_stages/ltx_2/denoising.py | 22 ++++++++++++++++--- 6 files changed, 32 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/sample/ltx_2.py b/python/sglang/multimodal_gen/configs/sample/ltx_2.py index 81233cb27..23989b820 100644 --- a/python/sglang/multimodal_gen/configs/sample/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/sample/ltx_2.py @@ -69,6 +69,9 @@ class LTX23SamplingParams(LTX2SamplingParams): def build_request_extra(self) -> dict[str, Any]: extra = super().build_request_extra() + # RL rollout uses the official CFG path (guidance_scale=1, no guider). + if self.rollout: + return extra extra["ltx2_stage1_guider_params"] = { "video_cfg_scale": self.video_cfg_scale, "video_stg_scale": self.video_stg_scale, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py index 159692d44..32e280896 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py @@ -218,7 +218,9 @@ def _build_response( responses: list[RolloutResponse] = [] for sample_idx in range(batch_size): - out_i = result.output[sample_idx].contiguous() + out_i = result.output[sample_idx] + if isinstance(out_i, torch.Tensor): + out_i = out_i.contiguous() serialized_generated_output = _maybe_serialize(out_i) if not rollout: responses.append( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py index d281cf6bc..f01f3857b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py @@ -5,6 +5,7 @@ from __future__ import annotations import base64 from typing import Any +import numpy as np import torch from safetensors.torch import load, save @@ -28,6 +29,8 @@ def _maybe_serialize(obj: Any) -> Any: "shape": list(obj.shape), "dtype": str(obj.dtype), } + if isinstance(obj, np.ndarray): + return _maybe_serialize(torch.from_numpy(obj)) if isinstance(obj, dict): return {k: _maybe_serialize(v) for k, v in obj.items()} if isinstance(obj, (list, tuple)): diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index ca5ca703a..da6cdd65f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -3,7 +3,6 @@ import os import numpy as np import torch -from diffusers import FlowMatchEulerDiscreteScheduler from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( STAGE_2_DISTILLED_SIGMA_VALUES as _SHARED_STAGE_2_DISTILLED_SIGMA_VALUES, @@ -23,6 +22,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im ComponentUse, ResidencyState, ) +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py index 8a7d6e866..0e0bb5ed7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py @@ -108,6 +108,7 @@ class LTX2AVDecodingStage(DecodingStage): trajectory_latents=batch.trajectory_latents, trajectory_decoded=None, metrics=batch.metrics, + rollout_trajectory_data=batch.rollout_trajectory_data, ) # 2. Decode Audio 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 9c8453cb9..8b213c9bc 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 @@ -137,6 +137,12 @@ class LTX2DenoisingStage(DenoisingStage): ) self.sampler_name = sampler_name + 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 @@ -1873,9 +1879,19 @@ class LTX2DenoisingStage(DenoisingStage): midpoint_model_call=_stage2_midpoint_model_call, ) else: - ctx.latents = ctx.scheduler.step( - model_video, step.t_device, ctx.latents, return_dict=False - )[0] + 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.audio_latents = ctx.audio_scheduler.step( model_audio, step.t_device, ctx.audio_latents, return_dict=False )[0]