[diffusion] rl: support cosmos3 (#34197)
This commit is contained in:
Andy Ye
+88
-11
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user