From bb8614595d44e578ce91d2e49fd44989150661e7 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:23:36 +0800 Subject: [PATCH] [BugFix][Cosmos3] Pad sound latents so video+sound runs under sequence parallelism (#4678) Signed-off-by: lishunyang12 <125541396+lishunyang12@users.noreply.github.com> Signed-off-by: lishunyang --- .../models/cosmos3/test_cosmos3_pipeline.py | 2 +- .../cosmos3/test_cosmos3_transformer.py | 21 ++++++++++++++++ .../models/cosmos3/pipeline_cosmos3.py | 25 ++++++++++++++----- .../models/cosmos3/transformer_cosmos3.py | 24 ++++++++++++++++++ 4 files changed, 65 insertions(+), 7 deletions(-) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index c4eab9812..da4ae32e3 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -1422,7 +1422,7 @@ class TestForwardRouting: pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3) sound_latents = torch.zeros(1, 3, 4) pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10) - pipeline._prepare_sound_latents = lambda *args: (sound_latents, 4) + pipeline._prepare_sound_latents = lambda *args, **kwargs: (sound_latents, 4) pipeline._decode_sound_latents = lambda *args: torch.ones(1, 2, 20) output = pipeline.forward( make_request_batch( diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index cb4e28b51..70598c952 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -330,6 +330,27 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch ) +def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + model = object.__new__(transformer_cosmos3.Cosmos3VFMTransformer) + model.latent_patch_size = 2 + vs = (3, 16, 16) + + monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 1) + assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 97 + + monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 2) + assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 98 + assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=98) == 98 + + # ulysses=4, with a transfer control folded into the vision base. + monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 4) + padded = model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97, num_vision_items=2) + assert (2 * 192 + padded) % 4 == 0 + + def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 890b76d6b..4178115d5 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1666,12 +1666,21 @@ class Cosmos3OmniDiffusersPipeline( self, target_audio_samples: int, generator: torch.Generator, + *, + sp_video_shape: tuple[int, int, int] | None = None, + sp_num_vision_items: int = 1, ) -> tuple[torch.Tensor, int]: sound_tokenizer = self._get_sound_tokenizer() hop_size = int( getattr(sound_tokenizer, "hop_size", None) or getattr(sound_tokenizer, "temporal_compression_factor") ) latent_frames = max(1, math.ceil(max(1, int(target_audio_samples)) / hop_size)) + if sp_video_shape is not None: + latent_frames = self.transformer.sound_latent_frames_for_sequence_parallel( + video_shape=sp_video_shape, + sound_frames=latent_frames, + num_vision_items=sp_num_vision_items, + ) sound_dim = int(getattr(sound_tokenizer, "latent_ch", 64)) transformer_sound_dim = int(getattr(self.transformer, "sound_dim", sound_dim)) if sound_dim != transformer_sound_dim: @@ -3119,17 +3128,21 @@ class Cosmos3OmniDiffusersPipeline( image_latent = None condition_latents = None + T_latent = latents.shape[2] + H_latent = latents.shape[3] + W_latent = latents.shape[4] + video_shape = (T_latent, H_latent, W_latent) + sound_latents = None target_audio_samples = None sound_sample_rate = None if sound_enabled: target_audio_samples, _, sound_sample_rate = self._resolve_sound_target_samples(sp, num_frames, frame_rate) - sound_latents, _ = self._prepare_sound_latents(target_audio_samples, generator) - - T_latent = latents.shape[2] - H_latent = latents.shape[3] - W_latent = latents.shape[4] - video_shape = (T_latent, H_latent, W_latent) + sound_latents, _ = self._prepare_sound_latents( + target_audio_samples, + generator, + sp_video_shape=video_shape, + ) # --- Denoising loop --- shared_kwargs = dict(video_shape=video_shape, fps=frame_rate) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 99e0eacf0..261573335 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -1373,6 +1373,30 @@ class Cosmos3VFMTransformer(nn.Module): f"ulysses_degree ({ulysses_size}). {adjust_detail}" ) + def sound_latent_frames_for_sequence_parallel( + self, + *, + video_shape: tuple[int, int, int], + sound_frames: int, + num_vision_items: int = 1, + ) -> int: + # Sound is the only modality the packed GEN sequence pairs with here: action and + # sound are never generated together (the pipeline rejects action+sound), so the + # base is just the vision tokens. + # + # Note: padded frames go through attention and are only trimmed on decode, so SP + # output is not bit-exact with non-SP (the extra frame perturbs the kept ones). + from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_world_size + + ulysses_size = get_ulysses_parallel_world_size() + if ulysses_size <= 1 or sound_frames <= 0: + return sound_frames + t, h, w = video_shape + hp, wp, _, _ = self._pad_to_patch_size(h, w) + base = num_vision_items * t * hp * wp + pad = (-(base + sound_frames)) % ulysses_size + return sound_frames + pad + # -- Forward ------------------------------------------------------------- def forward(