From 13c617a1407e59614fb7a7c02507eac9398b5c75 Mon Sep 17 00:00:00 2001 From: vOv <124694036+cr-gao@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:22:53 +0800 Subject: [PATCH] [Model] Add sequence parallelism to SANA-Video 2B (#5940) Signed-off-by: cr-gao --- docs/user_guide/diffusion_features.md | 2 +- recipes/NVIDIA/SANA-Video-2B.md | 43 +- .../distributed/test_sana_video_sp_numeric.py | 357 +++++++++++ tests/diffusion/models/sana_video/conftest.py | 13 + .../sana_video/test_sana_video_topology.py | 41 +- .../test_transformer_sana_video_sp.py | 567 ++++++++++++++++++ tests/model_tests/diffusion/model_settings.py | 5 + .../models/sana_video/pipeline_sana_video.py | 3 + .../sana_video/transformer_sana_video.py | 119 +++- 9 files changed, 1122 insertions(+), 28 deletions(-) create mode 100644 tests/diffusion/distributed/test_sana_video_sp_numeric.py create mode 100644 tests/diffusion/models/sana_video/conftest.py create mode 100644 tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index a06f53d3e..7d443c942 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -164,7 +164,7 @@ The following tables show which models support each feature: | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ | -| **SANA-Video-2B T2V I2V** | ❌ | ❌ | ❌ | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **SANA-Video-2B T2V I2V** | ❌ | ❌ | ✅ (`--usp`, frame-sharded) | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | diff --git a/recipes/NVIDIA/SANA-Video-2B.md b/recipes/NVIDIA/SANA-Video-2B.md index f26692867..8dd3e8e04 100644 --- a/recipes/NVIDIA/SANA-Video-2B.md +++ b/recipes/NVIDIA/SANA-Video-2B.md @@ -121,8 +121,8 @@ bash examples/online_serving/text_to_video/run_curl_sana_video.sh ##### Parallel native serving -The native pipelines support tensor parallelism (up to 2 GPUs) and CFG -parallelism. Tensor parallelism on 2 GPUs: +The native pipelines support tensor parallelism (up to 2 GPUs), sequence +parallelism (2 or 4 GPUs) and CFG parallelism. Tensor parallelism on 2 GPUs: ```bash vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ @@ -133,15 +133,32 @@ vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ --port 8091 ``` -For CFG parallelism on 2 GPUs, replace `--tensor-parallel-size 2` with -`--cfg-parallel-size 2`; it splits the guided and unguided branches across the -two GPUs and only helps when `guidance_scale` is above 1. Passing both flags -combines the two on 4 GPUs. +Sequence parallelism on 2 GPUs (replace `--tensor-parallel-size 2` with): -Whether tensor parallelism lowers latency depends on the interconnect. Each -transformer block adds an all-reduce, so it speeds up generation only on fast -GPU links such as NVLink and can be slower than a single GPU on PCIe-only -systems. Measure on your hardware before enabling it. +```bash + --usp 2 +``` + +The `--usp` flag carries the sequence parallel degree. On SANA-Video the +implementation is not Ulysses all-to-all: video frames are sharded across +ranks and the linear-attention state is combined with one small all-reduce +per block, which is far less communication than tensor parallelism needs. +Ring and AllGather-KV sequence parallel do not apply to SANA's linear +attention and are rejected. + +For CFG parallelism on 2 GPUs, use `--cfg-parallel-size 2` instead; it splits +the guided and unguided branches across the two GPUs and only helps when +`guidance_scale` is above 1. CFG parallelism combines with either tensor +parallelism or sequence parallelism on 4 GPUs. Tensor and sequence parallelism +also combine; on the measured 4-GPU NVLink box tp2 x sp2 reached 2.55x, below +the pure sp4 (3.24x) and sp2 x cfg2 (3.56x) configurations, so prefer those +when four GPUs serve one request. + +Whether these lower latency depends on the interconnect. Each transformer +block adds an all-reduce under tensor parallelism, so it speeds up generation +only on fast GPU links such as NVLink and can be slower than a single GPU on +PCIe-only systems. Sequence parallelism communicates far less per block, but +was only measured on NVLink. Measure on your hardware before enabling either. To run the black-box compatibility backend for T2V, replace the server script with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request @@ -287,9 +304,9 @@ for the tuning knobs. The denoising loop intentionally retains the checkpoint-compatible Diffusers `DPMSolverMultistepScheduler`. - Known limitations: - - Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for - the native pipeline. Sequence parallelism, TeaCache, and step execution - are not validated. + - Tensor parallelism (up to 2 GPUs), sequence parallelism (2 or 4 GPUs) + and CFG parallelism are supported for the native pipeline. TeaCache and + step execution are not validated. - Cache-DiT and CPU offload are limited to TP1/CFG1/SP1. Cache-DiT with distributed layerwise offload is not supported by the native pipeline. - Cache-DiT speedup and offload memory numbers above are single-GPU A800 diff --git a/tests/diffusion/distributed/test_sana_video_sp_numeric.py b/tests/diffusion/distributed/test_sana_video_sp_numeric.py new file mode 100644 index 000000000..b18cafe4c --- /dev/null +++ b/tests/diffusion/distributed/test_sana_video_sp_numeric.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real multi-process numeric equivalence gate for SANA-Video sequence parallel. + +Runs the native transformer directly (no pipeline, no topology validation) with +real NCCL groups and compares SP>1 output against a same-weight SP=1 baseline. +In fp32 the only difference is reduction reordering, so any halo, gather, +timestep-slice or state-reduction misalignment fails by orders of magnitude. +""" + +from __future__ import annotations + +import os +import tempfile + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + get_classifier_free_guidance_rank, + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.platforms import current_omni_platform + +_TINY_CONFIG = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 12, + "num_layers": 2, + "num_cross_attention_heads": 2, + "cross_attention_head_dim": 12, + "cross_attention_dim": 24, + "caption_channels": 8, + "mlp_ratio": 2.0, + "patch_size": (1, 2, 2), + "sample_size": 8, + "rope_max_seq_len": 512, +} + +_MODEL_SEED = 42 +_INPUT_SEED = 123 +_FP32_REL_L2 = 1e-5 +# bf16 collectives only reorder low-precision sums; informational sanity bound. +_BF16_REL_L2 = 5e-2 + + +def _make_omni_config(sp_size: int, cfg_size: int, dtype: torch.dtype, tp_size: int = 1) -> OmniDiffusionConfig: + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=tp_size, + sequence_parallel_size=sp_size, + ulysses_degree=sp_size, + ring_degree=1, + cfg_parallel_size=cfg_size, + ) + return OmniDiffusionConfig( + model="test", + dtype=dtype, + parallel_config=parallel_config, + # Backend-agnostic gate; SDPA also avoids platform quirks of the + # default backend on older GPUs. + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")), + ) + + +def _build_model(device: torch.device, dtype: torch.dtype): + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + torch.manual_seed(_MODEL_SEED) + model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval() + for _, param in sorted(model.named_parameters()): + torch.nn.init.normal_(param, mean=0.0, std=0.02) + return model.to(device=device, dtype=dtype) + + +def _make_inputs(task: str, num_frames: int, input_seed: int): + generator = torch.Generator().manual_seed(input_seed) + latent = torch.randn(2, 4, num_frames, 8, 8, generator=generator) + encoder_hidden_states = torch.randn(2, 6, 8, generator=generator) + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1]]) + if task == "t2v": + timestep = torch.tensor([500.0, 700.0]) + else: + timestep = torch.rand(2, 1, num_frames, 4, 4, generator=generator) * 1000.0 + return latent, encoder_hidden_states, encoder_attention_mask, timestep + + +def _worker( + local_rank: int, + world_size: int, + sp_size: int, + cfg_size: int, + tp_size: int, + dtype: torch.dtype, + task: str, + num_frames: int, + weights_path: str, + port: int, + result_queue, +): + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + os.environ.update( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": str(port), + } + ) + + init_distributed_environment() + initialize_model_parallel( + cfg_parallel_size=cfg_size, + sequence_parallel_size=sp_size, + ulysses_degree=sp_size, + tensor_parallel_size=tp_size, + ) + assert get_sequence_parallel_world_size() == sp_size + + cfg_rank = get_classifier_free_guidance_rank() if cfg_size > 1 else 0 + od_config = _make_omni_config(sp_size, cfg_size, dtype, tp_size) + with set_current_diffusion_config(od_config): + model = _build_model(device, dtype) + # Load the baseline's full weights through the real weight loaders so TP + # ranks hold proper shards; per-rank random init is not TP-consistent. + model.load_weights(iter(torch.load(weights_path, map_location="cpu").items())) + # Each CFG branch gets distinct inputs so a CFG/SP group mix-up cannot + # cancel out. + latent, encoder_hidden_states, encoder_attention_mask, timestep = _make_inputs( + task, num_frames, _INPUT_SEED + cfg_rank + ) + + with torch.no_grad(): + output = model( + latent.to(device=device, dtype=dtype), + encoder_hidden_states.to(device=device, dtype=dtype), + timestep.to(device=device), + encoder_attention_mask=encoder_attention_mask.to(device=device), + return_dict=False, + )[0] + + result_queue.put((cfg_rank, get_sequence_parallel_rank(), output.float().cpu())) + destroy_distributed_env() + + +def _rel_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + return ((actual - expected).double().norm() / expected.double().norm()).item() + + +def _spawn(worker_args: tuple, nprocs: int, result_queue) -> list: + torch.multiprocessing.spawn(_worker, args=(*worker_args, result_queue), nprocs=nprocs) + return [result_queue.get() for _ in range(nprocs)] + + +def _run_sp_parity( + *, + sp_size: int, + cfg_size: int, + dtype: torch.dtype, + task: str, + num_frames: int, + rel_l2_bound: float, + tp_size: int = 1, +) -> float: + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + + baselines = {} + with tempfile.TemporaryDirectory() as tmpdir: + weights_path = os.path.join(tmpdir, "baseline_weights.pt") + for cfg_rank in range(cfg_size): + queue = manager.Queue() + # A world-1 run sees input seed _INPUT_SEED + 0; shift the seed via + # the task inputs by spawning one baseline per CFG branch input. + torch.multiprocessing.spawn( + _baseline_worker, + args=( + dtype, + task, + num_frames, + _INPUT_SEED + cfg_rank, + weights_path if cfg_rank == 0 else "", + 29531, + queue, + ), + nprocs=1, + ) + baselines[cfg_rank] = queue.get() + + queue = manager.Queue() + world_size = sp_size * cfg_size * tp_size + results = _spawn( + (world_size, sp_size, cfg_size, tp_size, dtype, task, num_frames, weights_path, 29532), world_size, queue + ) + + worst = 0.0 + by_cfg: dict[int, list] = {} + for cfg_rank, sp_rank, output in results: + by_cfg.setdefault(cfg_rank, []).append((sp_rank, output)) + assert sorted(by_cfg) == list(range(cfg_size)) + for cfg_rank, rank_outputs in by_cfg.items(): + rank_outputs.sort(key=lambda item: item[0]) + first = rank_outputs[0][1] + for sp_rank, output in rank_outputs[1:]: + torch.testing.assert_close( + output, first, rtol=0, atol=0, msg=f"SP ranks disagree (cfg_rank={cfg_rank}, sp_rank={sp_rank})" + ) + drift = _rel_l2(first, baselines[cfg_rank]) + worst = max(worst, drift) + assert drift <= rel_l2_bound, ( + f"SP{sp_size} output drifted rel_l2={drift:.3e} from the SP1 baseline " + f"(bound {rel_l2_bound:.1e}, cfg_rank={cfg_rank}, task={task}, frames={num_frames})" + ) + return worst + + +def _baseline_worker( + local_rank: int, + dtype: torch.dtype, + task: str, + num_frames: int, + input_seed: int, + weights_path: str, + port: int, + result_queue, +): + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + os.environ.update( + { + "RANK": "0", + "LOCAL_RANK": "0", + "WORLD_SIZE": "1", + "MASTER_ADDR": "localhost", + "MASTER_PORT": str(port), + } + ) + + init_distributed_environment() + initialize_model_parallel() + + od_config = _make_omni_config(1, 1, dtype) + with set_current_diffusion_config(od_config): + model = _build_model(device, dtype) + if weights_path: + torch.save({name: param.detach().cpu() for name, param in model.named_parameters()}, weights_path) + latent, encoder_hidden_states, encoder_attention_mask, timestep = _make_inputs(task, num_frames, input_seed) + + with torch.no_grad(): + output = model( + latent.to(device=device, dtype=dtype), + encoder_hidden_states.to(device=device, dtype=dtype), + timestep.to(device=device), + encoder_attention_mask=encoder_attention_mask.to(device=device), + return_dict=False, + )[0] + + result_queue.put(output.float().cpu()) + destroy_distributed_env() + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=2) +@pytest.mark.parametrize("task", ["t2v", "i2v"]) +@pytest.mark.parametrize("num_frames", [21, 5]) +def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None: + _run_sp_parity( + sp_size=2, + cfg_size=1, + dtype=torch.float32, + task=task, + num_frames=num_frames, + rel_l2_bound=_FP32_REL_L2, + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("task", ["t2v", "i2v"]) +@pytest.mark.parametrize("num_frames", [21, 5]) +def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None: + _run_sp_parity( + sp_size=4, + cfg_size=1, + dtype=torch.float32, + task=task, + num_frames=num_frames, + rel_l2_bound=_FP32_REL_L2, + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=4) +def test_sp2_cfg2_matches_sp1_fp32() -> None: + """Four processes: CFG branches with distinct inputs each run their own SP + group; catches CFG/SP group crossover and deadlocks.""" + _run_sp_parity( + sp_size=2, + cfg_size=2, + dtype=torch.float32, + task="i2v", + num_frames=21, + rel_l2_bound=_FP32_REL_L2, + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=4) +def test_tp2_sp2_matches_sp1_fp32() -> None: + """Four processes: TP shards heads and channels while SP shards frames; the + combined mesh must still reproduce the serial output.""" + _run_sp_parity( + sp_size=2, + cfg_size=1, + tp_size=2, + dtype=torch.float32, + task="t2v", + num_frames=21, + rel_l2_bound=_FP32_REL_L2, + ) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=2) +def test_sp2_bf16_drift_informational() -> None: + """Records the bf16 reduction-reorder drift magnitude; the load-bearing + numeric gate is the fp32 suite above.""" + drift = _run_sp_parity( + sp_size=2, + cfg_size=1, + dtype=torch.bfloat16, + task="t2v", + num_frames=21, + rel_l2_bound=_BF16_REL_L2, + ) + print(f"\nSANA-Video SP2 bf16 drift vs SP1: rel_l2={drift:.3e}") diff --git a/tests/diffusion/models/sana_video/conftest.py b/tests/diffusion/models/sana_video/conftest.py new file mode 100644 index 000000000..06e5ba647 --- /dev/null +++ b/tests/diffusion/models/sana_video/conftest.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm_omni.diffusion.models.sana_video import transformer_sana_video + + +@pytest.fixture(autouse=True) +def sp_world_1(monkeypatch): + """L1 runs single-process without a sequence parallel group; keep the + transformer on its SP1 path unless a test mocks a larger SP world.""" + monkeypatch.setattr(transformer_sana_video, "get_sequence_parallel_world_size", lambda: 1) diff --git a/tests/diffusion/models/sana_video/test_sana_video_topology.py b/tests/diffusion/models/sana_video/test_sana_video_topology.py index ac0fa97e0..440de3be9 100644 --- a/tests/diffusion/models/sana_video/test_sana_video_topology.py +++ b/tests/diffusion/models/sana_video/test_sana_video_topology.py @@ -3,8 +3,10 @@ """Topology validation for SANA-Video parallelism. -The native pipeline only supports tensor/CFG parallel of size 1 or 2 and rejects -sequence parallel and text-encoder TP before any large weight is allocated. +The native pipeline only supports tensor/CFG parallel of size 1 or 2 and +sequence parallel of degree 1, 2 or 4 (Ulysses flag, frame-sharded manual +implementation); it rejects ring/AllGather-KV SP, the untested TP x SP combo +and text-encoder TP before any large weight is allocated. """ from types import SimpleNamespace @@ -24,6 +26,8 @@ def _pc( tp: int = 1, cfg: int = 1, sp: int | None = None, + ring: int = 1, + allgather: int = 1, text_tp: int = 1, pp: int = 1, hsdp: bool = False, @@ -32,6 +36,8 @@ def _pc( tensor_parallel_size=tp, cfg_parallel_size=cfg, sequence_parallel_size=sp, + ring_degree=ring, + allgather_degree=allgather, text_encoder_tp_size=text_tp, pipeline_parallel_size=pp, use_hsdp=hsdp, @@ -54,9 +60,30 @@ def test_cfg3_rejected() -> None: validate_sana_video_parallel_config(_pc(cfg=3)) -def test_sequence_parallel_rejected() -> None: - with pytest.raises(NotImplementedError, match="[Ss]equence parallel"): - validate_sana_video_parallel_config(_pc(sp=2)) +@pytest.mark.parametrize("sp", [None, 1, 2, 4]) +def test_supported_sp_sizes_pass(sp: int | None) -> None: + validate_sana_video_parallel_config(_pc(sp=sp)) + + +@pytest.mark.parametrize("sp", [3, 8]) +def test_unsupported_sp_sizes_rejected(sp: int) -> None: + with pytest.raises(NotImplementedError, match="sequence_parallel_size"): + validate_sana_video_parallel_config(_pc(sp=sp)) + + +def test_ring_sequence_parallel_rejected() -> None: + with pytest.raises(NotImplementedError, match="ring"): + validate_sana_video_parallel_config(_pc(sp=2, ring=2)) + + +def test_allgather_sequence_parallel_rejected() -> None: + with pytest.raises(NotImplementedError, match="AllGather"): + validate_sana_video_parallel_config(_pc(sp=2, allgather=2)) + + +@pytest.mark.parametrize("sp", [2, 4]) +def test_tp2_sp_combos_pass(sp: int) -> None: + validate_sana_video_parallel_config(_pc(tp=2, sp=sp)) def test_pipeline_parallel_rejected() -> None: @@ -78,6 +105,10 @@ def test_tp2_cfg2_combo_passes() -> None: validate_sana_video_parallel_config(_pc(tp=2, cfg=2)) +def test_sp2_cfg2_combo_passes() -> None: + validate_sana_video_parallel_config(_pc(sp=2, cfg=2)) + + @pytest.mark.parametrize("pipeline_cls", [SanaVideoPipeline, SanaImageToVideoPipeline]) def test_pipeline_rejects_unsupported_topology_before_loading(pipeline_cls, mocker) -> None: """The pipeline must reject an unsupported topology before it loads weights.""" diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py new file mode 100644 index 000000000..3f4c62075 --- /dev/null +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py @@ -0,0 +1,567 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""SP unit tests for SANA-Video: balanced frame split sizes, linear-attention +state reduction and the uneven frame gather. + +CPU-only. The SP world size, rank and the collectives are mocked so a single +process simulates the shard/reduce logic a real multi-rank run depends on. +""" + +import os +import threading +from concurrent.futures import ThreadPoolExecutor + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + GLUMBTempConv, + SanaAdaLayerNormSingle, + SanaLinearAttention, + _sp_frame_split_sizes, + _sp_gather_frames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_MODULE = "vllm_omni.diffusion.models.sana_video.transformer_sana_video" + + +@pytest.fixture +def tp1_group(): + """Real single-process TP group so the parallel linear layers construct and + run on CPU at tensor_parallel_size=1.""" + from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29501") + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") + initialize_model_parallel() + yield + cleanup_dist_env_and_memory() + + +@pytest.fixture +def force_default_gemm(monkeypatch): + """Force CPU-compatible GEMM dispatch for the parallel linear layers.""" + from vllm.model_executor.layers.utils import default_unquantized_gemm + + monkeypatch.setattr( + "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", + lambda: default_unquantized_gemm, + ) + + +# ── frame split sizes ── + + +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("num_frames", [5, 6, 9, 11, 21]) +def test_frame_split_sizes_are_balanced_and_cover_all_frames(num_frames, world_size) -> None: + """Exactly world_size non-empty near-equal chunks that sum to num_frames. + + torch.chunk's ceil semantics produce fewer than world_size chunks for + 5/6/9 frames at world size 4, which is why this helper exists. + """ + sizes = _sp_frame_split_sizes(num_frames, world_size) + + assert len(sizes) == world_size + assert all(size >= 1 for size in sizes) + assert sum(sizes) == num_frames + assert max(sizes) - min(sizes) <= 1 + + +def test_frame_split_sizes_known_layouts() -> None: + assert _sp_frame_split_sizes(21, 2) == [11, 10] + assert _sp_frame_split_sizes(11, 2) == [6, 5] + assert _sp_frame_split_sizes(9, 4) == [3, 2, 2, 2] + assert _sp_frame_split_sizes(5, 4) == [2, 1, 1, 1] + + +def test_frame_split_sizes_rejects_fewer_frames_than_ranks() -> None: + with pytest.raises(ValueError, match="latent frame per rank"): + _sp_frame_split_sizes(3, 4) + + +# ── linear attention state reduction ── + + +class _FakeSpAllReduce: + """Two-pass stand-in for the SP all-reduce: pass 1 records each rank's + partial, pass 2 replays their sum to every rank.""" + + def __init__(self) -> None: + self.partials: list[torch.Tensor] = [] + self.total: torch.Tensor | None = None + + def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor: + self.partials.append(tensor.clone()) + if self.total is not None: + return self.total.clone() + return tensor + + +def _make_attn(dim: int, num_heads: int, head_dim: int) -> SanaLinearAttention: + attn = SanaLinearAttention( + dim=dim, num_heads=num_heads, head_dim=head_dim, dropout=0.0, bias=False, qk_norm="rms_norm_across_heads" + ) + attn.eval() + torch.manual_seed(0) + for _, param in sorted(attn.named_parameters()): + param.data.normal_() + return attn + + +def _attn_inputs(seq_len: int, dim: int, head_dim: int): + """Second half of the sequence drawn from a shifted, scaled distribution so + per-rank partial sums differ; a missing state reduction cannot cancel out.""" + torch.manual_seed(1) + x = torch.randn(2, seq_len, dim) + x[:, seq_len // 2 :] = x[:, seq_len // 2 :] * 3.0 + 1.0 + freqs_cos = torch.randn(1, seq_len, 1, head_dim) + freqs_sin = torch.randn(1, seq_len, 1, head_dim) + return x, (freqs_cos, freqs_sin) + + +def test_attn1_sp2_state_reduction_matches_dense(tp1_group, force_default_gemm, mocker) -> None: + """Two token shards with a summed packed state must reproduce the dense + full-sequence linear attention output.""" + dim, num_heads, head_dim, seq_len = 24, 4, 6, 12 + attn = _make_attn(dim, num_heads, head_dim) + x, (freqs_cos, freqs_sin) = _attn_inputs(seq_len, dim, head_dim) + bounds = [(0, 8), (8, 12)] + + with torch.no_grad(): + ref = attn(x, rotary_emb=(freqs_cos, freqs_sin)) + + fake_group = _FakeSpAllReduce() + mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=2) + mocker.patch(f"{_MODULE}.get_sp_group", return_value=fake_group) + + def run_ranks() -> list[torch.Tensor]: + outs = [] + for start, stop in bounds: + local_rotary = (freqs_cos[:, start:stop], freqs_sin[:, start:stop]) + with torch.no_grad(): + outs.append(attn(x[:, start:stop], rotary_emb=local_rotary)) + return outs + + run_ranks() + fake_group.total = fake_group.partials[0] + fake_group.partials[1] + out = torch.cat(run_ranks(), dim=1) + + torch.testing.assert_close(out, ref, rtol=1e-5, atol=1e-5) + + +def test_attn1_issues_single_packed_all_reduce(tp1_group, force_default_gemm, mocker) -> None: + """One collective per forward, carrying scores and k_sum packed together.""" + dim, num_heads, head_dim, seq_len = 24, 4, 6, 8 + attn = _make_attn(dim, num_heads, head_dim) + x, rotary_emb = _attn_inputs(seq_len, dim, head_dim) + + fake_group = _FakeSpAllReduce() + mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=2) + mocker.patch(f"{_MODULE}.get_sp_group", return_value=fake_group) + + with torch.no_grad(): + attn(x, rotary_emb=rotary_emb) + + assert len(fake_group.partials) == 1 + packed = fake_group.partials[0] + assert packed.shape == (2, num_heads, head_dim, head_dim + 1) + assert packed.dtype == torch.float32 + + +def test_attn1_sp1_path_is_untouched(tp1_group, force_default_gemm, mocker) -> None: + """At SP1 no collective may run: the dense operator sequence must not + change shape or route through get_sp_group at all.""" + dim, num_heads, head_dim, seq_len = 24, 4, 6, 8 + attn = _make_attn(dim, num_heads, head_dim) + x, rotary_emb = _attn_inputs(seq_len, dim, head_dim) + + mock_group = mocker.patch(f"{_MODULE}.get_sp_group") + + with torch.no_grad(): + attn(x, rotary_emb=rotary_emb) + + mock_group.assert_not_called() + + +# ── uneven frame gather ── + + +class _FakeSpAllGather: + """Stand-in for the SP all-gather: returns precomputed equal-shape per-rank + contributions, enforcing the equal-shard contract of the real collective.""" + + def __init__(self, parts: list[torch.Tensor]) -> None: + self.parts = parts + + def all_gather(self, tensor: torch.Tensor, dim: int = 0, separate_tensors: bool = False): + assert separate_tensors, "frame gather must request per-rank tensors" + assert all(part.shape == tensor.shape for part in self.parts), "all_gather requires equal shards" + return [part.clone() for part in self.parts] + + +@pytest.mark.parametrize("num_frames,world_size", [(21, 2), (6, 2), (5, 4)]) +def test_gather_frames_roundtrips_uneven_shards(num_frames, world_size, mocker) -> None: + """shard -> pad -> all_gather -> per-rank narrow -> cat must reproduce the + full sequence bitwise, with communication pads never reaching the output.""" + batch, hw, channels = 2, 3, 4 + torch.manual_seed(0) + full = torch.randn(batch, num_frames * hw, channels) + + sizes = _sp_frame_split_sizes(num_frames, world_size) + frame_view = full.unflatten(1, (num_frames, hw)) + locals_ = [t.flatten(1, 2) for t in frame_view.split(sizes, dim=1)] + + # Per-rank padded contributions as the real collective would deliver them; + # NaN pads prove a leaked pad frame cannot go unnoticed. + parts = [] + for local, size in zip(locals_, sizes): + part = local.unflatten(1, (size, hw)) + pad = max(sizes) - size + if pad: + part = torch.cat([part, torch.full((batch, pad, hw, channels), torch.nan)], dim=1) + parts.append(part) + + mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=world_size) + mocker.patch(f"{_MODULE}.get_sp_group", return_value=_FakeSpAllGather(parts)) + for rank in range(world_size): + mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", return_value=rank) + + gathered = _sp_gather_frames(locals_[rank], sizes) + + torch.testing.assert_close(gathered, full, rtol=0, atol=0) + + +# ── GLUMB conv_temp halo exchange ── + + +class _NoopReq: + def wait(self) -> None: + pass + + +class _ReplayP2PDist: + """Two-pass stand-in for torch.distributed P2P: pass 1 records each rank's + boundary sends, pass 2 replays them into the matching neighbour recvs.""" + + irecv = "irecv" + isend = "isend" + + def __init__(self) -> None: + self.rank = 0 + self.sent: dict[tuple[int, int], torch.Tensor] = {} + self.replay = False + self.batches = 0 + self.peers: set[tuple[int, int]] = set() + + def P2POp(self, op, tensor, peer, group): + return (op, tensor, peer) + + def batch_isend_irecv(self, ops): + self.batches += 1 + for op, tensor, peer in ops: + self.peers.add((self.rank, peer)) + if op == "isend": + self.sent[(self.rank, peer)] = tensor.clone() + elif self.replay: + tensor.copy_(self.sent[(peer, self.rank)]) + return [_NoopReq() for _ in ops] + + +class _FakeSpGroup: + def __init__(self, world_size: int) -> None: + self.ranks = list(range(world_size)) + self.device_group = None + + +def _make_glumb(mocker) -> GLUMBTempConv: + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1) + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=0) + torch.manual_seed(0) + glumb = GLUMBTempConv(4, 4, 2.0, norm_type=None, residual_connection=False) + glumb.eval() + for _, param in sorted(glumb.named_parameters()): + param.data.normal_() + return glumb + + +@pytest.mark.parametrize("world_size", [2, 4]) +def test_glumb_conv_temp_sp_halo_matches_dense(world_size, mocker) -> None: + """Frame shards with neighbour boundary frames exchanged must reproduce the + dense temporal conv, including zero padding at the global first/last frame. + The exchange must talk to the two neighbouring ranks only.""" + glumb = _make_glumb(mocker) + num_frames = 5 + torch.manual_seed(1) + x = torch.randn(2, num_frames, 3, 3, 4) + + with torch.no_grad(): + ref = glumb(x) + + sizes = _sp_frame_split_sizes(num_frames, world_size) + shards = [shard.contiguous() for shard in x.split(sizes, dim=1)] + fake_dist = _ReplayP2PDist() + mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=world_size) + mocker.patch(f"{_MODULE}.get_sp_group", return_value=_FakeSpGroup(world_size)) + mocker.patch(f"{_MODULE}.dist", fake_dist) + + def run_ranks() -> list[torch.Tensor]: + outs = [] + for rank in range(world_size): + fake_dist.rank = rank + mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", return_value=rank) + with torch.no_grad(): + outs.append(glumb(shards[rank])) + return outs + + run_ranks() + assert fake_dist.batches == world_size + assert fake_dist.peers == {(a, b) for a in range(world_size) for b in (a - 1, a + 1) if 0 <= b < world_size} + fake_dist.replay = True + out = torch.cat(run_ranks(), dim=1) + + torch.testing.assert_close(out, ref, rtol=1e-5, atol=1e-5) + + +def test_glumb_sp_halo_rejects_non_three_frame_kernel(mocker) -> None: + """The single-frame halo is only correct for a three-frame temporal kernel; + any other width must fail loudly instead of corrupting parity.""" + glumb = _make_glumb(mocker) + glumb.conv_temp = nn.Conv2d(4, 4, kernel_size=(5, 1), stride=1, padding=(2, 0), bias=False) + mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=2) + mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", return_value=0) + mocker.patch(f"{_MODULE}.get_sp_group", return_value=_FakeSpGroup(2)) + mocker.patch(f"{_MODULE}.dist", _ReplayP2PDist()) + with pytest.raises(AssertionError, match="temporal kernel"): + with torch.no_grad(): + glumb(torch.randn(1, 3, 3, 3, 4)) + + +# ── I2V per-token timestep frame slicing ── + + +def test_i2v_timestep_frame_slice_matches_dense_token_region() -> None: + """Slicing the 5D timestep along its frame dim before the per-token + time-embed MLP must match the dense modulation for the same token region. + + The sinusoidal projection is bitwise under slicing; the MLP is only equal + up to CPU GEMM blocking, which depends on the row count. A wrongly sliced + frame produces O(1) differences, far outside this tolerance.""" + dim = 24 + torch.manual_seed(0) + layer = SanaAdaLayerNormSingle(dim).eval() + + batch, num_frames, height, width = 2, 5, 2, 3 + timestep = torch.rand(batch, 1, num_frames, height, width) * 1000.0 + + with torch.no_grad(): + dense_mod, dense_emb = layer(timestep.flatten(), batch_size=batch, hidden_dtype=torch.float32) + dense_mod = dense_mod.view(batch, -1, dense_mod.size(-1)) + dense_emb = dense_emb.view(batch, -1, dense_emb.size(-1)) + + sizes = _sp_frame_split_sizes(num_frames, 2) + start = 0 + for rank, size in enumerate(sizes): + local = timestep.split(sizes, dim=2)[rank] + with torch.no_grad(): + local_mod, local_emb = layer(local.flatten(), batch_size=batch, hidden_dtype=torch.float32) + local_mod = local_mod.view(batch, -1, local_mod.size(-1)) + local_emb = local_emb.view(batch, -1, local_emb.size(-1)) + + lo, hi = start * height * width, (start + size) * height * width + torch.testing.assert_close(local_mod, dense_mod[:, lo:hi], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(local_emb, dense_emb[:, lo:hi], rtol=1e-5, atol=1e-6) + start += size + + +# ── full-model lockstep parity ── + + +_TINY_SP_CONFIG = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 12, + "num_layers": 2, + "num_cross_attention_heads": 2, + "cross_attention_head_dim": 12, + "cross_attention_dim": 24, + "caption_channels": 8, + "mlp_ratio": 2.0, + "patch_size": (1, 2, 2), + "sample_size": 4, + "rope_max_seq_len": 64, +} + + +class _LockstepSpGroup: + """Barrier-synchronized collectives so rank threads run true lockstep SP in + one process; records per-call inputs so tests can prove ranks computed on + different shards.""" + + def __init__(self, world_size: int, tls: threading.local) -> None: + self.world_size = world_size + self.ranks = list(range(world_size)) + self.device_group = None + self.tls = tls + self.barrier = threading.Barrier(world_size) + self.slots: list[torch.Tensor | None] = [None] * world_size + self.reduce_inputs: list[list[torch.Tensor]] = [] + self.gather_calls = 0 + self.lock = threading.Lock() + + def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor: + self.slots[self.tls.rank] = tensor + self.barrier.wait(timeout=60) + total = self.slots[0].clone() + for part in self.slots[1:]: + total = total + part + if self.tls.rank == 0: + self.reduce_inputs.append([part.clone() for part in self.slots]) + self.barrier.wait(timeout=60) + return total + + def all_gather(self, tensor: torch.Tensor, dim: int = 0, separate_tensors: bool = False): + assert separate_tensors + self.slots[self.tls.rank] = tensor + self.barrier.wait(timeout=60) + assert all(part.shape == tensor.shape for part in self.slots), "all_gather requires equal shards" + parts = [part.clone() for part in self.slots] + if self.tls.rank == 0: + with self.lock: + self.gather_calls += 1 + self.barrier.wait(timeout=60) + return parts + + +class _LockstepP2PDist: + """Barrier-synchronized torch.distributed P2P stand-in: sends land in a + mailbox, all rank threads rendezvous, then recvs read their neighbour's + send of the same batch.""" + + irecv = "irecv" + isend = "isend" + + def __init__(self, world_size: int, tls: threading.local) -> None: + self.tls = tls + self.barrier = threading.Barrier(world_size) + self.mailbox: dict[tuple[int, int], torch.Tensor] = {} + self.peers: set[tuple[int, int]] = set() + self.batches = 0 + self.lock = threading.Lock() + + def P2POp(self, op, tensor, peer, group): + return (op, tensor, peer) + + def batch_isend_irecv(self, ops): + rank = self.tls.rank + with self.lock: + self.batches += 1 + for op, tensor, peer in ops: + self.peers.add((rank, peer)) + if op == "isend": + self.mailbox[(rank, peer)] = tensor.clone() + self.barrier.wait(timeout=60) + for op, tensor, peer in ops: + if op == "irecv": + tensor.copy_(self.mailbox[(peer, rank)]) + self.barrier.wait(timeout=60) + return [_NoopReq() for _ in ops] + + +@pytest.mark.parametrize("task", ["t2v", "i2v"]) +def test_tiny_transformer_sp2_lockstep_matches_dense(tp1_group, force_default_gemm, mocker, task) -> None: + """Two lockstep rank threads over uneven frame shards must reproduce the + dense tiny-transformer output, and must actually have sharded the work.""" + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + torch.manual_seed(3) + model = SanaVideoTransformer3DModel(**_TINY_SP_CONFIG).eval() + for _, param in sorted(model.named_parameters()): + param.data.normal_() + batch, frames, height, width = 2, 5, 4, 4 + torch.manual_seed(11) + latent = torch.randn(batch, 4, frames, height, width) + encoder_hidden_states = torch.randn(batch, 6, 8) + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1]]) + if task == "t2v": + timestep = torch.tensor([500.0, 700.0]) + else: + timestep = torch.rand(batch, 1, frames, height // 2, width // 2) * 1000.0 + + with torch.no_grad(): + ref = model( + latent, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + world_size = 2 + tls = threading.local() + group = _LockstepSpGroup(world_size, tls) + fake_dist = _LockstepP2PDist(world_size, tls) + mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=world_size) + mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", side_effect=lambda: tls.rank) + mocker.patch(f"{_MODULE}.get_sp_group", return_value=group) + mocker.patch(f"{_MODULE}.dist", fake_dist) + + def run_rank(rank: int) -> torch.Tensor: + tls.rank = rank + with torch.no_grad(): + return model( + latent, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + with ThreadPoolExecutor(max_workers=world_size) as pool: + outs = list(pool.map(run_rank, range(world_size))) + + num_layers = _TINY_SP_CONFIG["num_layers"] + # One attn1 state reduce per block, its inputs distinct per rank (the ranks + # really processed different frames), one neighbour halo exchange per block + # per rank, and only the final frame gather is a full collective. + assert len(group.reduce_inputs) == num_layers + for parts in group.reduce_inputs: + assert not torch.equal(parts[0], parts[1]) + assert group.gather_calls == 1 + assert fake_dist.batches == world_size * num_layers + assert fake_dist.peers == {(0, 1), (1, 0)} + + torch.testing.assert_close(outs[1], outs[0], rtol=0, atol=0) + # Reduction-reorder noise amplified by the unit-normal test weights; any + # sharding misalignment is orders of magnitude above this. + torch.testing.assert_close(outs[0], ref, rtol=1e-3, atol=1e-3) + + +def test_transformer_declares_empty_sp_plan() -> None: + """Empty plan: the registry takes its normal SP enablement path with zero + hooks and no misleading 'no _sp_plan found' warning; sharding is manual.""" + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + assert SanaVideoTransformer3DModel._sp_plan == {} + + +def test_dummy_run_num_frames_covers_max_sp_degree() -> None: + """Engine warmup must produce enough latent frames for every allowed SP + degree on both VAE variants (temporal /4 and /8); one pixel frame would + fail the per-rank frame check at startup.""" + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline, SanaVideoPipeline + + for temporal_scale in (4, 8): + latent_frames = (SanaVideoPipeline.dummy_run_num_frames - 1) // temporal_scale + 1 + assert latent_frames >= 4 + assert SanaImageToVideoPipeline.dummy_run_num_frames == SanaVideoPipeline.dummy_run_num_frames diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index 122138dc2..083dba0ed 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -57,6 +57,9 @@ DIFFUSION_TEST_SETTINGS = { [DiffusionAccs.TENSOR_PARALLEL], [DiffusionAccs.CFG_PARALLEL], [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL], + [DiffusionAccs.SEQUENCE_PARALLEL], + [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CFG_PARALLEL], + [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.SEQUENCE_PARALLEL], [DiffusionAccs.CACHE_DIT], [DiffusionAccs.CPU_OFFLOAD], [DiffusionAccs.LAYERWISE_OFFLOAD], @@ -72,6 +75,8 @@ DIFFUSION_TEST_SETTINGS = { extra_test_groups=[ [DiffusionAccs.CFG_PARALLEL], [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL], + [DiffusionAccs.SEQUENCE_PARALLEL], + [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CFG_PARALLEL], ], check_i2v_t2v_divergence=False, check_multi_output=False, diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py index 48db6254f..f2d86622a 100644 --- a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py @@ -333,6 +333,9 @@ class SanaVideoPipeline( _encoder_modules = ["text_encoder"] _vae_modules = ["vae"] supports_step_execution = False + # Warmup must yield >= 4 latent frames on both VAE variants (temporal /4 + # and /8) so every allowed SP degree has one frame per rank. + dummy_run_num_frames = 25 default_num_inference_steps = 50 def __init__( diff --git a/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py index 1b64790db..cbad6703d 100644 --- a/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py +++ b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py @@ -18,6 +18,7 @@ from collections.abc import Iterable from dataclasses import dataclass, fields import torch +import torch.distributed as dist from cache_dit import ForwardPattern from torch import nn from vllm.distributed import ( @@ -32,6 +33,11 @@ from vllm.model_executor.utils import set_weight_attrs from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as OmniAttention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group, +) def validate_sana_video_parallel_config(parallel_config) -> None: @@ -46,13 +52,18 @@ def validate_sana_video_parallel_config(parallel_config) -> None: raise NotImplementedError( f"SANA-Video supports cfg_parallel_size 1 or 2, got {cfg_size}. Set --cfg-parallel-size to 1 or 2." ) - sp_size = parallel_config.sequence_parallel_size - if sp_size is not None and sp_size > 1: + sp_size = parallel_config.sequence_parallel_size or 1 + if sp_size not in (1, 2, 4): raise NotImplementedError( - "Sequence parallel is not supported for SANA-Video: its linear attention and " - "GLUMB temporal conv have no SANA-specific SP implementation. Use tensor and/or " - "CFG parallel instead." + f"SANA-Video supports sequence_parallel_size 1, 2 or 4, got {sp_size}. Set --usp to 1, 2 or 4." ) + if parallel_config.ring_degree > 1: + raise NotImplementedError( + "SANA-Video does not support ring sequence parallel: its linear attention has no " + "log-sum-exp state to merge across ring stages. Use --usp for sequence parallel." + ) + if parallel_config.allgather_degree > 1: + raise NotImplementedError("SANA-Video does not support AllGather-KV sequence parallel. Use --usp instead.") if parallel_config.pipeline_parallel_size > 1: raise NotImplementedError("SANA-Video does not support pipeline parallel. Set --pipeline-parallel-size to 1.") if parallel_config.use_hsdp: @@ -63,6 +74,32 @@ def validate_sana_video_parallel_config(parallel_config) -> None: ) +def _sp_frame_split_sizes(num_frames: int, world_size: int) -> list[int]: + """Balanced per-rank frame counts: the first ``num_frames % world_size`` + ranks take one extra frame, so exactly ``world_size`` non-empty chunks.""" + if num_frames < world_size: + raise ValueError( + f"SANA-Video sequence parallel needs at least one latent frame per rank, " + f"got {num_frames} frames for sequence parallel size {world_size}. " + "Lower the sequence parallel degree or increase num_frames." + ) + base, rem = divmod(num_frames, world_size) + return [base + 1 if rank < rem else base for rank in range(world_size)] + + +def _sp_gather_frames(hidden_states: torch.Tensor, sizes: list[int]) -> torch.Tensor: + """Gather per-rank frame shards of shape (B, sizes[rank]*hw, C) into the + full (B, sum(sizes)*hw, C) sequence. Shards are padded to the widest rank + only for the collective; pads are dropped before concatenation.""" + rank = get_sequence_parallel_rank() + frames = hidden_states.unflatten(1, (sizes[rank], -1)) + pad = max(sizes) - sizes[rank] + if pad: + frames = torch.cat([frames, frames.new_zeros((frames.shape[0], pad, *frames.shape[2:]))], dim=1) + parts = get_sp_group().all_gather(frames, dim=0, separate_tensors=True) + return torch.cat([part.narrow(1, 0, size) for part, size in zip(parts, sizes)], dim=1).flatten(1, 2) + + @dataclass class SanaVideoTransformerConfig: in_channels: int = 16 @@ -521,7 +558,37 @@ class GLUMBTempConv(nn.Module): hidden_states_temporal = hidden_states.view(batch_size, num_frames, num_channels, height * width).permute( 0, 2, 1, 3 ) - hidden_states = hidden_states_temporal + self.conv_temp(hidden_states_temporal) + if get_sequence_parallel_world_size() > 1: + # The temporal conv reads one frame across each shard boundary, so + # each rank exchanges single frames with its neighbours only; the + # global first/last rank keeps the dense zero padding on its outer + # side. + assert self.conv_temp.kernel_size == (3, 1), ( + "SP halo exchange assumes a temporal kernel of 3 (one frame per side), " + f"got {self.conv_temp.kernel_size}" + ) + x = hidden_states_temporal + rank = get_sequence_parallel_rank() + world_size = get_sequence_parallel_world_size() + sp_group = get_sp_group() + left = x.new_zeros((x.shape[0], x.shape[1], 1, x.shape[3])) + right = torch.zeros_like(left) + p2p_ops = [] + if rank > 0: + prev_rank = sp_group.ranks[rank - 1] + p2p_ops.append(dist.P2POp(dist.irecv, left, prev_rank, sp_group.device_group)) + p2p_ops.append(dist.P2POp(dist.isend, x[:, :, :1].contiguous(), prev_rank, sp_group.device_group)) + if rank < world_size - 1: + next_rank = sp_group.ranks[rank + 1] + p2p_ops.append(dist.P2POp(dist.isend, x[:, :, -1:].contiguous(), next_rank, sp_group.device_group)) + p2p_ops.append(dist.P2POp(dist.irecv, right, next_rank, sp_group.device_group)) + for req in dist.batch_isend_irecv(p2p_ops): + req.wait() + padded = torch.cat([left, x, right], dim=2) + temporal_out = nn.functional.conv2d(padded, self.conv_temp.weight, padding=0) + else: + temporal_out = self.conv_temp(hidden_states_temporal) + hidden_states = hidden_states_temporal + temporal_out hidden_states = hidden_states.permute(0, 2, 3, 1).view(batch_size, num_frames, height, width, num_channels) if self.norm_type == "rms_norm": @@ -621,9 +688,15 @@ class SanaLinearAttention(nn.Module): # Keep the unrotated denominator contraction in the input dtype. This # intentionally matches Diffusers 0.38.0's SanaLinearAttnProcessor3_0; # only the rotated numerator path is accumulated in FP32. - z = 1 / (key.sum(dim=-1, keepdim=True).transpose(-2, -1) @ query + 1e-15) - + k_sum = key.sum(dim=-1, keepdim=True) scores = torch.matmul(value, key_rotate.transpose(-1, -2)) + if get_sequence_parallel_world_size() > 1: + # Token sums are the only sequence coupling; one packed all-reduce + # turns the per-rank partial state into the global one. + packed = torch.cat([scores, k_sum.float()], dim=-1) + packed = get_sp_group().all_reduce(packed) + scores, k_sum = packed[..., :-1], packed[..., -1:].to(key.dtype) + z = 1 / (k_sum.transpose(-2, -1) @ query + 1e-15) hidden_states = torch.matmul(scores, query_rotate) hidden_states = hidden_states * z @@ -971,6 +1044,10 @@ class SanaVideoTransformer3DModel(nn.Module): ) _layerwise_offload_blocks_attrs = ["transformer_blocks"] + # SP is implemented manually in forward (frame-aligned uneven sharding); + # the empty plan enables the registry's SP path with zero generic hooks. + _sp_plan = {} + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return AutoWeightsLoader(self).load_weights(weights) @@ -1151,6 +1228,26 @@ class SanaVideoTransformer3DModel(nn.Module): hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) + sp_world_size = get_sequence_parallel_world_size() + local_frames = post_patch_num_frames + if sp_world_size > 1: + # Frame-aligned sharding: every rank keeps full spatial extent, so + # the per-frame convs, the GLUMB unflatten and the per-token + # modulation stay local. The I2V per-token timestep is sliced + # before the time embed so its MLP also runs on local tokens only. + sp_sizes = _sp_frame_split_sizes(post_patch_num_frames, sp_world_size) + sp_rank = get_sequence_parallel_rank() + local_frames = sp_sizes[sp_rank] + hidden_states = ( + hidden_states.unflatten(1, (post_patch_num_frames, -1)).split(sp_sizes, dim=1)[sp_rank].flatten(1, 2) + ) + rotary_emb = tuple( + freqs.unflatten(1, (post_patch_num_frames, -1)).split(sp_sizes, dim=1)[sp_rank].flatten(1, 2) + for freqs in rotary_emb + ) + if timestep.ndim == 5: + timestep = timestep.split(sp_sizes, dim=2)[sp_rank] + if guidance is not None: timestep, embedded_timestep = self.time_embed( timestep.flatten(), guidance=guidance, hidden_dtype=hidden_states.dtype @@ -1176,7 +1273,7 @@ class SanaVideoTransformer3DModel(nn.Module): encoder_hidden_states, encoder_attention_mask, timestep, - post_patch_num_frames, + local_frames, post_patch_height, post_patch_width, rotary_emb, @@ -1189,6 +1286,10 @@ class SanaVideoTransformer3DModel(nn.Module): hidden_states = self.proj_out(hidden_states) + if sp_world_size > 1: + # Gather after proj_out where the channel width is smallest. + hidden_states = _sp_gather_frames(hidden_states, sp_sizes) + # 5. Unpatchify hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1