From bad50980c9e2679d4fa6352b0fbb121af9dfbb40 Mon Sep 17 00:00:00 2001 From: Yan Cao <31481315+yancaocn@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:43:34 +0800 Subject: [PATCH] [Bugfix] Fix Helios USP: per-component split for correct sequence parallelism (#6930) Signed-off-by: yancaocn Co-authored-by: yancaocn --- tests/diffusion/models/helios/__init__.py | 0 .../models/helios/test_helios_usp_2rank.py | 258 ++++++++++++++++++ .../models/helios/test_helios_usp_shard.py | 225 +++++++++++++++ .../models/helios/helios_transformer.py | 90 +++++- 4 files changed, 561 insertions(+), 12 deletions(-) create mode 100644 tests/diffusion/models/helios/__init__.py create mode 100644 tests/diffusion/models/helios/test_helios_usp_2rank.py create mode 100644 tests/diffusion/models/helios/test_helios_usp_shard.py diff --git a/tests/diffusion/models/helios/__init__.py b/tests/diffusion/models/helios/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/helios/test_helios_usp_2rank.py b/tests/diffusion/models/helios/test_helios_usp_2rank.py new file mode 100644 index 000000000..7b16d1d4f --- /dev/null +++ b/tests/diffusion/models/helios/test_helios_usp_2rank.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Two-process USP smoke test for the Helios transformer. + +Pins the multi-rank mechanism that ``test_helios_usp_shard.py`` (unit tests +of ``_sp_split_seq`` alone) cannot see: + + * the manual ``_sp_shard_depth`` bump in ``forward()`` must make attn1 run + the Ulysses all-to-all (not ``NoParallelAttention``); + * attn2 cross-attention must stay rank-local (``skip_sequence_parallel= + True``): if replicated text K/V were fed through Ulysses they would be + ws-fold duplicated and outputs would diverge from the single-rank run; + * end-to-end correctness: ws=2 output equals the ws=1 baseline for the same + weights/inputs (divisible shapes, so the padding path is not involved). + +Runs ``HeliosTransformer3DModel.forward`` with tiny random weights via +``torch.multiprocessing.spawn``: first a baseline (world_size=1), then an +SP run (world_size=2, ulysses_degree=2); rank-0 outputs are compared. + +Requires >= 2 devices; skipped otherwise. +""" + +import os +import tempfile + +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelConfig, + get_sp_plan_from_model, +) +from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context +from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel +from vllm_omni.platforms import current_omni_platform + +_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2) +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, *_L4_TWO_GPU] + +RANDOM_SEED = 42 +DTYPE = torch.float32 # keep both runs on the SDPA fallback for tight tolerance + +# Tiny but architecturally faithful config: +# inner_dim = 4 * 32 = 128, rope_dim sums to head_dim (32), +# current chunk 2x4x6 latent -> (1,2,2) patch -> 2*2*3 = 12 tokens, +# short history 1x4x6 latent -> (1,2,2) patch -> 6 tokens. +# All component lengths are even, so ws=2 needs no padding replicas. +TINY_MODEL_KWARGS = dict( + patch_size=(1, 2, 2), + num_attention_heads=4, + attention_head_dim=32, + in_channels=4, + out_channels=4, + text_dim=64, + freq_dim=32, + ffn_dim=128, + num_layers=2, + rope_dim=(12, 10, 10), + rope_theta=10000.0, + guidance_cross_attn=True, + zero_history_timestep=True, + has_multi_term_memory_patch=True, + is_amplify_history=False, +) + +BATCH = 1 +CUR_FRAMES, LAT_H, LAT_W = 2, 4, 6 # -> 12 current tokens +HIST_FRAMES = 1 # short history -> 6 tokens +TEXT_LEN = 5 + + +def _update_env(envs_dict: dict[str, str]) -> None: + for k, v in envs_dict.items(): + os.environ[k] = v + + +def _make_inputs(device: torch.device) -> dict[str, torch.Tensor]: + torch.manual_seed(RANDOM_SEED + 1) + return { + "hidden_states": torch.randn(BATCH, 4, CUR_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE), + "timestep": torch.full((BATCH,), 500.0, device=device), + "encoder_hidden_states": torch.randn( + BATCH, TEXT_LEN, TINY_MODEL_KWARGS["text_dim"], device=device, dtype=DTYPE + ), + "indices_latents_history_short": torch.tensor([[0]], device=device), + "latents_history_short": torch.randn(BATCH, 4, HIST_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE), + } + + +def _spy_strategies(model: HeliosTransformer3DModel) -> dict[str, list[str]]: + """Record the parallel strategy each Attention resolves to during forward.""" + seen: dict[str, list[str]] = {"attn1": [], "attn2": []} + for name, attn in (("attn1", model.blocks[0].attn1.attn), ("attn2", model.blocks[0].attn2.attn)): + original = attn._get_active_parallel_strategy + + def spy(original=original, name=name): + strategy = original() + seen[name].append(getattr(strategy, "name", type(strategy).__name__)) + return strategy + + attn._get_active_parallel_strategy = spy + return seen + + +def helios_forward_on_model( + local_rank: int, + world_size: int, + ulysses_degree: int, + output_file: str, + model_state_file: str, + is_baseline: bool, +) -> None: + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + torch.set_default_device(device) + torch.set_default_dtype(DTYPE) + + _update_env( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": "12361", + } + ) + init_distributed_environment() + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=ulysses_degree, + ulysses_degree=ulysses_degree, + ring_degree=1, + allgather_degree=1, + cfg_parallel_size=1, + ) + od_config = OmniDiffusionConfig.from_kwargs( + model="test_model", + dtype=DTYPE, + parallel_config=parallel_config, + diffusion_attention_backend="TORCH_SDPA", + ) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=ulysses_degree, + ulysses_degree=ulysses_degree, + ring_degree=1, + allgather_degree=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config): + torch.manual_seed(RANDOM_SEED) + model = HeliosTransformer3DModel(**TINY_MODEL_KWARGS).to(device).to(DTYPE) + + if is_baseline and local_rank == 0: + with open(model_state_file, "wb") as f: + torch.save({k: v.cpu() for k, v in model.state_dict().items()}, f) + + if world_size > 1: + torch.distributed.barrier() + + with open(model_state_file, "rb") as f: + model_state = torch.load(f) + model.load_state_dict({k: v.to(device).to(DTYPE) for k, v in model_state.items()}) + + if not is_baseline: + # Mirror registry._apply_sequence_parallel_if_enabled: install the + # proj_out gather hook from the model's own _sp_plan and flip + # sp_plan_hooks_applied so sp_active follows _sp_shard_depth + # (i.e. the manual bump in Helios forward is what enables Ulysses). + plan = get_sp_plan_from_model(model) + assert plan is not None, "HeliosTransformer3DModel must define _sp_plan" + apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=ulysses_degree), plan) + get_forward_context().sp_plan_hooks_applied = True + + seen = _spy_strategies(model) + inputs = _make_inputs(device) + output = model(**inputs).sample + + assert output.shape == (BATCH, 4, CUR_FRAMES, LAT_H, LAT_W), f"unexpected output shape {output.shape}" + + if is_baseline: + assert seen["attn1"] and all(n == "none" for n in seen["attn1"]), ( + f"baseline attn1 must use NoParallelAttention ('none'), saw {seen['attn1']}" + ) + else: + # attn1: SP communication must be active inside the block loop + # (regression here = the manual _sp_shard_depth bump was lost). + assert seen["attn1"] and all(n == "ulysses" for n in seen["attn1"]), ( + f"attn1 must run Ulysses under USP (manual _sp_shard_depth bump), saw {seen['attn1']}" + ) + # attn2: replicated text K/V must stay rank-local + # (regression here = the ws-fold duplicated-keys bug). + assert seen["attn2"] and all(n == "none" for n in seen["attn2"]), ( + f"attn2 must skip sequence parallel (replicated text K/V), saw {seen['attn2']}" + ) + + if local_rank == 0: + with open(output_file, "wb") as f: + torch.save(output.detach().cpu().float(), f) + + destroy_distributed_env() + + +def test_helios_usp2_matches_single_rank(): + """ws=2 (Ulysses) must reproduce the ws=1 output for identical weights/inputs.""" + available_devices = current_omni_platform.get_device_count() + if available_devices < 2: + pytest.skip(f"Test requires 2 devices but only {available_devices} available") + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + baseline_output_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + sp_output_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + model_state_file = f.name + + try: + torch.multiprocessing.spawn( + helios_forward_on_model, + args=(1, 1, baseline_output_file, model_state_file, True), + nprocs=1, + ) + torch.multiprocessing.spawn( + helios_forward_on_model, + args=(2, 2, sp_output_file, model_state_file, False), + nprocs=2, + ) + + with open(baseline_output_file, "rb") as f: + baseline = torch.load(f) + with open(sp_output_file, "rb") as f: + sp = torch.load(f) + + assert baseline.shape == sp.shape + max_abs_diff = (baseline - sp).abs().max().item() + assert torch.allclose(baseline, sp, atol=1e-3, rtol=1e-3), ( + f"USP=2 output diverges from single-rank baseline: max_abs_diff={max_abs_diff:.6e}" + ) + finally: + for path in [baseline_output_file, sp_output_file, model_state_file]: + if os.path.exists(path): + os.remove(path) diff --git a/tests/diffusion/models/helios/test_helios_usp_shard.py b/tests/diffusion/models/helios/test_helios_usp_shard.py new file mode 100644 index 000000000..88301a851 --- /dev/null +++ b/tests/diffusion/models/helios/test_helios_usp_shard.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for Helios per-component sequence splitting. + +These tests verify the manual sharding logic that replaces the broken +rope/blocks.0 hooks in the Helios ``_sp_plan``. They exercise the *real* +``HeliosTransformer3DModel._sp_split_seq`` (the parallel-state getters it +imports lazily inside the method body are monkeypatched), so the model-side +helper cannot drift silently from what is tested. For the multi-rank +mechanism around the helper (manual ``_sp_shard_depth`` bump, attn1/attn2 +parallel strategy, ws=2 vs ws=1 output equivalence) see +``test_helios_usp_2rank.py``. + +Coverage: + 1. USP disabled (ws == 1) returns the tensor unchanged. + 2. Divisible split gives each rank the expected slice. + 3. Non-divisible split pads with replicated tokens (no crash). + 4. Per-component split: history + current tokens are distributed so + that *both* ranks receive some current tokens (no mosaic regression). + 5. Edge cases: 1-D tensor, empty seq, 4-D tensor. +""" + +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +import torch + +import vllm_omni.diffusion.distributed.parallel_state as sp_parallel_state +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@contextmanager +def _fake_sp_state(ws: int, rank: int) -> Iterator[None]: + """Patch the parallel-state getters used by the real helper. + + ``_sp_split_seq`` imports these names from the module *inside* its + function body, so replacing the module attributes takes effect. + """ + prev_ws = sp_parallel_state.get_sequence_parallel_world_size + prev_rank = sp_parallel_state.get_sequence_parallel_rank + sp_parallel_state.get_sequence_parallel_world_size = lambda: ws + sp_parallel_state.get_sequence_parallel_rank = lambda: rank + try: + yield + finally: + sp_parallel_state.get_sequence_parallel_world_size = prev_ws + sp_parallel_state.get_sequence_parallel_rank = prev_rank + + +def _sp_split_seq(x: torch.Tensor, ws: int, rank: int) -> torch.Tensor: + """Call the real model-side helper without constructing a model.""" + with _fake_sp_state(ws, rank): + return HeliosTransformer3DModel._sp_split_seq(None, x) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestSpSplitSeqDisabled: + """When USP is off (ws == 1) the helper must be a no-op.""" + + def test_ws1_returns_unchanged(self): + x = torch.randn(2, 100, 64) + out = _sp_split_seq(x, ws=1, rank=0) + assert torch.equal(out, x) + assert out is x + + +class TestSpSplitSeqDivisible: + """Divisible sequences: each rank gets exactly seq_len / ws tokens.""" + + def test_each_rank_gets_correct_slice(self): + for ws in [2, 4, 8]: + seq_len = ws * 10 + x = torch.arange(seq_len * 4).float().reshape(1, seq_len, 4) + for rank in range(ws): + out = _sp_split_seq(x, ws=ws, rank=rank) + n = seq_len // ws + expected = x[:, rank * n : (rank + 1) * n, :] + assert out.shape == (1, n, 4) + assert torch.equal(out, expected) + + def test_output_is_contiguous(self): + x = torch.randn(1, 20, 8) + out = _sp_split_seq(x, ws=2, rank=0) + assert out.is_contiguous() + + +class TestSpSplitSeqNonDivisible: + """Non-divisible sequences are padded with replicated tokens.""" + + def test_non_divisible_pads_instead_of_crash(self): + """400x272 → 9*25*17 = 3825 tokens, USP=2 → not divisible. + Should pad to 3826 and split, not crash.""" + x = torch.randn(1, 3825, 8) + for rank in range(2): + out = _sp_split_seq(x, ws=2, rank=rank) + # 3825 → pad 1 → 3826 → /2 = 1913 + assert out.shape[1] == 1913 + + def test_non_divisible_pads_with_replicated_last_token(self): + """Verify the padding is the replicated last token, not zeros.""" + x = torch.randn(1, 7, 4) # 7 % 2 = 1, pad 1 + last_token = x[:, -1:, :].clone() + out = _sp_split_seq(x, ws=2, rank=0) + # rank 0 gets tokens [0:4] (first 4 of the padded 8) + assert out.shape == (1, 4, 4) + # rank 1 gets tokens [4:8] — token 7 (the padded one) = last token + out1 = _sp_split_seq(x, ws=2, rank=1) + assert out1.shape == (1, 4, 4) + assert torch.equal(out1[:, -1:, :], last_token) + + def test_non_divisible_7_4(self): + """7 tokens, ws=4 → pad 1 → 8 → /4 = 2 per rank.""" + x = torch.randn(1, 7, 8) + for rank in range(4): + out = _sp_split_seq(x, ws=4, rank=rank) + assert out.shape[1] == 2 + + def test_non_divisible_15_4(self): + """15 tokens, ws=4 → pad 1 → 16 → /4 = 4 per rank.""" + x = torch.randn(1, 15, 8) + for rank in range(4): + out = _sp_split_seq(x, ws=4, rank=rank) + assert out.shape[1] == 4 + + +class TestSpSplitSeqPerComponent: + """Simulate Helios forward: split current + history components so that + *both* ranks receive some current tokens (no mosaic regression).""" + + def test_both_ranks_get_current_tokens(self): + """history=2400, current=540, ws=2. + After per-component split, each rank must have > 0 current tokens.""" + current = torch.randn(1, 540, 64) + history = torch.randn(1, 2400, 64) + + for rank in range(2): + cur_local = _sp_split_seq(current, ws=2, rank=rank) + hist_local = _sp_split_seq(history, ws=2, rank=rank) + + assert cur_local.shape[1] > 0, f"Rank {rank} received 0 current tokens" + assert cur_local.shape[1] == 270 # 540 / 2 + assert hist_local.shape[1] == 1200 # 2400 / 2 + + original_context_length = cur_local.shape[1] + assert original_context_length > 0 # no mosaic regression + + def test_gather_restores_full_current(self): + """After splitting current into 2 ranks, simulating proj_out + gather must restore the full current sequence length so that + the original unpatchify reshape succeeds.""" + current = torch.randn(1, 540, 64) + local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] + gathered = torch.cat(local_chunks, dim=1) + assert gathered.shape[1] == 540 # full current restored + + def test_gather_with_padding_truncates_correctly(self): + """Non-divisible current (3825) → pad to 3826 → split → + gather → 3826. Slice to 3825 for unpatchify.""" + current = torch.randn(1, 3825, 64) + local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] + gathered = torch.cat(local_chunks, dim=1) + assert gathered.shape[1] == 3826 # padded size + # Simulate the slice before unpatchify + expected_seq = 3825 + if gathered.shape[1] > expected_seq: + gathered = gathered[:, :expected_seq, :] + assert gathered.shape[1] == 3825 # correct after slice + + def test_rotary_emb_aligned_with_hidden_states(self): + """rotary_emb and hidden_states must be split at the same + positions so that tokens align inside the attention layer.""" + hidden = torch.randn(1, 540, 5120) + rotary = torch.randn(1, 540, 256) + for rank in range(2): + h_local = _sp_split_seq(hidden, ws=2, rank=rank) + r_local = _sp_split_seq(rotary, ws=2, rank=rank) + assert h_local.shape[1] == r_local.shape[1] + + def test_whole_split_causes_mosaic_but_per_component_does_not(self): + """Demonstrates why per-component split is needed: + with whole split (simulated), rank 0 gets 0 current tokens; + with per-component split, both ranks get current tokens.""" + history = torch.randn(1, 2400, 64) + current = torch.randn(1, 540, 64) + full = torch.cat([history, current], dim=1) # 2940 + + # Whole split (the OLD broken behavior) + whole_n = full.shape[1] // 2 # 1470 + rank0_whole = full[:, :whole_n, :] + rank0_current = rank0_whole.shape[1] - 2400 # negative → 0 current + assert rank0_current <= 0, "Whole split should give rank 0 no current" + + # Per-component split (the NEW correct behavior) + for rank in range(2): + cur_local = _sp_split_seq(current, ws=2, rank=rank) + assert cur_local.shape[1] == 270 # both ranks have current + + +class TestSpSplitSeqEdgeCases: + """Edge cases that should not crash or silently misbehave.""" + + def test_1d_tensor_skipped(self): + """1-D tensors must be returned unchanged (dim < 2 guard).""" + x = torch.randn(100) + out = _sp_split_seq(x, ws=2, rank=0) + assert torch.equal(out, x) + + def test_empty_seq_skipped(self): + """Empty sequence dimension must be returned unchanged.""" + x = torch.randn(1, 0, 64) + out = _sp_split_seq(x, ws=2, rank=0) + assert out.shape == x.shape + + def test_4d_tensor_split_on_dim1(self): + """4-D tensor [B, S, H, D] should split along dim=1.""" + x = torch.randn(1, 20, 8, 32) + out = _sp_split_seq(x, ws=2, rank=1) + assert out.shape == (1, 10, 8, 32) diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index 4e3cf5078..6c6ec223d 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/vllm_omni/diffusion/models/helios/helios_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from Helios (https://github.com/BestWishYsh/Helios) import math @@ -27,10 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.distributed.sp_plan import ( - SequenceParallelInput, - SequenceParallelOutput, -) +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelOutput if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -448,6 +445,12 @@ class HeliosCrossAttention(nn.Module): num_kv_heads=self.num_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, + # Text K/V is replicated on every SP rank, so Ulysses must not be + # used here: its all-to-all would treat the per-rank replicas as + # sequence shards and ws-fold duplicate the keys. Sharded Q over + # full replicated K/V is correct locally and needs no + # communication (same as Wan2.2 cross-attn). + skip_sequence_parallel=True, ) def project_kv(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -624,16 +627,50 @@ class HeliosTransformer3DModel(nn.Module): _hsdp_shard_conditions = [_is_transformer_block] _sp_plan = { - "rope": { - 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), - 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), - }, - "blocks.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), - }, + # "rope" and "blocks.0" are intentionally omitted. + # The rope hook splits along dim=1 of the 5D output [B, D, T, H, W], + # which is the frequency dimension, NOT the sequence dimension. + # The blocks.0 hook splits the concatenated hidden_states as a whole, + # which can put all history tokens in one rank (original_context_length=0). + # Instead, hidden_states and rotary_emb are split per-component in + # forward() after flatten+transpose, ensuring each rank gets half + # of each component (history and current). "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } + def _sp_split_seq(self, x: torch.Tensor) -> torch.Tensor: + """Split tensor along sequence dim (dim=1) for Ulysses SP. + + If seq_len is not divisible by world_size, replicate the last + token(s) to pad to a divisible length. This avoids crashing + on non-divisible history components (which use different patch + sizes and thus produce different token counts). + + Residual caveat: the replicated padding tokens attend unmasked + in attn1 on the last rank (at most ws-1 replicas per component). + The impact is bounded, but this is the first place to look if + odd resolutions ever show quality drift. + """ + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + ) + + ws = get_sequence_parallel_world_size() + if ws > 1 and x.dim() >= 2 and x.shape[1] > 0: + seq_len = x.shape[1] + remainder = seq_len % ws + if remainder != 0: + pad = ws - remainder + # Replicate the last token to fill the padding + last = x[:, -1:, ...].expand(-1, pad, *([-1] * (x.dim() - 2))) + x = torch.cat([x, last], dim=1) + seq_len = x.shape[1] + r = get_sequence_parallel_rank() + n = seq_len // ws + x = x[:, r * n : (r + 1) * n, ...].contiguous() + return x + def __init__( self, patch_size: tuple[int, ...] = (1, 2, 2), @@ -860,6 +897,9 @@ class HeliosTransformer3DModel(nn.Module): device=hidden_states.device, ) rotary_emb = rotary_emb.flatten(2).transpose(1, 2) + # USP: per-component split (each rank gets half of current frames) + hidden_states = self._sp_split_seq(hidden_states) + rotary_emb = self._sp_split_seq(rotary_emb) original_context_length = hidden_states.shape[1] # 2. Process short history latents @@ -876,6 +916,9 @@ class HeliosTransformer3DModel(nn.Module): device=latents_history_short.device, ) rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_short = self._sp_split_seq(latents_history_short) + rotary_emb_history_short = self._sp_split_seq(rotary_emb_history_short) hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) @@ -896,6 +939,9 @@ class HeliosTransformer3DModel(nn.Module): rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_mid = self._sp_split_seq(latents_history_mid) + rotary_emb_history_mid = self._sp_split_seq(rotary_emb_history_mid) hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) @@ -916,6 +962,9 @@ class HeliosTransformer3DModel(nn.Module): rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_long = self._sp_split_seq(latents_history_long) + rotary_emb_history_long = self._sp_split_seq(rotary_emb_history_long) hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) @@ -958,6 +1007,18 @@ class HeliosTransformer3DModel(nn.Module): timestep_proj = timestep_proj.permute(0, 2, 1, 3) # 6. Transformer blocks + # Manually increment _sp_shard_depth so that attention layers + # know SP is active and perform Ulysses All-to-All communication. + # Normally this is done by split_output=True hooks (e.g. rope), + # but those are removed because they split the wrong dimension. + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, + ) + + if is_forward_context_available(): + get_forward_context()._sp_shard_depth += 1 + hidden_states = hidden_states.contiguous() encoder_hidden_states = encoder_hidden_states.contiguous() rotary_emb = rotary_emb.contiguous() @@ -979,6 +1040,11 @@ class HeliosTransformer3DModel(nn.Module): hidden_states = self.proj_out(hidden_states) # 8. Unpatchify + # proj_out gather may include padded tokens from non-divisible + # sequences. Slice to the expected size before reshape. + expected_seq = post_patch_num_frames * post_patch_height * post_patch_width + if hidden_states.shape[1] > expected_seq: + hidden_states = hidden_states[:, :expected_seq, :] hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 )