[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 <lishunyang12@163.com>

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(

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
 

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)

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(