From 47ac830c07eedfe0b1b7e36e568de0c6fdd73600 Mon Sep 17 00:00:00 2001 From: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Date: Tue, 14 Apr 2026 19:10:38 -0700 Subject: [PATCH] [diffusion] rl: support standalone rollout api, denoising environment backpass and sp-aligned log-prob for T2I post-training (#22604) Co-authored-by: MikukuOvO Co-authored-by: Claude Opus 4.6 (1M context) --- .../configs/pipeline_configs/base.py | 3 + .../configs/pipeline_configs/qwen_image.py | 5 +- .../configs/pipeline_configs/zimage.py | 5 +- .../pipeline_configs/__init__.py | 14 + .../qwen_image_rollout_pipeline_mixin.py | 27 ++ .../zimage_rollout_pipeline_mixin.py | 27 ++ .../configs/sample/sampling_params.py | 6 + .../runtime/entrypoints/http_server.py | 6 +- .../entrypoints/post_training/io_struct.py | 56 +++ .../entrypoints/post_training/rollout_api.py | 309 +++++++++++++++ .../entrypoints/post_training/utils.py | 45 +++ .../pipelines_core/stages/denoising.py | 80 ++-- .../runtime/post_training/rl_dataclasses.py | 20 +- .../post_training/rollout_denoising_mixin.py | 197 ++++++++++ .../post_training/scheduler_rl_debug_mixin.py | 9 +- .../post_training/scheduler_rl_mixin.py | 45 ++- .../runtime/post_training/sp_utils.py | 61 +++ .../test/unit/test_rollout_api.py | 360 ++++++++++++++++++ .../test/unit/test_scheduler_rollout_unit.py | 106 +++++- 19 files changed, 1305 insertions(+), 76 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/post_training/pipeline_configs/__init__.py create mode 100644 python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py create mode 100644 python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py create mode 100644 python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py create mode 100644 python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py create mode 100644 python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py create mode 100644 python/sglang/multimodal_gen/runtime/post_training/sp_utils.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_rollout_api.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 46086cba7..86febb0e0 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -508,6 +508,9 @@ class PipelineConfig: def _unpad_and_unpack_latents(self, latents, audio_latents, batch, vae, audio_vae): raise NotImplementedError("not yet implemented") + def gather_dit_env_static_for_sp(self, batch, cond_kwargs: dict | None): + return cond_kwargs + @staticmethod def add_cli_args( parser: FlexibleArgumentParser, prefix: str = "" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index bf058ad05..c6ff94aca 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -18,6 +18,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( maybe_unpad_latents, shard_rotary_emb_for_sp, ) +from sglang.multimodal_gen.configs.post_training.pipeline_configs import ( + QwenImageRolloutPipelineMixin, +) from sglang.multimodal_gen.runtime.models.vision_utils import resize from sglang.multimodal_gen.utils import calculate_dimensions @@ -127,7 +130,7 @@ def _pack_latents(latents, batch_size, num_channels_latents, height, width): @dataclass -class QwenImagePipelineConfig(ImagePipelineConfig): +class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig): """Configuration for the QwenImage pipeline.""" should_use_guidance: bool = False diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index de08a27ef..dada61710 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py @@ -15,6 +15,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( ImagePipelineConfig, ModelTaskType, ) +from sglang.multimodal_gen.configs.post_training.pipeline_configs import ( + ZImageRolloutPipelineMixin, +) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_group, get_sp_parallel_rank, @@ -40,7 +43,7 @@ class TransformersModelConfig(EncoderConfig): @dataclass -class ZImagePipelineConfig(ImagePipelineConfig): +class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig): should_use_guidance: bool = False task_type: ModelTaskType = ModelTaskType.T2I dit_config: DiTConfig = field(default_factory=ZImageDitConfig) diff --git a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/__init__.py new file mode 100644 index 000000000..bc8a1d7b5 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rollout / RL hooks mixed into multimodal pipeline configs.""" + +from sglang.multimodal_gen.configs.post_training.pipeline_configs.qwen_image_rollout_pipeline_mixin import ( + QwenImageRolloutPipelineMixin, +) +from sglang.multimodal_gen.configs.post_training.pipeline_configs.zimage_rollout_pipeline_mixin import ( + ZImageRolloutPipelineMixin, +) + +__all__ = [ + "QwenImageRolloutPipelineMixin", + "ZImageRolloutPipelineMixin", +] diff --git a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py new file mode 100644 index 000000000..6682c6e73 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rollout / RL hooks for Qwen-Image pipeline configs.""" + +from __future__ import annotations + +import torch + +from sglang.multimodal_gen.runtime.post_training.sp_utils import ( + all_gather_if_sp_sharded, + maybe_trim_sp_rope_seq_for_batch, +) + + +class QwenImageRolloutPipelineMixin: + + def gather_dit_env_static_for_sp(self, batch, cond_kwargs: dict | None): + if cond_kwargs is None: + return None + out = dict(cond_kwargs) + freqs = out.get("freqs_cis") + if freqs is not None: + img_cache, txt_cache = freqs[0], freqs[1] + if isinstance(img_cache, torch.Tensor) and img_cache.dim() == 2: + img_g = all_gather_if_sp_sharded(batch, img_cache, dim=0) + img_g = maybe_trim_sp_rope_seq_for_batch(batch, img_g) + out["freqs_cis"] = (img_g, txt_cache) + return out diff --git a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py new file mode 100644 index 000000000..89f9ae455 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rollout / RL hooks for Z-Image pipeline configs.""" + +from __future__ import annotations + +import torch + +from sglang.multimodal_gen.runtime.post_training.sp_utils import ( + all_gather_if_sp_sharded, + maybe_trim_sp_rope_seq_for_batch, +) + + +class ZImageRolloutPipelineMixin: + + def gather_dit_env_static_for_sp(self, batch, cond_kwargs: dict | None): + if cond_kwargs is None: + return None + out = dict(cond_kwargs) + freqs = out.get("freqs_cis") + if freqs is not None: + cap_freqs, x_freqs = freqs[0], freqs[1] + if isinstance(x_freqs, torch.Tensor) and x_freqs.dim() >= 2: + x_g = all_gather_if_sp_sharded(batch, x_freqs, dim=0) + x_g = maybe_trim_sp_rope_seq_for_batch(batch, x_g) + out["freqs_cis"] = (cap_freqs, x_g) + return out diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 3fae7d089..48434c209 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -187,6 +187,12 @@ class SamplingParams: ) return_trajectory_latents: bool = False # returns all latents for each timestep return_trajectory_decoded: bool = False # returns decoded latents for each timestep + rollout_return_denoising_env: bool = ( + False # populate ``denoising_env`` (image/pos/neg kwargs, guidance) for RL replay + ) + rollout_return_dit_trajectory: bool = ( + False # per-step noisy latents + final latent + timesteps (RolloutDitTrajectory) + ) # if True, disallow user params to override subclass-defined protected fields no_override_protected_fields: bool = False # whether to adjust num_frames for multi-GPU friendly splitting (default: True) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 303cea786..cc6af08fc 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -16,7 +16,10 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( VertexGenerateReqInput, ) from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params -from sglang.multimodal_gen.runtime.entrypoints.post_training import weights_api +from sglang.multimodal_gen.runtime.entrypoints.post_training import ( + rollout_api, + weights_api, +) from sglang.multimodal_gen.runtime.entrypoints.utils import ( prepare_request, save_outputs, @@ -282,6 +285,7 @@ def create_app(server_args: ServerArgs): app.include_router(video_api.router) app.include_router(mesh_api.router) app.include_router(weights_api.router) + app.include_router(rollout_api.router) app.state.server_args = server_args return app diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py index bda72df12..ece590a30 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py @@ -1,6 +1,11 @@ """Request/response data structures for post-training APIs.""" +from __future__ import annotations + from dataclasses import dataclass +from typing import Any, Optional + +from pydantic import BaseModel @dataclass @@ -17,3 +22,54 @@ class GetWeightsChecksumReqInput: """Compute SHA-256 checksum of loaded module weights for verification.""" module_names: list[str] | None = None + + +class RolloutRequest(BaseModel): + prompt: str + negative_prompt: Optional[str] = None + seed: int = 1024 + generator_device: str = "cuda" + + width: Optional[int] = None + height: Optional[int] = None + num_inference_steps: Optional[int] = None + num_outputs_per_prompt: Optional[int] = None + + guidance_scale: Optional[float] = None + true_cfg_scale: Optional[float] = None + + # video-specific (ignored by image pipelines) + num_frames: Optional[int] = None + fps: Optional[int] = None + + rollout: bool = True + rollout_sde_type: str = "sde" + rollout_noise_level: float = 0.7 + rollout_log_prob_no_const: bool = False + rollout_debug_mode: bool = True + + rollout_return_denoising_env: bool = False + rollout_return_dit_trajectory: bool = False + + image_path: Optional[list[str]] = None + + # suppress verbose per-request logging (also gates peak_memory_mb collection) + suppress_logs: bool = False + + extra_sampling_params: Optional[dict[str, Any]] = None + + +class RolloutResponse(BaseModel): + request_id: str + prompt: str + seed: int + + generated_output: Any = None + + rollout_log_probs: Optional[dict[str, Any]] = None + rollout_debug_tensors: Optional[dict[str, Any]] = None + denoising_env: Optional[dict[str, Any]] = None + dit_trajectory: Optional[dict[str, Any]] = None + + inference_time_s: Optional[float] = None + peak_memory_mb: Optional[float] = None 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 new file mode 100644 index 000000000..29536c3c7 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py @@ -0,0 +1,309 @@ +"""Rollout HTTP API (``POST /rollout/generate``).""" + +from __future__ import annotations + +from typing import Any + +import torch +from fastapi import APIRouter, HTTPException +from fastapi.responses import ORJSONResponse + +from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params +from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + RolloutRequest, + RolloutResponse, +) +from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import ( + _maybe_serialize, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch +from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutDebugTensors, + RolloutDenoisingEnv, + RolloutDitTrajectory, + RolloutTrajectoryData, +) +from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client +from sglang.multimodal_gen.runtime.server_args import get_global_server_args +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) +router = APIRouter(prefix="/rollout", tags=["rollout"]) + + +def _extract_single_sample_tensor(obj: Any, sample_idx: int, batch_size: int) -> Any: + if isinstance(obj, torch.Tensor): + if obj.dim() >= 1 and obj.shape[0] == batch_size: + return obj[sample_idx].contiguous() + return obj + if isinstance(obj, dict): + return { + k: _extract_single_sample_tensor(v, sample_idx, batch_size) + for k, v in obj.items() + } + if isinstance(obj, list): + return [_extract_single_sample_tensor(v, sample_idx, batch_size) for v in obj] + if isinstance(obj, tuple): + return tuple( + _extract_single_sample_tensor(v, sample_idx, batch_size) for v in obj + ) + return obj + + +def _slice_rollout_trajectory_for_sample( + rtd: RolloutTrajectoryData | None, + sample_idx: int, + batch_size: int, +) -> RolloutTrajectoryData | None: + if rtd is None: + return None + log_probs = rtd.rollout_log_probs + if ( + isinstance(log_probs, torch.Tensor) + and log_probs.dim() >= 1 + and log_probs.shape[0] == batch_size + ): + log_probs = log_probs[sample_idx].contiguous() + debug_tensors = None + if rtd.rollout_debug_tensors: + rd = rtd.rollout_debug_tensors + debug_tensors = RolloutDebugTensors( + rollout_variance_noises=_extract_single_sample_tensor( + rd.rollout_variance_noises, sample_idx, batch_size + ), + rollout_prev_sample_means=_extract_single_sample_tensor( + rd.rollout_prev_sample_means, sample_idx, batch_size + ), + rollout_noise_std_devs=_extract_single_sample_tensor( + rd.rollout_noise_std_devs, sample_idx, batch_size + ), + rollout_model_outputs=_extract_single_sample_tensor( + rd.rollout_model_outputs, sample_idx, batch_size + ), + ) + denoising_env = None + if rtd.denoising_env: + env = rtd.denoising_env + denoising_env = RolloutDenoisingEnv( + image_kwargs=( + _extract_single_sample_tensor(env.image_kwargs, sample_idx, batch_size) + if env.image_kwargs + else None + ), + pos_cond_kwargs=( + _extract_single_sample_tensor( + env.pos_cond_kwargs, sample_idx, batch_size + ) + if env.pos_cond_kwargs + else None + ), + neg_cond_kwargs=( + _extract_single_sample_tensor( + env.neg_cond_kwargs, sample_idx, batch_size + ) + if env.neg_cond_kwargs + else None + ), + guidance=( + _extract_single_sample_tensor(env.guidance, sample_idx, batch_size) + if env.guidance is not None + else None + ), + ) + dit_trajectory = None + if rtd.dit_trajectory: + dit = rtd.dit_trajectory + dit_trajectory = RolloutDitTrajectory( + latents=_extract_single_sample_tensor(dit.latents, sample_idx, batch_size), + timesteps=dit.timesteps, + ) + return RolloutTrajectoryData( + rollout_log_probs=log_probs, + rollout_debug_tensors=debug_tensors, + denoising_env=denoising_env, + dit_trajectory=dit_trajectory, + ) + + +def _serialize_rollout_trajectory( + rtd: RolloutTrajectoryData | None, + *, + serialized_dit_timesteps: dict | None = None, +) -> tuple[dict | None, dict | None, dict | None, dict | None]: + """Return order: rollout_log_probs, rollout_debug_tensors, denoising_env, dit_trajectory.""" + if rtd is None: + return None, None, None, None + serialized_log_probs = _maybe_serialize(rtd.rollout_log_probs) + serialized_debug_tensors = None + if rtd.rollout_debug_tensors: + rd = rtd.rollout_debug_tensors + serialized_debug_tensors = { + "rollout_variance_noises": _maybe_serialize(rd.rollout_variance_noises), + "rollout_prev_sample_means": _maybe_serialize(rd.rollout_prev_sample_means), + "rollout_noise_std_devs": _maybe_serialize(rd.rollout_noise_std_devs), + "rollout_model_outputs": _maybe_serialize(rd.rollout_model_outputs), + } + serialized_denoising_env = None + if rtd.denoising_env: + env = rtd.denoising_env + serialized_denoising_env = { + "image_kwargs": ( + _maybe_serialize(env.image_kwargs) if env.image_kwargs else None + ), + "pos_cond_kwargs": ( + _maybe_serialize(env.pos_cond_kwargs) if env.pos_cond_kwargs else None + ), + "neg_cond_kwargs": ( + _maybe_serialize(env.neg_cond_kwargs) if env.neg_cond_kwargs else None + ), + "guidance": ( + _maybe_serialize(env.guidance) if env.guidance is not None else None + ), + } + serialized_dit_trajectory = None + if rtd.dit_trajectory: + dit = rtd.dit_trajectory + serialized_dit_trajectory = { + "latents": ( + _maybe_serialize(dit.latents) if dit.latents is not None else None + ), + "timesteps": serialized_dit_timesteps, + } + return ( + serialized_log_probs, + serialized_debug_tensors, + serialized_denoising_env, + serialized_dit_trajectory, + ) + + +def _build_response( + request_id: str, prompt: str, seed: int, rollout: bool, result: OutputBatch +) -> list[RolloutResponse]: + """ + rollout: bool - set to False when evaluating the model + """ + batch_size = result.output.shape[0] + inference_time_s = ( + result.metrics.total_duration_s + if result.metrics and result.metrics.total_duration_s > 0 + else None + ) + peak_memory_mb = result.peak_memory_mb if result.peak_memory_mb > 0 else None + rollout_trajectory_data = result.rollout_trajectory_data + if rollout: + assert ( + rollout_trajectory_data is not None + ), "rollout_trajectory_data must be present when rollout=True" + + serialized_dit_timesteps = None + if rollout and rollout_trajectory_data and rollout_trajectory_data.dit_trajectory: + serialized_dit_timesteps = _maybe_serialize( + rollout_trajectory_data.dit_trajectory.timesteps + ) + + responses: list[RolloutResponse] = [] + for sample_idx in range(batch_size): + out_i = result.output[sample_idx].contiguous() + serialized_generated_output = _maybe_serialize(out_i) + if not rollout: + responses.append( + RolloutResponse( + request_id=request_id, + prompt=prompt, + seed=seed, + generated_output=serialized_generated_output, + inference_time_s=inference_time_s, + peak_memory_mb=peak_memory_mb, + ) + ) + continue + per_sample_trajectory = _slice_rollout_trajectory_for_sample( + result.rollout_trajectory_data, sample_idx, batch_size + ) + ( + serialized_log_probs, + serialized_debug_tensors, + serialized_denoising_env, + serialized_dit_trajectory, + ) = _serialize_rollout_trajectory( + per_sample_trajectory, + serialized_dit_timesteps=serialized_dit_timesteps, + ) + responses.append( + RolloutResponse( + request_id=request_id, + prompt=prompt, + seed=seed, + generated_output=serialized_generated_output, + rollout_log_probs=serialized_log_probs, + rollout_debug_tensors=serialized_debug_tensors, + denoising_env=serialized_denoising_env, + dit_trajectory=serialized_dit_trajectory, + inference_time_s=inference_time_s, + peak_memory_mb=peak_memory_mb, + ) + ) + return responses + + +@router.post("/generate", response_model=list[RolloutResponse]) +async def rollout_generate(request: RolloutRequest): + request_id = generate_request_id() + server_args = get_global_server_args() + sampling_kwargs: dict = dict( + prompt=request.prompt, + negative_prompt=request.negative_prompt, + seed=request.seed, + generator_device=request.generator_device, + width=request.width, + height=request.height, + num_inference_steps=request.num_inference_steps, + num_outputs_per_prompt=request.num_outputs_per_prompt, + guidance_scale=request.guidance_scale, + true_cfg_scale=request.true_cfg_scale, + num_frames=request.num_frames, + fps=request.fps, + image_path=request.image_path, + rollout=request.rollout, + rollout_sde_type=request.rollout_sde_type, + rollout_noise_level=request.rollout_noise_level, + rollout_log_prob_no_const=request.rollout_log_prob_no_const, + rollout_debug_mode=request.rollout_debug_mode, + rollout_return_denoising_env=request.rollout_return_denoising_env, + rollout_return_dit_trajectory=request.rollout_return_dit_trajectory, + suppress_logs=request.suppress_logs, + save_output=False, + return_trajectory_latents=False, + return_trajectory_decoded=False, + ) + if request.extra_sampling_params: + sampling_kwargs.update(request.extra_sampling_params) + sampling_kwargs["rollout"] = request.rollout + sampling_kwargs = {k: v for k, v in sampling_kwargs.items() if v is not None} + try: + sampling_params = build_sampling_params(request_id, **sampling_kwargs) + except Exception as exc: + raise HTTPException( + status_code=400, detail=f"Invalid sampling params: {exc}" + ) from exc + pipeline_request = prepare_request( + server_args=server_args, sampling_params=sampling_params + ) + try: + output_batch: OutputBatch = await async_scheduler_client.forward( + pipeline_request + ) + except Exception as exc: + logger.error("Rollout generation failed: %s", exc, exc_info=True) + raise HTTPException( + status_code=500, detail=f"Generation failed: {exc}" + ) from exc + if output_batch.error: + raise HTTPException(status_code=500, detail=output_batch.error) + rollout_responses = _build_response( + request_id, request.prompt, request.seed, request.rollout, output_batch + ) + return ORJSONResponse(content=[r.model_dump() for r in rollout_responses]) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py new file mode 100644 index 000000000..d281cf6bc --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py @@ -0,0 +1,45 @@ +"""Tensor serialization for post-training / rollout HTTP responses.""" + +from __future__ import annotations + +import base64 +from typing import Any + +import torch +from safetensors.torch import load, save + + +def tensor_to_base64(t: torch.Tensor) -> str: + t = t.detach().contiguous().cpu() + raw = save({"t": t}) + return base64.b64encode(raw).decode("ascii") + + +def base64_to_tensor(s: str) -> torch.Tensor: + raw = base64.b64decode(s) + return load(raw)["t"] + + +def _maybe_serialize(obj: Any) -> Any: + if isinstance(obj, torch.Tensor): + return { + "__tensor__": True, + "data": tensor_to_base64(obj), + "shape": list(obj.shape), + "dtype": str(obj.dtype), + } + if isinstance(obj, dict): + return {k: _maybe_serialize(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_maybe_serialize(v) for v in obj] + return obj + + +def _maybe_deserialize(obj: Any) -> Any: + if isinstance(obj, dict): + if obj.get("__tensor__"): + return base64_to_tensor(obj["data"]) + return {k: _maybe_deserialize(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_maybe_deserialize(v) for v in obj] + return obj 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 4db8c6816..afd2bcc44 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -82,11 +82,8 @@ from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( - RolloutTrajectoryData, -) -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( - SchedulerRLMixin, +from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( + RolloutDenoisingMixin, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin @@ -147,7 +144,7 @@ class DenoisingStepState: attn_metadata: Any | None -class DenoisingStage(PipelineStage): +class DenoisingStage(PipelineStage, RolloutDenoisingMixin): """ Stage for running the denoising loop in diffusion pipelines. @@ -192,43 +189,6 @@ class DenoisingStage(PipelineStage): self._cached_num_steps = None self._is_warmed_up = False - def _maybe_prepare_rollout(self, batch: Req): - """Prepare denoising loop for rollout.""" - if not isinstance(self.scheduler, SchedulerRLMixin): - if batch.rollout: - raise ValueError( - f"Scheduler {type(self.scheduler)} does not support rollout" - ) - return - - self.scheduler.release_rollout_resources(batch) - if batch.rollout: - self.scheduler.prepare_rollout( - batch=batch, - pipeline_config=self.server_args.pipeline_config, - ) - - def _maybe_collect_rollout_log_probs(self, batch: Req): - """Get rollout log probs and store into batch for reward calculation.""" - if not isinstance(self.scheduler, SchedulerRLMixin): - if batch.rollout: - raise ValueError( - f"Scheduler {type(self.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 = ( - self.scheduler.collect_rollout_log_probs(batch) - ) - if getattr(batch, "rollout_debug_mode", False): - batch.rollout_trajectory_data.rollout_debug_tensors = ( - self.scheduler.collect_rollout_debug_tensors(batch) - ) - self.scheduler.release_rollout_resources(batch) - def _maybe_enable_torch_compile(self, module: object) -> None: """ Compile a module with torch.compile, and enable inductor overlap tweak if available. @@ -901,10 +861,6 @@ class DenoisingStage(PipelineStage): trajectory_tensor = None trajectory_timesteps_tensor = None - # Gather log probs for rollout - if batch.rollout: - self._maybe_collect_rollout_log_probs(batch) - # Gather results if using sequence parallelism latents, trajectory_tensor = self._postprocess_sp_latents( batch, latents, trajectory_tensor @@ -1137,6 +1093,15 @@ class DenoisingStage(PipelineStage): 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, + ) denoising_start_time = time.time() self._before_denoising_loop(ctx, batch, server_args) # to avoid device-sync caused by timestep comparison @@ -1164,6 +1129,17 @@ class DenoisingStage(PipelineStage): 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 batch.rollout: + self._maybe_append_dit_trajectory_step( + batch=batch, + latents=ctx.latents, + timestep_value=step.t_host, + ) self._run_denoising_step(ctx, step, batch, server_args) self._record_trajectory(ctx, step, batch, server_args) @@ -1185,6 +1161,16 @@ class DenoisingStage(PipelineStage): (denoising_end_time - denoising_start_time) / len(ctx.timesteps), ) + # Rollout postprocessing must run BEFORE _finalize_denoising_loop so + # the final scheduler.step output (ctx.latents) is still SP-sharded and + # can be gathered uniformly alongside the per-step dit_trajectory via + # gather_stacked_latents_for_sp. + if batch.rollout: + self._postprocess_rollout_outputs( + batch=batch, + latents=ctx.latents, + server_args=server_args, + ) self._finalize_denoising_loop(ctx, batch, server_args) return batch 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 4ce8b0d94..351c17736 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py @@ -39,8 +39,24 @@ class RolloutDebugTensors: @dataclass -class RolloutTrajectoryData: - """Container for rollout-specific trajectory outputs.""" +class RolloutDenoisingEnv: + image_kwargs: dict[str, Any] | None = None + pos_cond_kwargs: dict[str, Any] | None = None + neg_cond_kwargs: dict[str, Any] | None = None + guidance: torch.Tensor | None = None + +@dataclass +class RolloutDitTrajectory: + # [B, T+1, ...]: per-step noisy latents x_{t_0..t_{T-1}} followed by the + # final denoised latent x_{t_T} (last scheduler.step output). + latents: torch.Tensor | None = None + timesteps: torch.Tensor | None = None # [T] + + +@dataclass +class RolloutTrajectoryData: rollout_log_probs: torch.Tensor | None = None rollout_debug_tensors: RolloutDebugTensors | None = None + denoising_env: RolloutDenoisingEnv | None = None + dit_trajectory: RolloutDitTrajectory | None = 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 new file mode 100644 index 000000000..1673bc046 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -0,0 +1,197 @@ +"""Mixin for rollout-related denoising hooks. + +Moved out of DenoisingStage to keep the core stage lean. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +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, +) +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 + + +class RolloutDenoisingMixin: + + def _maybe_prepare_rollout(self, batch: Req): + """Prepare denoising loop for rollout.""" + if not isinstance(self.scheduler, SchedulerRLMixin): + if batch.rollout: + raise ValueError( + f"Scheduler {type(self.scheduler)} does not support rollout" + ) + return + + self.scheduler.release_rollout_resources(batch) + if batch.rollout: + self.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): + if batch.rollout: + raise ValueError( + f"Scheduler {type(self.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 = ( + self.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) + ) + self.scheduler.release_rollout_resources(batch) + + def _postprocess_rollout_outputs( + self, + batch: Req, + latents: 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 the final denoised latent as the (T+1)-th entry of the + # dit-trajectory latents list. + state = getattr(batch, "_rollout_dit_env_state", None) + if state is not None and batch.rollout and batch.rollout_return_dit_trajectory: + state["step_latents"].append(latents.detach()) + self._maybe_finalize_dit_env_collection( + batch=batch, + pipeline_config=server_args.pipeline_config, + ) + + 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_dit_env_state = None + return + + sanitize = getattr(pipeline_config, "sanitize_dit_env_kwargs", lambda x: x) + if collect_env: + env = RolloutDenoisingEnv( + image_kwargs=_kwargs_to_cpu(sanitize(image_kwargs)), + pos_cond_kwargs=_kwargs_to_cpu(sanitize(pos_cond_kwargs)), + neg_cond_kwargs=( + _kwargs_to_cpu(sanitize(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_dit_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, + ) -> None: + if not batch.rollout or not batch.rollout_return_dit_trajectory: + return + state = getattr(batch, "_rollout_dit_env_state", None) + if state is None: + return + + state["step_latents"].append(latents.detach()) + state["step_timesteps"].append(timestep_value.detach().cpu()) + + def _maybe_finalize_dit_env_collection(self, batch, pipeline_config) -> None: + state = getattr(batch, "_rollout_dit_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"] + + 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, + ) + batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( + latents=step_latents_tensor.cpu(), + timesteps=torch.stack(step_timesteps, dim=0).cpu(), + ) + + if env is not None and batch.rollout_return_denoising_env: + sanitize = getattr(pipeline_config, "sanitize_dit_env_kwargs", lambda x: x) + gather_fn = getattr(pipeline_config, "gather_dit_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(sanitize(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(sanitize(gathered_neg)) + + batch.rollout_trajectory_data.denoising_env = env + + batch._rollout_dit_env_state = None 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 index acdacae71..e6113876a 100644 --- 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 @@ -33,7 +33,8 @@ class SchedulerRLDebugMixin: ) -> None: rollout_session_data = batch._rollout_session_data batch_size = variance_noise.shape[0] - rollout_session_data.local_variance_noises.append(variance_noise) + # 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)) @@ -86,13 +87,13 @@ class SchedulerRLDebugMixin: # Gather on packed tensors first. variance_noises_packed = pipeline_config.gather_latents_for_sp( - variance_noises_packed + variance_noises_packed, batch=batch ) prev_sample_means_packed = pipeline_config.gather_latents_for_sp( - prev_sample_means_packed + prev_sample_means_packed, batch=batch ) model_outputs_packed = pipeline_config.gather_latents_for_sp( - model_outputs_packed + model_outputs_packed, batch=batch ) # Unpack back to [B, T, ...]. diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py index 76ffaec77..f6eadf8db 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py @@ -7,12 +7,8 @@ from typing import Any, Union import torch from sglang.multimodal_gen.runtime.distributed import ( - get_local_torch_device, get_sp_world_size, ) -from sglang.multimodal_gen.runtime.distributed.communication_op import ( - sequence_model_parallel_all_reduce, -) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutSessionData, @@ -107,7 +103,7 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): ) sharded_noise, _ = rollout_session_data.pipeline_config.shard_latents_for_sp( - batch, buffer + batch=batch, latents=buffer ) if tuple(sharded_noise.shape) != local_shape: raise ValueError( @@ -145,10 +141,19 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): ), "True log-probability computation requires a non-zero noise level." dt = next_sigma - current_sigma + # sde/cps: cast to fp32 to match flowGRPO semantics and avoid the + # 0-dim-fp32 wrapped-scalar promotion demoting log-prob to bf16. + # ode: keep dtypes unchanged so rollout(ode) stays bit-exact with + # rollout=False (scheduling_flow_match_euler_discrete.step()). + # log_prob is computed on the full pre-shard noise buffer so SP ranks + # produce identical sums — see collect_rollout_log_probs(). if sde_type == "sde": + model_output = model_output.float() + sample = sample.float() variance_noise = self._rollout_variance_noise( batch, model_output, generator ) + full_variance_noise = rollout_session_data.noise_buffer std_dev_t = ( torch.sqrt( current_sigma @@ -173,12 +178,15 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): weighted_variance_noise = variance_noise * noise_std_dev prev_sample = prev_sample_mean + weighted_variance_noise - log_prob_no_const_val = -(weighted_variance_noise**2) + log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2) elif sde_type == "cps": + model_output = model_output.float() + sample = sample.float() variance_noise = self._rollout_variance_noise( batch, model_output, generator ) + full_variance_noise = rollout_session_data.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 @@ -189,7 +197,7 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): weighted_variance_noise = variance_noise * noise_std_dev prev_sample = prev_sample_mean + weighted_variance_noise - log_prob_no_const_val = -(weighted_variance_noise**2) + log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2) elif sde_type == "ode": prev_sample = sample + dt * model_output @@ -198,7 +206,11 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): noise_std_dev = torch.zeros( (), device=model_output.device, dtype=model_output.dtype ) - log_prob_no_const_val = torch.zeros_like(model_output) + log_prob_no_const_val = torch.zeros( + rollout_session_data.latents_shape, + device=model_output.device, + dtype=torch.float32, + ) assert ( 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" @@ -245,25 +257,20 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): self, batch ) -> tuple[torch.Tensor, torch.Tensor]: rollout_session_data = self._get_rollout_session_data(batch) - 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) + # [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 = [] return values_sum, values_count def collect_rollout_log_probs(self, batch: Req) -> torch.Tensor | None: - """Consume local rollout log probs and merge for all SP ranks.""" + """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) ) - if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): - packed = torch.stack( - [trajectory_log_prob_sum, trajectory_log_prob_count], dim=0 - ).to(get_local_torch_device()) - sequence_model_parallel_all_reduce(packed) - trajectory_log_prob_sum = packed[0] - trajectory_log_prob_count = packed[1] - 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/sp_utils.py b/python/sglang/multimodal_gen/runtime/post_training/sp_utils.py new file mode 100644 index 000000000..d05ad0924 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/sp_utils.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sequence Parallel helpers for post-training rollout code.""" + +from __future__ import annotations + +import torch + +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, + get_sp_world_size, +) +from sglang.multimodal_gen.runtime.distributed.communication_op import ( + sequence_model_parallel_all_gather, + sequence_model_parallel_all_reduce, +) + + +def should_do_sp_collective(batch) -> bool: + return get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False) + + +def gather_stacked_latents_for_sp( + pipeline_config, + batch, + stacked_latents: torch.Tensor, +) -> torch.Tensor: + if not should_do_sp_collective(batch): + return stacked_latents + if stacked_latents.dim() < 2: + return stacked_latents + bsz, t_steps = stacked_latents.shape[0], stacked_latents.shape[1] + flat_inputs = stacked_latents.flatten(0, 1).contiguous() + gathered_flat_inputs = pipeline_config.gather_latents_for_sp( + flat_inputs, batch=batch + ) + return gathered_flat_inputs.unflatten(0, (bsz, t_steps)) + + +def all_reduce_if_sp_sharded(batch, tensor: torch.Tensor) -> torch.Tensor: + if not should_do_sp_collective(batch): + return tensor + tensor = tensor.to(get_local_torch_device()) + sequence_model_parallel_all_reduce(tensor) + return tensor + + +def all_gather_if_sp_sharded(batch, x: torch.Tensor, dim: int = 0) -> torch.Tensor: + if not should_do_sp_collective(batch): + return x + x = x.to(get_local_torch_device()).contiguous() + return sequence_model_parallel_all_gather(x, dim=dim) + + +def maybe_trim_sp_rope_seq_for_batch(batch, rope: torch.Tensor) -> torch.Tensor: + raw = getattr(batch, "raw_latent_shape", None) + if raw is None or len(raw) < 2: + return rope + target = int(raw[1]) + if rope.shape[0] > target: + return rope[:target] + return rope diff --git a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py new file mode 100644 index 000000000..e37108760 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py @@ -0,0 +1,360 @@ +"""Unit tests for the rollout generate API (serialization, io_struct, rollout_api).""" + +import types +import unittest + +import torch + +from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import ( + _maybe_deserialize, + _maybe_serialize, + base64_to_tensor, + tensor_to_base64, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch +from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutDebugTensors, + RolloutDenoisingEnv, + RolloutDitTrajectory, + RolloutTrajectoryData, +) + + +class TestTensorToBase64Roundtrip(unittest.TestCase): + + def _roundtrip(self, t: torch.Tensor): + encoded = tensor_to_base64(t) + self.assertIsInstance(encoded, str) + decoded = base64_to_tensor(encoded) + self.assertTrue( + torch.equal(t, decoded), f"Mismatch for shape={t.shape} dtype={t.dtype}" + ) + + def test_float32_1d(self): + self._roundtrip(torch.randn(16)) + + def test_float32_nd(self): + self._roundtrip(torch.randn(2, 4, 8, 8)) + + def test_float16(self): + self._roundtrip(torch.randn(3, 5).half()) + + def test_int64(self): + self._roundtrip(torch.arange(10)) + + def test_bool(self): + self._roundtrip(torch.tensor([True, False, True])) + + def test_scalar(self): + self._roundtrip(torch.tensor(3.14)) + + def test_empty(self): + self._roundtrip(torch.empty(0)) + + def test_cuda_tensor_moves_to_cpu(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA not available") + t = torch.randn(4, device="cuda") + encoded = tensor_to_base64(t) + decoded = base64_to_tensor(encoded) + self.assertTrue(torch.equal(t.cpu(), decoded)) + + def test_non_contiguous(self): + t = torch.randn(4, 6)[:, ::2] + self.assertFalse(t.is_contiguous()) + self._roundtrip(t.contiguous()) + decoded = base64_to_tensor(tensor_to_base64(t)) + self.assertTrue(torch.equal(t.contiguous(), decoded)) + + def test_grad_tensor_detaches(self): + t = torch.randn(3, requires_grad=True) + encoded = tensor_to_base64(t) + decoded = base64_to_tensor(encoded) + self.assertFalse(decoded.requires_grad) + self.assertTrue(torch.equal(t.detach(), decoded)) + + +class TestMaybeSerialize(unittest.TestCase): + def test_tensor(self): + t = torch.randn(2, 3) + result = _maybe_serialize(t) + self.assertIsInstance(result, dict) + self.assertTrue(result["__tensor__"]) + self.assertEqual(result["shape"], [2, 3]) + self.assertEqual(result["dtype"], "torch.float32") + decoded = base64_to_tensor(result["data"]) + self.assertTrue(torch.equal(t, decoded)) + + def test_dict_with_tensors(self): + d = {"a": torch.tensor([1.0]), "b": "hello", "c": 42} + result = _maybe_serialize(d) + self.assertIsInstance(result, dict) + self.assertTrue(result["a"]["__tensor__"]) + self.assertEqual(result["b"], "hello") + self.assertEqual(result["c"], 42) + + def test_list_with_tensors(self): + lst = [torch.tensor(1.0), "text", torch.tensor(2.0)] + result = _maybe_serialize(lst) + self.assertIsInstance(result, list) + self.assertTrue(result[0]["__tensor__"]) + self.assertEqual(result[1], "text") + self.assertTrue(result[2]["__tensor__"]) + + def test_nested_structure(self): + nested = { + "level1": {"level2": [torch.tensor(1.0), {"level3": torch.tensor(2.0)}]} + } + result = _maybe_serialize(nested) + self.assertTrue(result["level1"]["level2"][0]["__tensor__"]) + self.assertTrue(result["level1"]["level2"][1]["level3"]["__tensor__"]) + + def test_none_passthrough(self): + self.assertIsNone(_maybe_serialize(None)) + + def test_plain_values_passthrough(self): + self.assertEqual(_maybe_serialize(42), 42) + self.assertEqual(_maybe_serialize("hello"), "hello") + self.assertAlmostEqual(_maybe_serialize(3.14), 3.14) + + def test_tuple_becomes_list(self): + result = _maybe_serialize((torch.tensor(1.0), 2)) + self.assertIsInstance(result, list) + self.assertEqual(len(result), 2) + + +from sglang.multimodal_gen.runtime.entrypoints.post_training.rollout_api import ( + _build_response, + _serialize_rollout_trajectory, +) + + +class TestSerializeRolloutTrajectory(unittest.TestCase): + def test_none_input(self): + log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(None) + self.assertIsNone(log_probs) + self.assertIsNone(debug) + self.assertIsNone(env) + self.assertIsNone(dit_traj) + + def test_log_probs_only(self): + rtd = RolloutTrajectoryData( + rollout_log_probs=torch.tensor([-1.0, -2.0]), + ) + log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd) + self.assertIsNotNone(log_probs) + self.assertTrue(log_probs["__tensor__"]) + self.assertIsNone(debug) + self.assertIsNone(env) + self.assertIsNone(dit_traj) + + def test_log_probs_none_in_rtd(self): + rtd = RolloutTrajectoryData(rollout_log_probs=None) + log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd) + self.assertIsNone(log_probs) + self.assertIsNone(debug) + self.assertIsNone(env) + self.assertIsNone(dit_traj) + + def test_with_debug_tensors(self): + dt = RolloutDebugTensors( + rollout_variance_noises=torch.randn(2, 5, 4, 8, 8), + rollout_prev_sample_means=torch.randn(2, 5, 4, 8, 8), + rollout_noise_std_devs=torch.randn(2, 5, 1), + rollout_model_outputs=torch.randn(2, 5, 4, 8, 8), + ) + rtd = RolloutTrajectoryData( + rollout_log_probs=torch.tensor([-0.5, -0.6]), + rollout_debug_tensors=dt, + ) + log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd) + self.assertIsNotNone(log_probs) + self.assertIsNotNone(debug) + self.assertIsNone(env) + self.assertIsNone(dit_traj) + self.assertIn("rollout_variance_noises", debug) + self.assertIn("rollout_prev_sample_means", debug) + self.assertIn("rollout_noise_std_devs", debug) + self.assertIn("rollout_model_outputs", debug) + self.assertTrue(debug["rollout_variance_noises"]["__tensor__"]) + + def test_debug_tensors_with_none_fields(self): + dt = RolloutDebugTensors( + rollout_variance_noises=None, + rollout_prev_sample_means=torch.randn(1, 2, 4, 4, 4), + rollout_noise_std_devs=None, + rollout_model_outputs=None, + ) + rtd = RolloutTrajectoryData( + rollout_log_probs=torch.tensor([-0.3]), + rollout_debug_tensors=dt, + ) + log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd) + self.assertIsNotNone(debug) + self.assertIsNone(debug["rollout_variance_noises"]) + self.assertTrue(debug["rollout_prev_sample_means"]["__tensor__"]) + self.assertIsNone(env) + self.assertIsNone(dit_traj) + + def test_with_denoising_env(self): + rtd = RolloutTrajectoryData( + denoising_env=RolloutDenoisingEnv( + image_kwargs={"encoder_hidden_states_image": [torch.randn(1, 8)]}, + pos_cond_kwargs={"encoder_hidden_states": torch.randn(1, 8)}, + neg_cond_kwargs={"encoder_hidden_states": torch.randn(1, 8)}, + guidance=torch.tensor([3.5]), + ), + dit_trajectory=RolloutDitTrajectory( + latents=torch.randn(1, 5, 4, 2, 2, 2), + timesteps=torch.tensor([1.0, 0.75, 0.5, 0.25]), + ), + ) + _, _, env, dit_traj = _serialize_rollout_trajectory( + rtd, + serialized_dit_timesteps=_maybe_serialize(rtd.dit_trajectory.timesteps), + ) + self.assertIsNotNone(env) + self.assertIn("pos_cond_kwargs", env) + self.assertNotIn("trajectory", env) + self.assertIsNotNone(dit_traj) + self.assertIn("latents", dit_traj) + self.assertIn("timesteps", dit_traj) + self.assertTrue(dit_traj["latents"]["__tensor__"]) + self.assertTrue(dit_traj["timesteps"]["__tensor__"]) + + +class TestBuildResponse(unittest.TestCase): + def _make_metrics(self, duration_s: float = 1.0): + return types.SimpleNamespace(total_duration_s=duration_s) + + def test_minimal_output(self): + batch = OutputBatch( + output=torch.randn(1, 3, 1, 64, 64), + rollout_trajectory_data=RolloutTrajectoryData( + rollout_log_probs=torch.tensor([0.0]), + ), + ) + batch.metrics = self._make_metrics(2.5) + resps = _build_response("r1", "prompt", 42, True, batch) + self.assertEqual(len(resps), 1) + resp = resps[0] + self.assertEqual(resp.request_id, "r1") + self.assertEqual(resp.prompt, "prompt") + self.assertEqual(resp.seed, 42) + self.assertIsNotNone(resp.generated_output) + self.assertIsNotNone(resp.rollout_log_probs) + lp = base64_to_tensor(resp.rollout_log_probs["data"]) + self.assertEqual(lp.shape, ()) + self.assertAlmostEqual(resp.inference_time_s, 2.5) + + def test_full_response(self): + batch = OutputBatch( + output=torch.randn(1, 3, 1, 64, 64), + rollout_trajectory_data=RolloutTrajectoryData( + rollout_log_probs=torch.tensor([-0.5]), + ), + peak_memory_mb=8192.0, + ) + batch.metrics = self._make_metrics(5.0) + resps = _build_response("r2", "test", 99, True, batch) + self.assertEqual(len(resps), 1) + resp = resps[0] + self.assertIsNotNone(resp.rollout_log_probs) + self.assertIsNone(resp.rollout_debug_tensors) + self.assertAlmostEqual(resp.peak_memory_mb, 8192.0) + + def test_no_metrics(self): + batch = OutputBatch( + output=torch.randn(1, 3, 1, 64, 64), + rollout_trajectory_data=RolloutTrajectoryData( + rollout_log_probs=torch.tensor([0.0]), + ), + ) + batch.metrics = None + resp = _build_response("r3", "p", 1, True, batch)[0] + self.assertIsNone(resp.inference_time_s) + + def test_zero_metrics(self): + batch = OutputBatch( + output=torch.randn(1, 3, 1, 64, 64), + rollout_trajectory_data=RolloutTrajectoryData( + rollout_log_probs=torch.tensor([0.0]), + ), + ) + batch.metrics = self._make_metrics(0.0) + resp = _build_response("r4", "p", 1, True, batch)[0] + self.assertIsNone(resp.inference_time_s) + + def test_zero_peak_memory(self): + batch = OutputBatch( + output=torch.randn(1, 3, 1, 64, 64), + peak_memory_mb=0.0, + rollout_trajectory_data=RolloutTrajectoryData( + rollout_log_probs=torch.tensor([0.0]), + ), + ) + batch.metrics = None + resp = _build_response("r6", "p", 1, True, batch)[0] + self.assertIsNone(resp.peak_memory_mb) + + def test_batch_splits_log_probs_and_output(self): + B, T = 2, 3 + batch = OutputBatch( + output=torch.randn(B, 1, 8, 8), + rollout_trajectory_data=RolloutTrajectoryData( + rollout_log_probs=torch.randn(B, T), + ), + ) + batch.metrics = self._make_metrics(1.0) + resps = _build_response("rb", "p", 0, True, batch) + self.assertEqual(len(resps), B) + lp0 = base64_to_tensor(resps[0].rollout_log_probs["data"]) + lp1 = base64_to_tensor(resps[1].rollout_log_probs["data"]) + self.assertEqual(lp0.shape, (T,)) + self.assertEqual(lp1.shape, (T,)) + g0 = base64_to_tensor(resps[0].generated_output["data"]) + g1 = base64_to_tensor(resps[1].generated_output["data"]) + self.assertEqual(g0.shape, (1, 8, 8)) + self.assertEqual(g1.shape, (1, 8, 8)) + self.assertFalse(torch.equal(g0, g1)) + + def test_batch_dit_timesteps_on_each_row_one_serialize(self): + B, T, D = 2, 4, 3 + batch = OutputBatch( + output=torch.randn(B, 1, 8, 8), + rollout_trajectory_data=RolloutTrajectoryData( + rollout_log_probs=torch.randn(B, T), + dit_trajectory=RolloutDitTrajectory( + latents=torch.randn(B, T + 1, D), + timesteps=torch.linspace(1.0, 0.0, T), + ), + ), + ) + batch.metrics = self._make_metrics(1.0) + resps = _build_response("r", "p", 0, True, batch) + self.assertEqual(len(resps), B) + self.assertIsNotNone(resps[0].dit_trajectory) + self.assertIsNotNone(resps[1].dit_trajectory) + ts0 = base64_to_tensor(resps[0].dit_trajectory["timesteps"]["data"]) + ts1 = base64_to_tensor(resps[1].dit_trajectory["timesteps"]["data"]) + self.assertEqual(ts0.shape, (T,)) + self.assertTrue(torch.equal(ts0, ts1)) + self.assertEqual( + _maybe_deserialize(resps[1].dit_trajectory["latents"]).shape, (T + 1, D) + ) + + def test_rollout_false_omits_trajectory(self): + batch = OutputBatch( + output=torch.randn(2, 1, 8, 8), + rollout_trajectory_data=None, + ) + batch.metrics = self._make_metrics(1.0) + resps = _build_response("r0", "p", 0, False, batch) + self.assertEqual(len(resps), 2) + self.assertIsNone(resps[0].rollout_log_probs) + self.assertIsNone(resps[1].rollout_log_probs) + self.assertIsNotNone(resps[0].generated_output) + + +if __name__ == "__main__": + unittest.main() 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 8231b3634..00b4eb5cc 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 @@ -71,6 +71,48 @@ class TestSchedulerRolloutOdeUnit(unittest.TestCase): self.assertEqual(tuple(local_elem_count.shape), (sample.shape[0],)) self.assertTrue(torch.all(local_elem_count == float(sample[0].numel()))) + def test_ode_bit_exact_with_non_rollout_path(self): + """ODE rollout must produce the exact same prev_sample as the + non-rollout deterministic branch in + ``scheduling_flow_match_euler_discrete.step`` (``prev_sample = + sample + dt * model_output``). Uses bf16 model_output because the + 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) + + sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) + model_output = torch.randn_like(sample).to(torch.bfloat16) + current_sigma = torch.tensor(0.6, dtype=torch.float32) + next_sigma = torch.tensor(0.4, dtype=torch.float32) + dt = next_sigma - current_sigma + + rollout_prev = scheduler.flow_sde_sampling( + batch, + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + generator=torch.Generator(device=sample.device).manual_seed(0), + ) + # Exact expression used by the non-rollout branch at + # scheduling_flow_match_euler_discrete.py `prev_sample = sample + + # dt * model_output` (after the shared ``sample.to(fp32)`` cast). + non_rollout_prev = sample + dt * model_output + + self.assertEqual(rollout_prev.dtype, non_rollout_prev.dtype) + self.assertTrue(torch.equal(rollout_prev, non_rollout_prev)) + # Also verify the post-cast to model_output.dtype (what scheduler.step + # returns downstream) is bit-exact. + self.assertTrue( + torch.equal( + rollout_prev.to(model_output.dtype), + non_rollout_prev.to(model_output.dtype), + ) + ) + def test_ode_debug_tensors_have_shape_safe_noise_std(self): scheduler = _DummyScheduler() batch = self._build_batch(debug_mode=True) @@ -215,8 +257,19 @@ class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase): 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] - lambda _batch, *_args, **_kwargs: variance_noise + _mock_rollout_variance_noise ) prev_sgl = scheduler.flow_sde_sampling( @@ -277,6 +330,57 @@ class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase): msg=f"{sde_type} seed={seed} {name} max_abs={err:.9f}", ) + def test_sde_cps_force_fp32_with_bf16_model_output(self): + """Regression for PyTorch's wrapped-scalar promotion trap: a 0-dim + fp32 ``noise_std_dev`` multiplied by an N-dim bf16 tensor silently + demotes to bf16, which would corrupt log-prob precision. SDE/CPS + 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) + pipeline_config = types.SimpleNamespace( + shard_latents_for_sp=lambda batch, latents: (latents, False) + ) + + 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) + + g = torch.Generator(device="cpu").manual_seed(0) + model_output = torch.randn(shape, generator=g, dtype=torch.float32).to( + torch.bfloat16 + ) + sample = torch.randn(shape, generator=g, dtype=torch.float32) + + # Use the real _rollout_variance_noise (no mock) so its dtype + # 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( + batch, + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + generator=g, + ) + 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 + self.assertEqual( + noise_buffer.dtype, + torch.float32, + msg=f"{sde_type}: noise_buffer must be fp32 with bf16 model_output", + ) + if __name__ == "__main__": unittest.main()