From 57c08d34c359cce233990fd138a8dc9f46973e34 Mon Sep 17 00:00:00 2001 From: vOv <124694036+cr-gao@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:44:41 +0800 Subject: [PATCH] [Model] Add TP and CFG parallelism to SANA-Video 2B (#5861) Signed-off-by: cr-gao --- docs/user_guide/diffusion_features.md | 2 +- recipes/NVIDIA/SANA-Video-2B.md | 29 +- .../test_diffusers_backend.py | 16 +- .../test_pipeline_sana_video_cfg.py | 582 ++++++++++++++++++ .../sana_video/test_sana_video_topology.py | 90 +++ .../sana_video/test_transformer_sana_video.py | 35 +- .../test_transformer_sana_video_tp.py | 445 +++++++++++++ tests/model_tests/diffusion/case_filtering.py | 2 + tests/model_tests/diffusion/config_types.py | 5 + .../diffusion/diff_model_builders.py | 22 +- tests/model_tests/diffusion/model_settings.py | 16 + tests/model_tests/diffusion/task_runners.py | 8 +- tests/model_tests/diffusion/test_alignment.py | 6 - .../diffusion/test_common_offline.py | 5 +- .../pipeline_diffusers_adapter.py | 5 + .../models/sana_video/pipeline_sana_video.py | 165 +++-- .../sana_video/pipeline_sana_video_i2v.py | 138 ++++- .../sana_video/transformer_sana_video.py | 166 ++++- 18 files changed, 1621 insertions(+), 116 deletions(-) create mode 100644 tests/diffusion/models/sana_video/test_pipeline_sana_video_cfg.py create mode 100644 tests/diffusion/models/sana_video/test_sana_video_topology.py create mode 100644 tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 0830a280d..3f6a33878 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -162,7 +162,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** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | +| **SANA-Video-2B T2V I2V** | ❌ | ❌ | ❌ | ✅ | ✅ (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 c85d8ca40..ca4b44ee8 100644 --- a/recipes/NVIDIA/SANA-Video-2B.md +++ b/recipes/NVIDIA/SANA-Video-2B.md @@ -119,6 +119,30 @@ MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ 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: + +```bash +vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --omni \ + --model-class-name SanaVideoPipeline \ + --tensor-parallel-size 2 \ + --dtype bfloat16 \ + --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. + +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. + 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 works; `num_frames` is adapted to Diffusers' `frames` argument. The script @@ -194,8 +218,9 @@ validated compatibility/reference backend. The denoising loop intentionally retains the checkpoint-compatible Diffusers `DPMSolverMultistepScheduler`. - Known limitations: - - Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution - are not validated for the native pipeline. + - Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for + the native pipeline. Sequence parallelism, Cache-DiT, TeaCache, and step + execution are not validated. - The Diffusers backend is a compatibility path and does not provide native vLLM-Omni parallelism or continuous batching. - Native describes pipeline and Transformer ownership, not a zero-Diffusers diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index 3b094c37f..6311d494d 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py @@ -143,7 +143,16 @@ class TestPipelineArgumentsHandling: @pytest.mark.parametrize( "feature_id", - ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"], + [ + "cfg_parallel", + "tensor_parallel", + "ulysses", + "ring", + "teacache", + "cache_dit", + "enforce_eager", + "unsupported_quantization", + ], ) def test_adapter_guard_unsupported_feature(self, feature_id): if feature_id == "cfg_parallel": @@ -151,6 +160,11 @@ class TestPipelineArgumentsHandling: parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1), cache_backend="none", ) + elif feature_id == "tensor_parallel": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, sequence_parallel_size=1), + cache_backend="none", + ) elif feature_id == "ulysses": od_config = _make_od_config( parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2), diff --git a/tests/diffusion/models/sana_video/test_pipeline_sana_video_cfg.py b/tests/diffusion/models/sana_video/test_pipeline_sana_video_cfg.py new file mode 100644 index 000000000..8f1c1ea5d --- /dev/null +++ b/tests/diffusion/models/sana_video/test_pipeline_sana_video_cfg.py @@ -0,0 +1,582 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CFG dispatch tests for the native SANA-Video pipelines. + +These exercise the CFG world-size-1 batch path, the CFG world-size-2 dispatch, +and the fp32 combine semantics on CPU / mocked distributed. The real all-gather +path is covered by the multi-process suite, not here. +""" + +import pytest +import torch + +from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline +from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline +from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaVideoTransformerOutput, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_predict_noise_calls_transformer_with_return_dict_false(): + """predict_noise must call the transformer with return_dict=False. + + The transformer returns a SanaVideoTransformerOutput dataclass when + return_dict is True (its default), which is not subscriptable. The CFG + mixin's default predict_noise does ``result[0]`` and would therefore + raise TypeError on SANA; the pipeline override must force return_dict=False + so the tuple form is returned and indexed. + """ + captured = {} + expected = torch.full((1, 2, 3), 7.0) + + class _SpyTransformer: + def __call__(self, *args, return_dict=True, **kwargs): + captured["return_dict"] = return_dict + # Mirror the real transformer: tuple on False, dataclass on True. + return (expected,) if not return_dict else SanaVideoTransformerOutput(sample=expected) + + pipe = SanaVideoPipeline.__new__(SanaVideoPipeline) # bypass heavy __init__ + pipe.transformer = _SpyTransformer() + + out = pipe.predict_noise( + torch.zeros(1, 2, 3), + encoder_hidden_states=torch.zeros(1, 1, 4), + encoder_attention_mask=None, + timestep=torch.zeros(1), + ) + + assert captured["return_dict"] is False + assert torch.equal(out, expected) + + +class _IdentityScheduler: + """Scheduler stub whose step is a no-op, so diffuse tests isolate the CFG + dispatch from real scheduler math.""" + + order = 1 + + def step(self, noise_pred, t, latents, **kwargs): + return (latents,) + + +def _set_cfg_world_size(monkeypatch, n): + """Pipelines always run inside an initialized CFG group; unit tests stub its + world size instead of spinning up torch.distributed. Patch both pipeline + modules since each resolves the symbol in its own namespace.""" + import vllm_omni.diffusion.distributed.cfg_parallel as cfg + import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v + import vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v as i2v + + monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False) + monkeypatch.setattr(i2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False) + monkeypatch.setattr(cfg, "get_classifier_free_guidance_world_size", lambda: n, raising=False) + monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False) + monkeypatch.setattr(i2v, "model_parallel_is_initialized", lambda: True, raising=False) + + +def _bare_t2v_pipeline(transformer, guidance_scale): + pipe = SanaVideoPipeline.__new__(SanaVideoPipeline) # bypass heavy __init__ + pipe.transformer = transformer + pipe.scheduler = _IdentityScheduler() + pipe._interrupt = False + pipe._guidance_scale = guidance_scale + return pipe + + +def test_cfg_world_size_1_does_single_batch2_forward_per_step(monkeypatch): + """CFG world size 1 with guidance must do exactly one batch-2 transformer + forward per step (the batch path), not the mixin's sequential two-forward + path.""" + _set_cfg_world_size(monkeypatch, 1) + batch_dims = [] + + class _SpyTransformer: + dtype = torch.float32 + + def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict): + batch_dims.append(hidden_states.shape[0]) + # out_channels == in_channels here (no learned-sigma split). + return (torch.zeros_like(hidden_states),) + + b = 1 + pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0) + + out = pipe.diffuse( + latents=torch.randn(b, 4, 2, 4, 4), + timesteps=torch.tensor([1000.0]), + prompt_embeds=torch.randn(b, 3, 8), + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=torch.randn(b, 3, 8), + negative_prompt_attention_mask=torch.ones(b, 3), + guidance_scale=6.0, + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=None, + ) + + assert len(batch_dims) == 1, f"expected 1 forward this step, got {len(batch_dims)}" + assert batch_dims[0] == 2 * b, f"expected batch-2 (neg+pos) forward, got batch {batch_dims[0]}" + assert out.shape == (b, 4, 2, 4, 4) + + +class _DeterministicTransformer: + """Learned-sigma transformer whose output depends on the row's encoder + embeds and timestep, so uncond != text and the CFG combine is non-trivial.""" + + dtype = torch.float32 + + def __init__(self, latent_channels): + self.out_channels = 2 * latent_channels # learned sigma + + def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict): + b, c = hidden_states.shape[0], hidden_states.shape[1] + sig = encoder_hidden_states.float().mean(dim=(1, 2)).view(b, 1, 1, 1, 1) + base = hidden_states.float() + sig + timestep.float().view(b, 1, 1, 1, 1) + return (base.repeat(1, self.out_channels // c, 1, 1, 1),) + + +class _LinearScheduler: + """Deterministic stepping scheduler so multi-step divergence compounds.""" + + order = 1 + + def step(self, noise_pred, t, latents, return_dict=False, **kwargs): + return (latents - 0.1 * noise_pred,) + + +def _pr1_reference_loop( + transformer, + scheduler, + latents, + timesteps, + prompt_embeds, + prompt_mask, + negative_embeds, + negative_mask, + guidance_scale, + latent_channels, + dtype, +): + """The original CFG1 denoise math, verbatim, as the parity oracle.""" + prompt_embeds = torch.cat([negative_embeds, prompt_embeds], dim=0) + prompt_mask = torch.cat([negative_mask, prompt_mask], dim=0) + for t in timesteps: + latent_model_input = torch.cat([latents] * 2) + timestep = t.expand(latent_model_input.shape[0]) + noise_pred = transformer( + latent_model_input.to(dtype=dtype), + encoder_hidden_states=prompt_embeds.to(dtype=dtype), + encoder_attention_mask=prompt_mask, + timestep=timestep, + return_dict=False, + )[0] + noise_pred = noise_pred.float() + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + if transformer.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0] + return latents + + +def test_cfg1_refactored_diffuse_is_bit_identical_to_pr1_loop(monkeypatch): + """The refactored diffuse must reproduce the original CFG1 latents bit-for-bit + across multiple steps, preserving the TP1/CFG1 golden output.""" + _set_cfg_world_size(monkeypatch, 1) + torch.manual_seed(0) + b, c = 2, 4 + latents0 = torch.randn(b, c, 2, 4, 4) + timesteps = torch.tensor([900.0, 600.0, 300.0]) + prompt_embeds = torch.randn(b, 3, 8) + negative_embeds = torch.randn(b, 3, 8) + prompt_mask = torch.ones(b, 3) + negative_mask = torch.ones(b, 3) + guidance_scale = 6.0 + + transformer = _DeterministicTransformer(latent_channels=c) + scheduler = _LinearScheduler() + + ref = _pr1_reference_loop( + transformer, + scheduler, + latents0.clone(), + timesteps, + prompt_embeds, + prompt_mask, + negative_embeds, + negative_mask, + guidance_scale, + latent_channels=c, + dtype=torch.float32, + ) + + pipe = _bare_t2v_pipeline(transformer, guidance_scale=guidance_scale) + pipe.scheduler = scheduler + got = pipe.diffuse( + latents=latents0.clone(), + timesteps=timesteps, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_mask, + negative_prompt_embeds=negative_embeds, + negative_prompt_attention_mask=negative_mask, + guidance_scale=guidance_scale, + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=c, + ) + + assert torch.equal(got, ref) + + +def test_combine_cfg_noise_upcasts_to_fp32_before_combine(): + """combine_cfg_noise must upcast positive/negative to fp32 before the CFG + formula, so the CFG2 (bf16 all-gather) combine is bit-identical to CFG1's + fp32 combine.""" + pipe = SanaVideoPipeline.__new__(SanaVideoPipeline) + torch.manual_seed(0) + pos = torch.randn(1, 4, 2, 4, 4).bfloat16() + neg = torch.randn(1, 4, 2, 4, 4).bfloat16() + scale = 6.0 + + out = pipe.combine_cfg_noise(pos, neg, scale, cfg_normalize=False) + + expected = neg.float() + scale * (pos.float() - neg.float()) + assert out.dtype == torch.float32 + assert torch.equal(out, expected) + + +def test_combine_cfg_noise_accepts_tuple_predictions_from_gather(): + """The CFG-parallel gather path hands combine_cfg_noise single-element + tuples, not bare tensors; the override must upcast and combine them.""" + pipe = SanaVideoPipeline.__new__(SanaVideoPipeline) + torch.manual_seed(0) + pos = torch.randn(1, 4, 2, 4, 4).bfloat16() + neg = torch.randn(1, 4, 2, 4, 4).bfloat16() + scale = 6.0 + + out = pipe.combine_cfg_noise((pos,), (neg,), scale, cfg_normalize=False) + + expected = neg.float() + scale * (pos.float() - neg.float()) + assert out.dtype == torch.float32 + assert torch.equal(out, expected) + + +def test_cfg2_dispatch_builds_independent_branch_kwargs(monkeypatch): + """In CFG world size 2, diffuse must build separate batch-B positive/negative + kwargs (no pre-cat), and call the mixin with cfg_normalize=False and the + learned-sigma output_slice.""" + _set_cfg_world_size(monkeypatch, 2) + + b, c = 1, 4 + captured = {} + + def _spy_cfg(**kwargs): + captured.update(kwargs) + return torch.zeros(b, c, 2, 4, 4) + + prompt_embeds = torch.randn(b, 3, 8) + negative_embeds = torch.randn(b, 3, 8) + dtype = torch.float32 + + pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0) + pipe.predict_noise_maybe_with_cfg = _spy_cfg + + pipe.diffuse( + latents=torch.randn(b, c, 2, 4, 4), + timesteps=torch.tensor([1000.0]), + prompt_embeds=prompt_embeds, + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=negative_embeds, + negative_prompt_attention_mask=torch.zeros(b, 3), + guidance_scale=6.0, + extra_step_kwargs={}, + dtype=dtype, + output_slice=c, + ) + + assert captured, "predict_noise_maybe_with_cfg was not called in CFG world size 2" + assert captured["do_true_cfg"] is True + assert captured["true_cfg_scale"] == 6.0 + assert captured["cfg_normalize"] is False + assert captured["output_slice"] == c + + pos = captured["positive_kwargs"] + neg = captured["negative_kwargs"] + assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B, not pre-concatenated" + assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B, not pre-concatenated" + assert pos["timestep"].shape[0] == b + assert torch.equal(pos["encoder_hidden_states"], prompt_embeds.to(dtype)) + assert torch.equal(neg["encoder_hidden_states"], negative_embeds.to(dtype)) + + +# ── I2V CFG dispatch ── + + +class _I2VSpyScheduler: + """Records the frame count it steps on and mutates non-first frames so + first-frame preservation is observable.""" + + order = 1 + + def __init__(self): + self.stepped_frame_counts = [] + + def step(self, noise_pred, t, latents, return_dict=False, **kwargs): + self.stepped_frame_counts.append(latents.shape[2]) + return (latents + 1.0,) + + +def _bare_i2v_pipeline(transformer, scheduler, guidance_scale): + pipe = SanaImageToVideoPipeline.__new__(SanaImageToVideoPipeline) + pipe.transformer = transformer + pipe.scheduler = scheduler + pipe._interrupt = False + pipe._guidance_scale = guidance_scale + return pipe + + +def _conditioning_mask(b, frames): + mask = torch.zeros(b, 1, frames, 1, 1) + mask[:, :, 0] = 1 + return mask + + +def test_i2v_cfg1_preserves_first_frame_and_steps_only_tail(monkeypatch): + """I2V CFG1 must keep the conditioning first frame bit-exact and only run the + scheduler on frames [1:].""" + _set_cfg_world_size(monkeypatch, 1) + b, c, frames = 1, 4, 3 + + class _SpyTransformer: + dtype = torch.float32 + out_channels = 2 * c + + def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict): + return (torch.zeros(hidden_states.shape[0], 2 * c, frames, 4, 4),) + + scheduler = _I2VSpyScheduler() + pipe = _bare_i2v_pipeline(_SpyTransformer(), scheduler, guidance_scale=6.0) + + latents0 = torch.randn(b, c, frames, 4, 4) + out = pipe.diffuse( + latents=latents0.clone(), + timesteps=torch.tensor([1000.0]), + prompt_embeds=torch.randn(b, 3, 8), + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=torch.randn(b, 3, 8), + negative_prompt_attention_mask=torch.ones(b, 3), + guidance_scale=6.0, + conditioning_mask=_conditioning_mask(b, frames), + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=c, + ) + + assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first (conditioning) frame must be unchanged" + assert torch.equal(out[:, :, 1:], latents0[:, :, 1:] + 1.0), "tail frames must be the scheduler output" + assert scheduler.stepped_frame_counts == [frames - 1], "scheduler must step only on frames [1:]" + + +def test_i2v_cfg2_dispatch_builds_masked_branch_kwargs(monkeypatch): + """I2V CFG world size 2 must build batch-B per-branch kwargs with the + conditioning-frame timestep zeroed (mask not concatenated) and preserve the + first frame.""" + _set_cfg_world_size(monkeypatch, 2) + b, c, frames = 1, 4, 3 + captured = {} + + def _spy_cfg(**kwargs): + captured.update(kwargs) + return torch.zeros(b, c, frames, 4, 4) + + scheduler = _I2VSpyScheduler() + pipe = _bare_i2v_pipeline(transformer=None, scheduler=scheduler, guidance_scale=6.0) + pipe.predict_noise_maybe_with_cfg = _spy_cfg + + prompt_embeds = torch.randn(b, 3, 8) + negative_embeds = torch.randn(b, 3, 8) + latents0 = torch.randn(b, c, frames, 4, 4) + mask = _conditioning_mask(b, frames) + + out = pipe.diffuse( + latents=latents0.clone(), + timesteps=torch.tensor([1000.0]), + prompt_embeds=prompt_embeds, + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=negative_embeds, + negative_prompt_attention_mask=torch.zeros(b, 3), + guidance_scale=6.0, + conditioning_mask=mask, + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=c, + ) + + assert captured["do_true_cfg"] is True + assert captured["cfg_normalize"] is False + assert captured["output_slice"] == c + pos = captured["positive_kwargs"] + neg = captured["negative_kwargs"] + assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B" + assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B" + assert tuple(pos["timestep"].shape) == tuple(mask.shape), "mask must not be concatenated in CFG2" + assert torch.equal(pos["timestep"][:, :, 0], torch.zeros(b, 1, 1, 1)), "conditioning frame timestep must be 0" + assert torch.equal(pos["encoder_hidden_states"], prompt_embeds) + assert torch.equal(neg["encoder_hidden_states"], negative_embeds) + assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first frame must be preserved" + + +def test_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker): + """diffuse must run the CFG-parallel validity check so a scale<=1 or missing + negative prompt warns the user instead of silently wasting a rank.""" + _set_cfg_world_size(monkeypatch, 2) + + b = 1 + pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0) + pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, 4, 2, 4, 4) + check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True) + + pipe.diffuse( + latents=torch.randn(b, 4, 2, 4, 4), + timesteps=torch.tensor([1000.0]), + prompt_embeds=torch.randn(b, 3, 8), + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=torch.randn(b, 3, 8), + negative_prompt_attention_mask=torch.ones(b, 3), + guidance_scale=6.0, + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=None, + ) + + check.assert_called_once_with(6.0, True) + + +def test_i2v_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker): + """I2V has its own diffuse; it must also run the CFG-parallel validity check.""" + _set_cfg_world_size(monkeypatch, 2) + b, c, frames = 1, 4, 3 + + pipe = _bare_i2v_pipeline(None, _I2VSpyScheduler(), guidance_scale=6.0) + pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, c, frames, 4, 4) + check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True) + + pipe.diffuse( + latents=torch.randn(b, c, frames, 4, 4), + timesteps=torch.tensor([1000.0]), + prompt_embeds=torch.randn(b, 3, 8), + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=torch.randn(b, 3, 8), + negative_prompt_attention_mask=torch.ones(b, 3), + guidance_scale=6.0, + conditioning_mask=_conditioning_mask(b, frames), + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=c, + ) + + check.assert_called_once_with(6.0, True) + + +def test_diffuse_runs_without_initialized_parallel_groups(): + """Single-process runs (unit tests, offline scripts) never initialize the + distributed groups; diffuse must fall back to the serial path instead of + tripping the parallel-state assertion. No world-size patching on purpose.""" + + class _SpyTransformer: + dtype = torch.float32 + + def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict): + return (torch.zeros_like(hidden_states),) + + b = 1 + pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0) + + out = pipe.diffuse( + latents=torch.randn(b, 4, 2, 4, 4), + timesteps=torch.tensor([1000.0]), + prompt_embeds=torch.randn(b, 3, 8), + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=torch.randn(b, 3, 8), + negative_prompt_attention_mask=torch.ones(b, 3), + guidance_scale=6.0, + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=None, + ) + + assert out.shape == (b, 4, 2, 4, 4) + + +def test_diffuse_propagates_group_error_when_parallel_initialized(monkeypatch): + """Inside an initialized parallel run a failing CFG-group lookup is a real + error and must propagate instead of silently flipping to the serial path.""" + import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v + + monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False) + + def _broken_group_lookup(): + raise AssertionError("classifier_free_guidance group is not initialized") + + monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", _broken_group_lookup) + + b = 1 + pipe = _bare_t2v_pipeline(None, guidance_scale=6.0) + + with pytest.raises(AssertionError, match="classifier_free_guidance"): + pipe.diffuse( + latents=torch.randn(b, 4, 2, 4, 4), + timesteps=torch.tensor([1000.0]), + prompt_embeds=torch.randn(b, 3, 8), + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=torch.randn(b, 3, 8), + negative_prompt_attention_mask=torch.ones(b, 3), + guidance_scale=6.0, + extra_step_kwargs={}, + dtype=torch.float32, + output_slice=None, + ) + + +def test_t2v_forwards_eta_and_generator_into_scheduler(monkeypatch): + """T2V must keep passing eta/generator (extra_step_kwargs) into scheduler.step.""" + _set_cfg_world_size(monkeypatch, 1) + captured = {} + + class _KwargScheduler: + order = 1 + + def step(self, noise_pred, t, latents, **kwargs): + captured.update(kwargs) + return (latents,) + + class _SpyTransformer: + dtype = torch.float32 + + def __call__(self, hidden_states, **kwargs): + return (torch.zeros_like(hidden_states),) + + b = 1 + pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0) + pipe.scheduler = _KwargScheduler() + generator = torch.Generator().manual_seed(0) + + pipe.diffuse( + latents=torch.randn(b, 4, 2, 4, 4), + timesteps=torch.tensor([1000.0]), + prompt_embeds=torch.randn(b, 3, 8), + prompt_attention_mask=torch.ones(b, 3), + negative_prompt_embeds=torch.randn(b, 3, 8), + negative_prompt_attention_mask=torch.ones(b, 3), + guidance_scale=6.0, + extra_step_kwargs={"eta": 0.5, "generator": generator}, + dtype=torch.float32, + output_slice=None, + ) + + assert captured.get("eta") == 0.5 + assert captured.get("generator") is generator + assert captured.get("return_dict") is False diff --git a/tests/diffusion/models/sana_video/test_sana_video_topology.py b/tests/diffusion/models/sana_video/test_sana_video_topology.py new file mode 100644 index 000000000..ac0fa97e0 --- /dev/null +++ b/tests/diffusion/models/sana_video/test_sana_video_topology.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""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. +""" + +from types import SimpleNamespace + +import pytest + +from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline +from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline +from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + validate_sana_video_parallel_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _pc( + tp: int = 1, + cfg: int = 1, + sp: int | None = None, + text_tp: int = 1, + pp: int = 1, + hsdp: bool = False, +) -> SimpleNamespace: + return SimpleNamespace( + tensor_parallel_size=tp, + cfg_parallel_size=cfg, + sequence_parallel_size=sp, + text_encoder_tp_size=text_tp, + pipeline_parallel_size=pp, + use_hsdp=hsdp, + ) + + +@pytest.mark.parametrize("tp", [1, 2]) +def test_supported_tp_sizes_pass(tp: int) -> None: + validate_sana_video_parallel_config(_pc(tp=tp)) + + +@pytest.mark.parametrize("tp", [3, 4, 8]) +def test_unsupported_tp_sizes_rejected(tp: int) -> None: + with pytest.raises(NotImplementedError, match="tensor_parallel_size"): + validate_sana_video_parallel_config(_pc(tp=tp)) + + +def test_cfg3_rejected() -> None: + with pytest.raises(NotImplementedError, match="cfg_parallel_size"): + 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)) + + +def test_pipeline_parallel_rejected() -> None: + with pytest.raises(NotImplementedError, match="pipeline parallel"): + validate_sana_video_parallel_config(_pc(pp=2)) + + +def test_hsdp_rejected() -> None: + with pytest.raises(NotImplementedError, match="HSDP"): + validate_sana_video_parallel_config(_pc(hsdp=True)) + + +def test_text_encoder_tp_rejected() -> None: + with pytest.raises(NotImplementedError, match="text encoder"): + validate_sana_video_parallel_config(_pc(text_tp=2)) + + +def test_tp2_cfg2_combo_passes() -> None: + validate_sana_video_parallel_config(_pc(tp=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.""" + load = mocker.patch.object(SanaVideoPipeline, "_load_components") + od_config = SimpleNamespace(parallel_config=_pc(tp=3)) + + with pytest.raises(NotImplementedError, match="tensor_parallel_size"): + pipeline_cls(od_config=od_config) + + load.assert_not_called() diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video.py b/tests/diffusion/models/sana_video/test_transformer_sana_video.py index 804692d7a..0bc2ddb8e 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video.py @@ -1,11 +1,37 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + import pytest import torch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +@pytest.fixture(autouse=True) +def _init_distributed(monkeypatch): + """The native transformer uses vLLM parallel linear layers, which require a + tensor-parallel group; initialize a single-process group for CPU tests.""" + from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.model_executor.layers.utils import default_unquantized_gemm + + monkeypatch.setattr( + "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", + lambda: default_unquantized_gemm, + ) + 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() + + _TINY_CONFIG = { "in_channels": 4, "out_channels": 4, @@ -323,6 +349,7 @@ def test_tiny_transformer_matches_diffusers_and_frozen_output(): from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaDistributedRMSNorm, SanaLinearAttention, SanaRMSNorm, SanaVideoTransformerConfig, @@ -342,10 +369,10 @@ def test_tiny_transformer_matches_diffusers_and_frozen_output(): assert isinstance(block.attn1, SanaLinearAttention) assert isinstance(block.attn2.attn, Attention) assert block.attn2.attn.role == "cross" - assert isinstance(block.attn1.norm_q, SanaRMSNorm) - assert isinstance(block.attn1.norm_k, SanaRMSNorm) - assert isinstance(block.attn2.norm_q, SanaRMSNorm) - assert isinstance(block.attn2.norm_k, SanaRMSNorm) + assert isinstance(block.attn1.norm_q, SanaDistributedRMSNorm) + assert isinstance(block.attn1.norm_k, SanaDistributedRMSNorm) + assert isinstance(block.attn2.norm_q, SanaDistributedRMSNorm) + assert isinstance(block.attn2.norm_k, SanaDistributedRMSNorm) assert isinstance(model.caption_norm, SanaRMSNorm) torch.manual_seed(11) diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py b/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py new file mode 100644 index 000000000..41ca717d6 --- /dev/null +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py @@ -0,0 +1,445 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""TP unit tests for SANA-Video: distributed RMSNorm, sharded attention and +sharded GLUMB temporal conv. + +CPU-only. The TP world size and the collective are mocked so a single process +simulates the shard/reduce logic a real multi-rank run depends on. +""" + +import os + +import pytest +import torch + +from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + GLUMBTempConv, + SanaCrossAttention, + SanaDistributedRMSNorm, + SanaLinearAttention, + SanaRMSNorm, +) + +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, + ) + + +def test_distributed_rms_norm_tp1_matches_sana_rms_norm(mocker) -> None: + """TP1 distributed norm must be bit-identical to the checkpoint SanaRMSNorm.""" + dim = 24 + torch.manual_seed(0) + + dense = SanaRMSNorm(dim, eps=1e-5, elementwise_affine=True) + dense.weight.data.normal_() + + dist = SanaDistributedRMSNorm(dim, eps=1e-5) + dist.weight.data.copy_(dense.weight.data) + + x = torch.randn(2, 5, dim) + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1) + out = dist(x) + ref = dense(x) + + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +def test_distributed_rms_norm_tp1_bf16_weight_matches_sana_rms_norm(mocker) -> None: + """bf16 weight must follow SanaRMSNorm's cast-to-weight-dtype branch. + + An unconditional cast back to the input dtype would diverge here; this pins + the fp16/bf16-only cast that SanaRMSNorm performs. + """ + dim = 24 + torch.manual_seed(0) + + dense = SanaRMSNorm(dim, eps=1e-5, elementwise_affine=True) + dense.weight.data.normal_() + dense.weight.data = dense.weight.data.to(torch.bfloat16) + + dist = SanaDistributedRMSNorm(dim, eps=1e-5) + dist.weight.data = dense.weight.data.clone() + + x = torch.randn(2, 5, dim, dtype=torch.bfloat16) + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1) + out = dist(x) + ref = dense(x) + + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +def test_distributed_rms_norm_tp2_shards_aggregate_to_dense(mocker) -> None: + """Two TP2 shards must reproduce dense SanaRMSNorm. + + The RMS must use the full (unsharded) hidden size: each rank reduces its + local sum-of-squares across ranks, not a per-shard mean over its half. + """ + full_dim, tp_size = 24, 2 + half = full_dim // tp_size + torch.manual_seed(0) + + dense = SanaRMSNorm(full_dim, eps=1e-5, elementwise_affine=True) + dense.weight.data.normal_() + + dist0 = SanaDistributedRMSNorm(half, eps=1e-5) + dist1 = SanaDistributedRMSNorm(half, eps=1e-5) + dist0.weight.data.copy_(dense.weight.data[:half]) + dist1.weight.data.copy_(dense.weight.data[half:]) + + x = torch.randn(2, 5, full_dim) + x0, x1 = x[..., :half], x[..., half:] + + # A real all-reduce returns the sum of every rank's local sum-of-squares. + global_sum_sq = x.to(torch.float32).pow(2).sum(-1, keepdim=True) + + def fake_all_reduce(tensor: torch.Tensor) -> torch.Tensor: + return global_sum_sq + + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size) + mocker.patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=fake_all_reduce) + out0 = dist0(x0) + out1 = dist1(x1) + + out = torch.cat([out0, out1], dim=-1) + ref = dense(x) + + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +def test_distributed_rms_norm_weight_loader_shards_dim0(mocker) -> None: + """Each rank loads its own dim-0 slice of the global affine weight.""" + full_dim, tp_size, rank = 24, 2, 1 + half = full_dim // tp_size + + norm = SanaDistributedRMSNorm(half, eps=1e-5) + full_weight = torch.randn(full_dim) + + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size) + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank) + norm.weight.weight_loader(norm.weight, full_weight) + + torch.testing.assert_close(norm.weight.data, full_weight[rank * half : (rank + 1) * half]) + + +def test_distributed_rms_norm_weight_loader_tp1_copies_full() -> None: + dim = 24 + norm = SanaDistributedRMSNorm(dim, eps=1e-5) + full_weight = torch.randn(dim) + + norm.weight.weight_loader(norm.weight, full_weight) + + torch.testing.assert_close(norm.weight.data, full_weight) + + +_SELF_ATTN_GOLDEN = torch.tensor( + [ + 5.0944157, + -16.9981823, + 55.4457436, + 42.5174408, + -9.3055859, + 19.3744793, + -47.4345360, + -5.8110709, + -14.0679598, + -22.3440742, + -8.3834534, + 8.9109459, + ] +) + + +def _fill_by_name(module: torch.nn.Module) -> None: + """Deterministic per-parameter fill that is stable across TP refactors.""" + torch.manual_seed(0) + for _, param in sorted(module.named_parameters()): + param.data.normal_() + + +def _self_attn_inputs(seq_len: int, dim: int, head_dim: int): + torch.manual_seed(1) + x = torch.randn(2, seq_len, dim) + 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_self_attention_tp1_matches_golden_and_keeps_param_names(tp1_group, force_default_gemm) -> None: + """TP1 self-attention stays bit-identical to PR1 and preserves checkpoint keys.""" + dim, num_heads, head_dim, seq_len = 24, 4, 6, 8 + 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() + _fill_by_name(attn) + x, rotary_emb = _self_attn_inputs(seq_len, dim, head_dim) + + with torch.no_grad(): + out = attn(x, rotary_emb=rotary_emb) + + assert {name for name, _ in attn.named_parameters()} == { + "to_q.weight", + "to_k.weight", + "to_v.weight", + "norm_q.weight", + "norm_k.weight", + "to_out.0.weight", + "to_out.0.bias", + } + torch.testing.assert_close(out.flatten()[:12], _SELF_ATTN_GOLDEN, rtol=1e-5, atol=1e-5) + + +def _mock_tp2(mocker): + """Mock a TP world size of 2 for module construction (no real group).""" + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2) + mocker.patch( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", + return_value=2, + ) + tp_group = mocker.MagicMock() + tp_group.world_size = 2 + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp_group) + + +def test_self_attention_tp2_shards_projections_and_heads(mocker) -> None: + """TP2 splits heads column-wise (q/k/v) and the output projection row-wise.""" + dim, num_heads, head_dim = 24, 4, 6 + inner_dim = num_heads * head_dim + _mock_tp2(mocker) + + attn = SanaLinearAttention( + dim=dim, num_heads=num_heads, head_dim=head_dim, dropout=0.0, bias=False, qk_norm="rms_norm_across_heads" + ) + + assert attn.heads == num_heads // 2 + assert attn.to_q.weight.shape == (inner_dim // 2, dim) + assert attn.to_k.weight.shape == (inner_dim // 2, dim) + assert attn.to_v.weight.shape == (inner_dim // 2, dim) + assert attn.to_out[0].weight.shape == (dim, inner_dim // 2) + assert attn.norm_q.weight.shape == (inner_dim // 2,) + assert attn.norm_k.weight.shape == (inner_dim // 2,) + + +def test_cross_attention_tp2_shards_projections_and_heads(mocker) -> None: + """TP2 cross-attention shards q/k/v column-wise, to_out row-wise, heads locally.""" + dim, cross_dim, num_heads, head_dim = 24, 16, 4, 6 + inner_dim = num_heads * head_dim + _mock_tp2(mocker) + + attn = SanaCrossAttention( + dim=dim, + cross_attention_dim=cross_dim, + num_heads=num_heads, + head_dim=head_dim, + dropout=0.0, + bias=False, + out_bias=True, + qk_norm="rms_norm_across_heads", + prefix="attn2", + ) + + assert attn.heads == num_heads // 2 + assert attn.to_q.weight.shape == (inner_dim // 2, dim) + assert attn.to_k.weight.shape == (inner_dim // 2, cross_dim) + assert attn.to_v.weight.shape == (inner_dim // 2, cross_dim) + assert attn.to_out[0].weight.shape == (dim, inner_dim // 2) + assert isinstance(attn.norm_q, SanaDistributedRMSNorm) + assert attn.norm_q.weight.shape == (inner_dim // 2,) + + +_TINY_TRANSFORMER_CONFIG = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 12, + "num_layers": 1, + "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": 32, +} + + +def test_tp2_load_weights_shards_full_checkpoint_without_missing(mocker) -> None: + """The full checkpoint loads under TP2 via weight_loaders with no missing/unexpected keys.""" + from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer + + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + _mock_tp2(mocker) + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=0) + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0) + + reference = DiffusersTransformer(**_TINY_TRANSFORMER_CONFIG) + model = SanaVideoTransformer3DModel(**_TINY_TRANSFORMER_CONFIG) + + loaded = model.load_weights(list(reference.state_dict().items())) + + assert loaded == set(model.state_dict()) + + +# ── GLUMB temporal conv TP ── + +# hidden_channels = expand_ratio * in_channels = 8, so conv_inverted emits 2h = 16. +_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO = 4, 4, 2.0 + +_GLUMB_GOLDEN = torch.tensor( + [ + 51.4757385, + -10.9912930, + -4.1482553, + -26.1580830, + -19.7727261, + 17.9901562, + -1.6007471, + 13.3993073, + -8.6459389, + -10.4008112, + -7.4962354, + 1.5446399, + ] +) + + +def _glumb_inputs() -> torch.Tensor: + torch.manual_seed(1) + return torch.randn(1, 2, 3, 3, _GLUMB_IN) # [B, F, H, W, C] + + +def test_glumb_tp1_matches_golden_and_keeps_param_names(tp1_group) -> None: + """TP1 GLUMB stays bit-identical to the PR1 conv output and preserves keys.""" + glumb = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False) + glumb.eval() + _fill_by_name(glumb) + x = _glumb_inputs() + + with torch.no_grad(): + out = glumb(x) + + assert {name for name, _ in glumb.named_parameters()} == { + "conv_inverted.weight", + "conv_inverted.bias", + "conv_depth.weight", + "conv_depth.bias", + "conv_point.weight", + "conv_temp.weight", + } + torch.testing.assert_close(out.flatten()[:12], _GLUMB_GOLDEN, rtol=1e-5, atol=1e-5) + + +def _random_glumb_checkpoint(hidden_channels: int) -> dict[str, torch.Tensor]: + """Checkpoint-format (nn.Conv2d) GLUMB weights with per-segment asymmetry. + + conv_inverted / conv_depth carry two h-channel segments (value then gate). + Filling the gate half from a different distribution than the value half is + what exposes a naive contiguous dim-0 split: that bug would place the whole + value segment on rank 0 and the whole gate segment on rank 1, so each rank's + local chunk(2) would slice value-vs-value instead of value-vs-gate. + """ + torch.manual_seed(0) + h = hidden_channels + value_w = torch.randn(h, _GLUMB_IN, 1, 1) + gate_w = torch.randn(h, _GLUMB_IN, 1, 1) * 3.0 + 5.0 + value_b = torch.randn(h) + gate_b = torch.randn(h) * 3.0 + 5.0 + value_dw = torch.randn(h, 1, 3, 3) + gate_dw = torch.randn(h, 1, 3, 3) * 3.0 + 5.0 + value_db = torch.randn(h) + gate_db = torch.randn(h) * 3.0 + 5.0 + return { + "conv_inverted.weight": torch.cat([value_w, gate_w], dim=0), + "conv_inverted.bias": torch.cat([value_b, gate_b], dim=0), + "conv_depth.weight": torch.cat([value_dw, gate_dw], dim=0), + "conv_depth.bias": torch.cat([value_db, gate_db], dim=0), + "conv_point.weight": torch.randn(_GLUMB_OUT, h, 1, 1), + "conv_temp.weight": torch.randn(_GLUMB_OUT, _GLUMB_OUT, 3, 1), + } + + +def _mock_glumb_tp(mocker, world_size: int, rank: int) -> None: + """Mock a TP group of ``world_size`` at ``rank`` with an identity all-reduce, + so a single process can build and run one rank's shard in isolation.""" + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=world_size) + mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank) + # Identity conv_point all-reduce returns each rank's local partial for the + # test to sum. + mocker.patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t) + + +def _load_glumb_checkpoint(glumb: GLUMBTempConv, checkpoint: dict[str, torch.Tensor]) -> None: + """Load like AutoWeightsLoader: custom weight_loader when present (the sharded + convs), plain copy for replicated params (conv_temp).""" + params = dict(glumb.named_parameters()) + for name, tensor in checkpoint.items(): + param = params[name] + loader = getattr(param, "weight_loader", None) + if loader is None: + param.data.copy_(tensor) + else: + loader(param, tensor) + + +def test_glumb_tp2_shards_aggregate_to_dense(mocker) -> None: + """The merged value/gate two-segment shard must aggregate to the dense conv. + + Everything after conv_point (temporal conv + residual) is linear and + bias-free, so summing each rank's full forward (with the conv_point + all-reduce mocked to identity) reproduces the replicated dense output. + """ + hidden_channels = int(_GLUMB_RATIO * _GLUMB_IN) + checkpoint = _random_glumb_checkpoint(hidden_channels) + x = _glumb_inputs() + + _mock_glumb_tp(mocker, world_size=1, rank=0) + dense = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False) + dense.eval() + _load_glumb_checkpoint(dense, checkpoint) + with torch.no_grad(): + ref = dense(x) + + partials = [] + for rank in (0, 1): + _mock_glumb_tp(mocker, world_size=2, rank=rank) + shard = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False) + shard.eval() + _load_glumb_checkpoint(shard, checkpoint) + with torch.no_grad(): + partials.append(shard(x)) + + aggregated = partials[0] + partials[1] + torch.testing.assert_close(aggregated, ref, rtol=1e-5, atol=1e-5) diff --git a/tests/model_tests/diffusion/case_filtering.py b/tests/model_tests/diffusion/case_filtering.py index a4fa5dc13..da1aaa552 100644 --- a/tests/model_tests/diffusion/case_filtering.py +++ b/tests/model_tests/diffusion/case_filtering.py @@ -76,6 +76,8 @@ def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts test_info.supported_tasks, test_info.check_multi_output and test_group is None, test_info.check_determinism and test_group is None, + # The divergence check compares generated outputs, which only the offline test does. + *([] if online else [test_info.check_i2v_t2v_divergence]), id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name, marks=get_test_group_marks(model_name, test_group, test_info.marks), ) diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py index 2872f7cec..634036ed4 100644 --- a/tests/model_tests/diffusion/config_types.py +++ b/tests/model_tests/diffusion/config_types.py @@ -76,6 +76,11 @@ class DiffusionModelTestOpts(NamedTuple): check_multi_output: bool = True # Runs multiple generations in one request check_determinism: bool = True # Runs 2 generations with the same seed and check determinism + # For IMAGE_TO_VIDEO, also run a T2V request and assert the outputs differ, to catch a + # silently dropped input image. Only valid for unified text/image pipelines (e.g. LTX2); + # set False for pipelines that mandate an image and fail closed without one (SANA-Video I2V). + check_i2v_t2v_divergence: bool = True + ### Mappings & utils for building offline Omni() instances given a list of enabled accelerations ACC_OMNI_KWARGS = { diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 7c1c9df67..094ef7d84 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -1,3 +1,4 @@ +import json from functools import partial from pathlib import Path @@ -49,15 +50,17 @@ def tiny_ltx2_builder() -> str: return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline") -def tiny_sana_video_builder() -> str: +def _build_tiny_sana_video(pipeline_name: str) -> str: """Build a tiny 480p SANA-Video model without downloading model weights. The tokenizer is the only component loaded from the upstream repository. All modules with weights are initialized locally from intentionally small - configs, and the scheduler is constructed from its weight-free config. + configs, and the scheduler is constructed from its weight-free config. The + I2V pipeline shares this diffusers layout, so it reuses the same components + and only rewrites model_index.json's ``_class_name`` to its own class. """ model_id = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers" - model_dir = _get_tiny_model_path("SanaVideoPipeline") + model_dir = _get_tiny_model_path(pipeline_name) tokenizer = GemmaTokenizerFast.from_pretrained(model_id, subfolder="tokenizer") scheduler = DPMSolverMultistepScheduler( @@ -112,9 +115,22 @@ def tiny_sana_video_builder() -> str: scheduler=scheduler, ) pipeline.to(dtype=torch.bfloat16).save_pretrained(model_dir) + if pipeline_name != "SanaVideoPipeline": + index_path = Path(model_dir) / "model_index.json" + index = json.loads(index_path.read_text()) + index["_class_name"] = pipeline_name + index_path.write_text(json.dumps(index, indent=2)) return model_dir +def tiny_sana_video_builder() -> str: + return _build_tiny_sana_video("SanaVideoPipeline") + + +def tiny_sana_video_i2v_builder() -> str: + return _build_tiny_sana_video("SanaImageToVideoPipeline") + + def _shrink_flux_clip_text_encoder(config: dict) -> dict: config["num_hidden_layers"] = 2 return config diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index f24564f62..5cba1e519 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -53,6 +53,22 @@ DIFFUSION_TEST_SETTINGS = { model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers", builder=diff_model_builders.tiny_sana_video_builder, supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO], + extra_test_groups=[ + [DiffusionAccs.TENSOR_PARALLEL], + [DiffusionAccs.CFG_PARALLEL], + [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL], + ], + check_multi_output=False, + ), + "SanaImageToVideoPipeline": DiffusionModelTestOpts( + model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + builder=diff_model_builders.tiny_sana_video_i2v_builder, + supported_tasks=[DiffusionTasks.IMAGE_TO_VIDEO], + extra_test_groups=[ + [DiffusionAccs.CFG_PARALLEL], + [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL], + ], + check_i2v_t2v_divergence=False, check_multi_output=False, ), "QwenImagePipeline": DiffusionModelTestOpts( diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py index 4fb31fb5c..010d73777 100644 --- a/tests/model_tests/diffusion/task_runners.py +++ b/tests/model_tests/diffusion/task_runners.py @@ -220,16 +220,20 @@ def run_and_validate_text_to_video_request(omni: Omni): _validate_video(_run_offline_t2v(omni)) -def run_and_validate_image_to_video_request(omni: Omni): +def run_and_validate_image_to_video_request(omni: Omni, check_t2v_divergence: bool = True): """Run and validate an image to video request. LTX2Pipeline uses a unified text/image entry and can fall back to T2V if the image never reaches conditioning, so also run T2V with the same seed and prompt and assert the outputs differ; a regression that silently drops the input image would otherwise still pass a "did we get a valid - video back" check alone. + video back" check alone. Pipelines that mandate an image and fail closed + without one (SANA-Video I2V) pass check_t2v_divergence=False: a dropped + image raises instead of silently degrading, so the T2V comparison is moot. """ i2v_video = _validate_video_frames(_get_offline_videos(_run_offline_i2v(omni)))[0] + if not check_t2v_divergence: + return t2v_video = _validate_video_frames(_get_offline_videos(_run_offline_t2v(omni)))[0] assert not np.array_equal(i2v_video, t2v_video), ( "I2V output is identical to T2V output with the same seed and prompt; " diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index f9a5072f9..acc1067a6 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -61,12 +61,6 @@ EXCLUDED_MODELS = [ "HunyuanVideo15Pipeline", "HunyuanVideo15ImageToVideoPipeline", "LingBotVideoPipeline", - # The common harness does not yet define an image-to-video task or supply - # an input image. SANA I2V is covered by - # tests/diffusion/models/sana_video/test_pipeline_sana_video.py and - # tests/e2e/online_serving/test_sana_video_expansion.py. SANA T2V is exercised through - # the common tiny-model setting in model_settings.py. - "SanaImageToVideoPipeline", # Requires camera/action assets and persistent AR state; covered by its # dedicated contract tests and real-checkpoint E2E smoke instead. "LingBotWorldCausalDMDPipeline", diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py index 884710f2e..9cecc31fd 100644 --- a/tests/model_tests/diffusion/test_common_offline.py +++ b/tests/model_tests/diffusion/test_common_offline.py @@ -24,7 +24,7 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.xdist] @pytest.mark.parametrize( - "model_name,accelerations,supported_tasks,check_multioutput,check_determinism", + "model_name,accelerations,supported_tasks,check_multioutput,check_determinism,check_i2v_t2v_divergence", get_parametrized_options(DIFFUSION_TEST_SETTINGS), ) def test_pipeline_on_supported_tasks( @@ -33,6 +33,7 @@ def test_pipeline_on_supported_tasks( supported_tasks: list[DiffusionTasks], check_multioutput: bool, check_determinism: bool, + check_i2v_t2v_divergence: bool, tiny_model_paths: dict[str, str], subtests, ): @@ -61,7 +62,7 @@ def test_pipeline_on_supported_tasks( elif task_type == DiffusionTasks.TEXT_TO_VIDEO: run_and_validate_text_to_video_request(omni) elif task_type == DiffusionTasks.IMAGE_TO_VIDEO: - run_and_validate_image_to_video_request(omni) + run_and_validate_image_to_video_request(omni, check_t2v_divergence=check_i2v_t2v_divergence) else: raise ValueError(f"Task type {task_type} is not yet supported") diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index 006af60e1..00c2399fa 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -230,6 +230,11 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): def _raise_unsupported_features(self) -> None: """Raise an error for incompatible feature switches.""" pc = self.od_config.parallel_config + if pc.tensor_parallel_size > 1: + raise NotImplementedError( + "Tensor parallel is not supported with the diffusers backend. " + "It requires model-specific layer sharding; use a native pipeline for TP." + ) if pc.cfg_parallel_size > 1: raise NotImplementedError( "CFG parallel is not supported with the diffusers backend. " 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 a1fa38d43..32a28ea60 100644 --- a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py @@ -37,6 +37,11 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, ) from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _wrap +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size, + model_parallel_is_initialized, +) from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders @@ -49,7 +54,7 @@ from vllm_omni.diffusion.request import resolve_video_num_frames from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from .pipeline_output import SanaVideoPipelineOutput -from .transformer_sana_video import SanaVideoTransformer3DModel +from .transformer_sana_video import SanaVideoTransformer3DModel, validate_sana_video_parallel_config def _resolve_vae_class_and_dtype( @@ -269,6 +274,7 @@ def retrieve_timesteps( class SanaVideoPipeline( nn.Module, + CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery, @@ -308,6 +314,9 @@ class SanaVideoPipeline( ): super().__init__() + if od_config is not None: + validate_sana_video_parallel_config(od_config.parallel_config) + self.od_config = od_config self.device = get_local_device() self.weights_sources: list[DiffusersPipelineLoader.ComponentSource] = [] @@ -724,6 +733,105 @@ class SanaVideoPipeline( def interrupt(self): return self._interrupt + def predict_noise(self, *args, **kwargs) -> torch.Tensor: + # return_dict=False: the dataclass output is not subscriptable. + return self.transformer(*args, **kwargs, return_dict=False)[0] + + def combine_cfg_noise( + self, + positive_noise_pred, + negative_noise_pred, + true_cfg_scale, + cfg_normalize=False, + kwargs=None, + ): + # fp32 combine so the bf16 all-gather path matches CFG1 bit-for-bit. + return super().combine_cfg_noise( + tuple(p.float() for p in _wrap(positive_noise_pred)), + tuple(n.float() for n in _wrap(negative_noise_pred)), + true_cfg_scale, + cfg_normalize, + kwargs, + ) + + def diffuse( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + negative_prompt_attention_mask: torch.Tensor | None, + guidance_scale: float, + extra_step_kwargs: dict, + dtype: torch.dtype, + output_slice: int | None, + ) -> torch.Tensor: + do_true_cfg = self.do_classifier_free_guidance + # Single-process runs never initialize the parallel groups. + cfg_parallel = model_parallel_is_initialized() and get_classifier_free_guidance_world_size() > 1 + if cfg_parallel: + self.check_cfg_parallel_validity(guidance_scale, negative_prompt_embeds is not None) + + if do_true_cfg and not cfg_parallel: + # Concatenate neg/pos for a single batch-2 forward. + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + with self.progress_bar(total=len(timesteps)) as progress_bar: + for t in timesteps: + if self.interrupt: + continue + + if cfg_parallel: + timestep = t.expand(latents.shape[0]) + latent_model_input = latents.to(dtype) + positive_kwargs = { + "hidden_states": latent_model_input, + "encoder_hidden_states": prompt_embeds.to(dtype), + "encoder_attention_mask": prompt_attention_mask, + "timestep": timestep, + } + negative_kwargs = ( + { + "hidden_states": latent_model_input, + "encoder_hidden_states": negative_prompt_embeds.to(dtype), + "encoder_attention_mask": negative_prompt_attention_mask, + "timestep": timestep, + } + if do_true_cfg + else None + ) + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + output_slice=output_slice, + ).float() + else: + latent_model_input = torch.cat([latents] * 2) if do_true_cfg else latents + timestep = t.expand(latent_model_input.shape[0]) + noise_pred = self.transformer( + latent_model_input.to(dtype), + encoder_hidden_states=prompt_embeds.to(dtype), + encoder_attention_mask=prompt_attention_mask, + timestep=timestep, + return_dict=False, + )[0] + noise_pred = noise_pred.float() + if do_true_cfg: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + if output_slice is not None: + noise_pred = noise_pred[:, :output_slice] + + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + progress_bar.update() + + return latents + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: transformer_weights = ( (name.removeprefix("transformer."), tensor) for name, tensor in weights if name.startswith("transformer.") @@ -947,10 +1055,6 @@ class SanaVideoPipeline( max_sequence_length=max_sequence_length, complex_human_instruction=complex_human_instruction, ) - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) - # 4. Prepare timesteps timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, device, timesteps, sigmas @@ -974,44 +1078,21 @@ class SanaVideoPipeline( extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 7. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) self._num_timesteps = len(timesteps) - - transformer_dtype = self.transformer.dtype - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - - # predict noise model_output - noise_pred = self.transformer( - latent_model_input.to(dtype=transformer_dtype), - encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), - encoder_attention_mask=prompt_attention_mask, - timestep=timestep, - return_dict=False, - )[0] - noise_pred = noise_pred.float() - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # learned sigma - if self.transformer.config.out_channels // 2 == latent_channels: - noise_pred = noise_pred.chunk(2, dim=1)[0] - - # compute previous image: x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() + # learned sigma: transformer predicts 2x latent channels; keep the first half. + output_slice = latent_channels if self.transformer.config.out_channels // 2 == latent_channels else None + latents = self.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_attention_mask=negative_prompt_attention_mask, + guidance_scale=guidance_scale, + extra_step_kwargs=extra_step_kwargs, + dtype=self.transformer.dtype, + output_slice=output_slice, + ) if output_type == "latent": video = latents diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py index 8d0fe9919..9475f306b 100644 --- a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py @@ -16,6 +16,10 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, ) from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size, + model_parallel_is_initialized, +) from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( ASPECT_RATIO_480_BIN, @@ -225,6 +229,91 @@ class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput): ) return DiffusionOutput(output=video) + def diffuse( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + negative_prompt_attention_mask: torch.Tensor | None, + guidance_scale: float, + conditioning_mask: torch.Tensor, + extra_step_kwargs: dict, + dtype: torch.dtype, + output_slice: int | None, + ) -> torch.Tensor: + do_true_cfg = self.do_classifier_free_guidance + # Single-process runs never initialize the parallel groups. + cfg_parallel = model_parallel_is_initialized() and get_classifier_free_guidance_world_size() > 1 + if cfg_parallel: + self.check_cfg_parallel_validity(guidance_scale, negative_prompt_embeds is not None) + + if do_true_cfg and not cfg_parallel: + # Concatenate neg/pos and the mask for a single batch-2 forward. + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) + conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) + + with self.progress_bar(total=len(timesteps)) as progress_bar: + for timestep_value in timesteps: + # Conditioning frames keep timestep 0 so they are not denoised. + timestep = timestep_value.expand(conditioning_mask.shape) * (1 - conditioning_mask) + + if cfg_parallel: + latent_model_input = latents.to(dtype) + positive_kwargs = { + "hidden_states": latent_model_input, + "encoder_hidden_states": prompt_embeds.to(dtype), + "encoder_attention_mask": prompt_attention_mask, + "timestep": timestep, + } + negative_kwargs = ( + { + "hidden_states": latent_model_input, + "encoder_hidden_states": negative_prompt_embeds.to(dtype), + "encoder_attention_mask": negative_prompt_attention_mask, + "timestep": timestep, + } + if do_true_cfg + else None + ) + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + output_slice=output_slice, + ).float() + else: + latent_input = torch.cat([latents, latents]) if do_true_cfg else latents + noise_pred = self.transformer( + latent_input.to(dtype), + encoder_hidden_states=prompt_embeds.to(dtype), + encoder_attention_mask=prompt_attention_mask, + timestep=timestep, + return_dict=False, + )[0].float() + if do_true_cfg: + uncond, cond = noise_pred.chunk(2) + noise_pred = uncond + guidance_scale * (cond - uncond) + if output_slice is not None: + noise_pred = noise_pred[:, :output_slice] + + # Only the non-conditioning frames are denoised; the first frame is carried through. + denoised = self.scheduler.step( + noise_pred[:, :, 1:], + timestep_value, + latents[:, :, 1:], + **extra_step_kwargs, + return_dict=False, + )[0] + latents = torch.cat([latents[:, :, :1], denoised], dim=2) + progress_bar.update() + + return latents + @torch.no_grad() def _generate_i2v( self, @@ -282,10 +371,6 @@ class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput): max_sequence_length=max_sequence_length, complex_human_instruction=complex_human_instruction, ) - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_embeds, prompt_embeds]) - prompt_mask = torch.cat([negative_mask, prompt_mask]) - timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps=num_inference_steps, @@ -319,39 +404,24 @@ class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput): latents.shape[4] // patch_w, ) conditioning_mask[:, :, 0] = 1 - if self.do_classifier_free_guidance: - conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) self._num_timesteps = len(timesteps) extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - transformer_dtype = self.transformer.dtype - with self.progress_bar(total=num_inference_steps) as progress_bar: - for index, timestep_value in enumerate(timesteps): - latent_input = torch.cat([latents, latents]) if self.do_classifier_free_guidance else latents - timestep = timestep_value.expand(conditioning_mask.shape) * (1 - conditioning_mask) - noise_pred = self.transformer( - latent_input.to(transformer_dtype), - encoder_hidden_states=prompt_embeds.to(transformer_dtype), - encoder_attention_mask=prompt_mask, - timestep=timestep, - return_dict=False, - )[0].float() - if self.do_classifier_free_guidance: - uncond, cond = noise_pred.chunk(2) - noise_pred = uncond + guidance_scale * (cond - uncond) - if self.transformer.config.out_channels // 2 == latent_channels: - noise_pred = noise_pred.chunk(2, dim=1)[0] - - denoised = self.scheduler.step( - noise_pred[:, :, 1:], - timestep_value, - latents[:, :, 1:], - **extra_step_kwargs, - return_dict=False, - )[0] - latents = torch.cat([latents[:, :, :1], denoised], dim=2) - if index == len(timesteps) - 1 or (index + 1) % self.scheduler.order == 0: - progress_bar.update() + # learned sigma: transformer predicts 2x latent channels; keep the first half. + output_slice = latent_channels if self.transformer.config.out_channels // 2 == latent_channels else None + latents = self.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_mask, + negative_prompt_embeds=negative_embeds, + negative_prompt_attention_mask=negative_mask, + guidance_scale=guidance_scale, + conditioning_mask=conditioning_mask, + extra_step_kwargs=extra_step_kwargs, + dtype=self.transformer.dtype, + output_slice=output_slice, + ) if output_type == "latent": return latents 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 eb0f5843b..db38cc099 100644 --- a/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py +++ b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py @@ -19,12 +19,48 @@ from dataclasses import dataclass, fields import torch from torch import nn +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.models.utils import AutoWeightsLoader +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 +def validate_sana_video_parallel_config(parallel_config) -> None: + """Reject unsupported SANA-Video parallel topologies before weights load.""" + tp_size = parallel_config.tensor_parallel_size + if tp_size not in (1, 2): + raise NotImplementedError( + f"SANA-Video supports tensor_parallel_size 1 or 2, got {tp_size}. Set --tensor-parallel-size to 1 or 2." + ) + cfg_size = parallel_config.cfg_parallel_size + if cfg_size not in (1, 2): + 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: + 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." + ) + 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: + raise NotImplementedError("SANA-Video does not support HSDP. Remove --use-hsdp.") + if parallel_config.text_encoder_tp_size > 1: + raise NotImplementedError( + "SANA-Video does not support text encoder tensor parallel. Set text_encoder_tp_size to 1." + ) + + @dataclass class SanaVideoTransformerConfig: in_channels: int = 16 @@ -109,6 +145,54 @@ class SanaRMSNorm(nn.Module): return hidden_states +class SanaDistributedRMSNorm(nn.Module): + """RMSNorm that computes global RMS across tensor parallel ranks.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + if param.shape == loaded_weight.shape: + param.data.copy_(loaded_weight) + return + + tp_size = get_tensor_model_parallel_world_size() + if loaded_weight.shape[0] % tp_size != 0: + raise ValueError( + f"Cannot shard RMSNorm weight of shape {tuple(loaded_weight.shape)} across tp_size={tp_size}." + ) + + shard_size = loaded_weight.shape[0] // tp_size + start_idx = get_tensor_model_parallel_rank() * shard_size + shard = loaded_weight.narrow(0, start_idx, shard_size) + if param.shape != shard.shape: + raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.") + param.data.copy_(shard) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + x_float = x.float() + sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True) + count = x.shape[-1] + + if tp_size > 1: + # Use vLLM's collective (custom all-reduce / symmetric-mem fast path + # for small tensors) instead of raw torch.distributed.all_reduce, and + # take the return value so the custom-AR path (which may return a new + # buffer) is handled correctly. No .clone() needed. + sum_sq = tensor_model_parallel_all_reduce(sum_sq) + count = count * tp_size + + hidden_states = x_float * torch.rsqrt(sum_sq / count + self.eps) + # Only cast to the weight dtype for fp16/bf16 weights; keep fp32 otherwise. + if self.weight.dtype in (torch.float16, torch.bfloat16): + hidden_states = hidden_states.to(self.weight.dtype) + return hidden_states * self.weight + + def _get_timestep_embedding( timesteps: torch.Tensor, embedding_dim: int, @@ -378,11 +462,22 @@ class GLUMBTempConv(nn.Module): hidden_channels = int(expand_ratio * in_channels) self.norm_type = norm_type self.residual_connection = residual_connection + self.hidden_channels = hidden_channels + + tp_size = get_tensor_model_parallel_world_size() + local_hidden_channels = hidden_channels // tp_size self.nonlinearity = nn.SiLU() - self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0) - self.conv_depth = nn.Conv2d(hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2) - self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) + self.conv_inverted = nn.Conv2d(in_channels, 2 * local_hidden_channels, 1, 1, 0) + self.conv_depth = nn.Conv2d( + 2 * local_hidden_channels, 2 * local_hidden_channels, 3, 1, 1, groups=2 * local_hidden_channels + ) + self.conv_point = nn.Conv2d(local_hidden_channels, out_channels, 1, 1, 0, bias=False) + # conv_inverted/conv_depth carry [value(h), gate(h)]; each rank must hold a + # slice of both segments so the local chunk(2) still splits value from gate. + for param in (self.conv_inverted.weight, self.conv_inverted.bias, self.conv_depth.weight, self.conv_depth.bias): + set_weight_attrs(param, {"weight_loader": self._segment_weight_loader}) + set_weight_attrs(self.conv_point.weight, {"weight_loader": self._point_weight_loader}) self.norm = None if norm_type == "rms_norm": @@ -390,6 +485,19 @@ class GLUMBTempConv(nn.Module): self.conv_temp = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 1), stride=1, padding=(1, 0), bias=False) + def _segment_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + h = self.hidden_channels + local = h // get_tensor_model_parallel_world_size() + start = get_tensor_model_parallel_rank() * local + value = loaded_weight.narrow(0, start, local) + gate = loaded_weight.narrow(0, h + start, local) + param.data.copy_(torch.cat([value, gate], dim=0)) + + def _point_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + local = self.hidden_channels // get_tensor_model_parallel_world_size() + start = get_tensor_model_parallel_rank() * local + param.data.copy_(loaded_weight.narrow(1, start, local)) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.residual_connection: residual = hidden_states @@ -404,6 +512,8 @@ class GLUMBTempConv(nn.Module): hidden_states = hidden_states * self.nonlinearity(gate) hidden_states = self.conv_point(hidden_states) + if get_tensor_model_parallel_world_size() > 1: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Temporal aggregation hidden_states_temporal = hidden_states.view(batch_size, num_frames, num_channels, height * width).permute( @@ -440,13 +550,20 @@ class SanaLinearAttention(nn.Module): ) -> None: super().__init__() inner_dim = num_heads * head_dim - self.heads = num_heads - self.to_q = nn.Linear(dim, inner_dim, bias=bias) - self.to_k = nn.Linear(dim, inner_dim, bias=bias) - self.to_v = nn.Linear(dim, inner_dim, bias=bias) - self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None - self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None - self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=True), nn.Dropout(dropout)]) + tp_size = get_tensor_model_parallel_world_size() + self.heads = num_heads // tp_size + local_inner_dim = inner_dim // tp_size + self.to_q = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False) + self.to_k = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False) + self.to_v = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False) + self.norm_q = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None + self.norm_k = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None + self.to_out = nn.ModuleList( + [ + RowParallelLinear(inner_dim, dim, bias=True, input_is_parallel=True, return_bias=False), + nn.Dropout(dropout), + ] + ) def forward( self, @@ -632,18 +749,29 @@ class SanaCrossAttention(nn.Module): ) -> None: super().__init__() inner_dim = num_heads * head_dim - self.heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + self.heads = num_heads // tp_size self.head_dim = head_dim - self.to_q = nn.Linear(dim, inner_dim, bias=bias) - self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias) - self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias) - self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None - self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None - self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=out_bias), nn.Dropout(dropout)]) + local_inner_dim = inner_dim // tp_size + self.to_q = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False) + self.to_k = ColumnParallelLinear( + cross_attention_dim, inner_dim, bias=bias, gather_output=False, return_bias=False + ) + self.to_v = ColumnParallelLinear( + cross_attention_dim, inner_dim, bias=bias, gather_output=False, return_bias=False + ) + self.norm_q = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None + self.norm_k = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None + self.to_out = nn.ModuleList( + [ + RowParallelLinear(inner_dim, dim, bias=out_bias, input_is_parallel=True, return_bias=False), + nn.Dropout(dropout), + ] + ) self.attn = OmniAttention( - num_heads=num_heads, + num_heads=self.heads, head_size=head_dim, - num_kv_heads=num_heads, + num_kv_heads=self.heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross",