From 94eef833fe40460230cd6856e2602e038c8305e0 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 18 Aug 2026 05:42:27 -0700 Subject: [PATCH] [diffusion] rl: support cosmos3 (#34197) --- .../stages/model_specific_stages/cosmos3.py | 99 ++++++++-- .../post_training/rollout_scheduler.py | 31 ++++ .../runtime/post_training/weights_updater.py | 31 +++- .../test/unit/test_cosmos3_rollout.py | 169 ++++++++++++++++++ 4 files changed, 312 insertions(+), 18 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py 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 697ce3bbf..80316b67c 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 @@ -56,6 +56,12 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( 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_scheduler import ( + prepare_rollout_request_scheduler, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler @@ -478,6 +484,8 @@ class Cosmos3LatentPreparationStage(PipelineStage): generator = batch.generator if generator is None and batch.seed is not None: generator = torch.Generator(device=device).manual_seed(batch.seed) + # The rollout SDE step draws its variance noise from this generator. + batch.generator = generator noise = torch.randn(shape, generator=generator, device=device, dtype=dtype) @@ -763,9 +771,10 @@ class Cosmos3TimestepPreparationStage(PipelineStage): return batch num_inference_steps = batch.num_inference_steps - flow_shift = getattr(batch, "flow_shift", None) - if flow_shift is None: - flow_shift = pipeline_config.flow_shift + explicit_flow_shift = getattr(batch, "flow_shift", None) + if explicit_flow_shift is None: + explicit_flow_shift = pipeline_config.flow_shift + flow_shift = explicit_flow_shift if flow_shift is None: flow_shift = self._default_flow_shift_for_mode( batch, bool(pipeline_config.is_edge) @@ -776,13 +785,22 @@ class Cosmos3TimestepPreparationStage(PipelineStage): self.scheduler.set_timesteps(num_inference_steps, device=device) batch.timesteps = self.scheduler.timesteps + if batch.rollout: + prepare_rollout_request_scheduler( + batch, + self.scheduler, + explicit_shift=explicit_flow_shift, + num_inference_steps=num_inference_steps, + device=device, + ) + self.log_info( f"Prepared {len(batch.timesteps)} timesteps (flow_shift={flow_shift})" ) return batch -class Cosmos3DenoisingStage(PipelineStage): +class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): """Cosmos3 denoise loop, including CFG and the parallelism modes. The UND pathway runs once and its K/V is cached per cache_key (``cond`` / @@ -985,6 +1003,37 @@ class Cosmos3DenoisingStage(PipelineStage): condition_latents = batch.extra.get("condition_latents") guidance_interval = getattr(batch.sampling_params, "guidance_interval", None) + # Rollout requests carry a per-request scheduler bound by the timestep stage. + scheduler = batch.scheduler if batch.scheduler is not None else self.scheduler + if batch.rollout: + if velocity_mask is not None or condition_latents is not None: + raise ValueError( + "Cosmos3 rollout supports T2V/T2I only; I2V/V2V " + "conditioned-frame re-blending breaks the Gaussian " + "transition assumption of the SDE log-prob math." + ) + if action_latents is not None or sound_latents is not None: + 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, + ) + do_cfg = guidance_scale > 1.0 enable_cfg_parallel = server_args.enable_cfg_parallel and do_cfg @@ -1154,13 +1203,31 @@ class Cosmos3DenoisingStage(PipelineStage): if velocity_mask is not None: noise_pred = noise_pred * velocity_mask - latents = self.scheduler.step( - noise_pred, - t, - latents, - generator=generator, - return_dict=False, - )[0] + 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, + 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] if action_noise_pred is not None: # Zero the velocity at conditioned (clean) action tokens and at @@ -1204,6 +1271,15 @@ class Cosmos3DenoisingStage(PipelineStage): if batch.profile and not batch.is_warmup: self.step_profile() + 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, + ) + batch.latents = latents if action_latents is not None: batch.action_latents = action_latents @@ -1641,6 +1717,7 @@ class Cosmos3DecodingStage(PipelineStage): audio_sample_rate=audio_sample_rate, action_pred=action_pred, metrics=batch.metrics if hasattr(batch, "metrics") else None, + rollout_trajectory_data=batch.rollout_trajectory_data, **action_metadata, ) diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_scheduler.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_scheduler.py index da2d3b7bc..152ae7abe 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_scheduler.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_scheduler.py @@ -1,5 +1,7 @@ from typing import Any +import torch + from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) @@ -39,3 +41,32 @@ def get_or_create_rollout_request_scheduler( scheduler, isolate=needs_clone, ) + + +def prepare_rollout_request_scheduler( + batch: Req, + serving_scheduler: Any, + *, + explicit_shift: float | None, + num_inference_steps: int, + device: torch.device, +) -> None: + """Bind the per-request rollout scheduler and hand it a sigma grid. + + Without an explicit shift the rollout scheduler inherits the serving grid + verbatim; an explicit shift selects a plain shifted grid instead. + """ + scheduler = get_or_create_rollout_request_scheduler(batch, serving_scheduler) + if scheduler is not serving_scheduler: + if explicit_shift is not None: + # Unwarped base; the mapped scheduler's sigmas are already serving-shift warped. + shift = float(explicit_shift) + num_train_timesteps = scheduler.config.num_train_timesteps + sigmas = torch.linspace(1.0, 1.0 / num_train_timesteps, num_inference_steps) + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + else: + sigmas = serving_scheduler.sigmas[:-1] + # shift=1.0 so set_timesteps keeps the explicit sigmas verbatim. + scheduler.set_shift(1.0) + scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device) + batch.timesteps = scheduler.timesteps diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py index fcc50e1d6..0f9453e9e 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -194,7 +194,13 @@ def _load_weights_into_module(module: torch.nn.Module, weights_iter) -> None: ] if offload_managers: - weight_dict = dict(weights_iter) + entries = list(weights_iter) + if any(shard_id is not None for _, _, shard_id in entries): + raise NotImplementedError( + "Fused-parameter weight updates are not supported for " + "layerwise-offloaded modules." + ) + weight_dict = {n: w for n, w, _ in entries} offloaded_names: set[str] = set() for manager in offload_managers: offloaded_names.update(manager.update_cpu_weights(weight_dict)) @@ -267,17 +273,22 @@ def _iter_module_weight_updates( weights_iter, model_params: dict, ): + """Yield (mapped_name, weight, shard_id); shard_id is the merge index for + weights that map into a fused parameter (e.g. q/k/v -> to_qkv), else None. + """ map_name = _build_module_weight_name_mapper(module) module_name = type(module).__name__ for name, loaded_weight in weights_iter: if name in model_params: - yield name, loaded_weight + yield name, loaded_weight, None continue - mapped_name = map_name(name)[0] if map_name is not None else name + mapped_name, merge_index = ( + map_name(name) if map_name is not None else (name, None) + ) if mapped_name in model_params: - yield mapped_name, loaded_weight + yield mapped_name, loaded_weight, merge_index continue logger.warning( @@ -291,15 +302,21 @@ def _iter_module_weight_updates( def load_weights_into_model( weights_iter, model_params: dict, module_name: str | None = None ) -> None: - """Copy weights from weights_iter into model_params in-place.""" - for name, loaded_weight in weights_iter: + """Copy weights into model_params in-place; entries are (name, weight) or + (name, weight, shard_id), shard_id routing fused parts via weight_loader.""" + for entry in weights_iter: + name, loaded_weight, *rest = entry + shard_id = rest[0] if rest else None if name not in model_params: logger.warning("Skipping weight update: parameter %r not found", name) continue param = model_params[name] weight_loader = getattr(param, "weight_loader", None) if callable(weight_loader): - weight_loader(param, loaded_weight.to(param.dtype)) + if shard_id is not None: + weight_loader(param, loaded_weight.to(param.dtype), shard_id) + else: + weight_loader(param, loaded_weight.to(param.dtype)) else: dtensor_param = param if isinstance(param, DTensor) else None if dtensor_param is None and isinstance( diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py new file mode 100644 index 000000000..70ca35733 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the Cosmos3 RL rollout path: rollout scheduler grid +selection and fused-parameter shard-id plumbing in the weights updater.""" + +import types +import unittest + +import torch + +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.post_training.rollout_scheduler import ( + prepare_rollout_request_scheduler, +) +from sglang.multimodal_gen.runtime.post_training.weights_updater import ( + _load_weights_into_module, +) + +NUM_STEPS = 16 +NUM_TRAIN_TIMESTEPS = 1000 + + +def _serving_scheduler() -> FlowUniPCMultistepScheduler: + """Mirror the Cosmos3 checkpoint's serving scheduler (UniPC, flow + prediction) with the T2I per-mode flow_shift default applied.""" + scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=NUM_TRAIN_TIMESTEPS) + scheduler.set_shift(3.0) + scheduler.set_timesteps(NUM_STEPS, device="cpu") + return scheduler + + +def _rollout_batch() -> types.SimpleNamespace: + return types.SimpleNamespace(rollout=True, scheduler=None, timesteps=None) + + +def _prepare(serving, batch, explicit_shift): + prepare_rollout_request_scheduler( + batch, + serving, + explicit_shift=explicit_shift, + num_inference_steps=NUM_STEPS, + device=torch.device("cpu"), + ) + + +class TestPrepareRolloutRequestScheduler(unittest.TestCase): + + def test_inherits_serving_grid_without_explicit_shift(self): + serving = _serving_scheduler() + batch = _rollout_batch() + + _prepare(serving, batch, None) + + self.assertIsInstance(batch.scheduler, FlowMatchEulerDiscreteScheduler) + self.assertIsNot(batch.scheduler, serving) + # Serving grid inherited verbatim, terminal sigma re-appended as 0. + torch.testing.assert_close( + batch.scheduler.sigmas[:-1].float(), + serving.sigmas[:-1].float(), + atol=1e-6, + rtol=0, + ) + self.assertEqual(batch.scheduler.sigmas[-1].item(), 0.0) + # batch.timesteps switches to the rollout grid (t = sigma * T). + torch.testing.assert_close( + batch.timesteps.float(), + batch.scheduler.sigmas[:-1].float() * NUM_TRAIN_TIMESTEPS, + atol=1e-3, + rtol=0, + ) + + def test_explicit_flow_shift_selects_plain_shifted_grid(self): + serving = _serving_scheduler() + batch = _rollout_batch() + shift = 2.0 + + _prepare(serving, batch, shift) + + # Plain flow-match Euler grid under the requested shift, not the serving grid. + base = torch.linspace(1.0, 1.0 / NUM_TRAIN_TIMESTEPS, NUM_STEPS) + expected = shift * base / (1 + (shift - 1) * base) + torch.testing.assert_close( + batch.scheduler.sigmas[:-1].float(), expected.float(), atol=1e-5, rtol=0 + ) + self.assertFalse( + torch.allclose(batch.scheduler.sigmas[:-1], serving.sigmas[:-1].float()) + ) + + def test_rl_capable_serving_scheduler_passes_through(self): + # An RL-capable serving scheduler is shared and keeps its serving grid. + serving = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=NUM_TRAIN_TIMESTEPS + ) + serving.set_shift(3.0) + serving.set_timesteps(NUM_STEPS, device="cpu") + expected = serving.sigmas.clone() + batch = _rollout_batch() + + _prepare(serving, batch, 3.0) + + self.assertIs(batch.scheduler, serving) + torch.testing.assert_close(serving.sigmas, expected) + + +class _FusedParamModule(torch.nn.Module): + """Diffusers-style q/k/v weights that map into one fused to_qkv param.""" + + def __init__(self): + super().__init__() + qkv = torch.nn.Module() + qkv.weight = torch.nn.Parameter(torch.zeros(6, 2)) + attn = torch.nn.Module() + attn.to_qkv = qkv + self.attn = attn + self.out_proj = torch.nn.Linear(2, 2, bias=False) + self.param_names_mapping = { + r"^attn\.q\.(weight)$": (r"attn.to_qkv.\1", 0, 3), + r"^attn\.k\.(weight)$": (r"attn.to_qkv.\1", 1, 3), + r"^attn\.v\.(weight)$": (r"attn.to_qkv.\1", 2, 3), + r"^proj\.(weight)$": r"out_proj.\1", + } + + +class TestWeightsUpdaterFusedParams(unittest.TestCase): + + def test_merge_index_reaches_weight_loader_as_shard_id(self): + module = _FusedParamModule() + calls = [] + + def loader(param, weight, *args): + calls.append((weight.clone(), args)) + + module.attn.to_qkv.weight.weight_loader = loader + + _load_weights_into_module( + module, + [ + ("attn.q.weight", torch.full((2, 2), 1.0)), + ("attn.k.weight", torch.full((2, 2), 2.0)), + ("attn.v.weight", torch.full((2, 2), 3.0)), + ("proj.weight", torch.full((2, 2), 4.0)), + ], + ) + + # Fused q/k/v parts arrive with their merge index as shard id. + self.assertEqual([args for _, args in calls], [(0,), (1,), (2,)]) + # Renamed non-fused weight lands via plain in-place copy. + torch.testing.assert_close(module.out_proj.weight.data, torch.full((2, 2), 4.0)) + + def test_direct_hit_keeps_two_arg_weight_loader_call(self): + module = _FusedParamModule() + calls = [] + + def loader(param, weight, *args): + calls.append(args) + + module.attn.to_qkv.weight.weight_loader = loader + + _load_weights_into_module(module, [("attn.to_qkv.weight", torch.zeros(6, 2))]) + + self.assertEqual(calls, [()]) + + +if __name__ == "__main__": + unittest.main()