diff --git a/docs/.nav.yml b/docs/.nav.yml index 68d019d9c..1e3907a6e 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -68,7 +68,7 @@ nav: - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md - - VAE Patch Parallel: user_guide/diffusion/parallelism/vae_patch_parallel.md + - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md - LoRA: user_guide/diffusion/lora.md - Custom Pipeline: features/custom_pipeline.md diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md index e330b41a6..00f41378a 100644 --- a/docs/design/feature/vae_parallel.md +++ b/docs/design/feature/vae_parallel.md @@ -448,6 +448,42 @@ Complete examples in the codebase: --- +## Spatially-Sharded Decode (Wan) + +The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`). + +### How it differs from tile parallel + +| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) | +|--------|--------------------------|-----------------------------------------------| +| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W | +| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) | +| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) | +| Scope | Decode + encode | Decode only | + +Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`. + +### Wiring + +`vae_parallel_mode` flows through the same path as `vae_patch_parallel_size`: + +```text +serve.py (--vae-parallel-mode) / OmniEngineArgs + -> DiffusionParallelConfig.vae_parallel_mode + -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...) + -> DistributedVaeExecutor.parallel_mode + -> DistributedAutoencoderKLWan.tiled_decode dispatch +``` + +`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode. + +### Notes + +- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance. +- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`). + +--- + ## Summary Adding VAE Patch Parallel support to diffusion model: diff --git a/docs/user_guide/diffusion/parallelism/overview.md b/docs/user_guide/diffusion/parallelism/overview.md index d482b5ea2..8419b4800 100644 --- a/docs/user_guide/diffusion/parallelism/overview.md +++ b/docs/user_guide/diffusion/parallelism/overview.md @@ -10,7 +10,7 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion | **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos | | **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation | | **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory | -| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory | +| **[VAE Parallelism](vae_parallelism.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory | | **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs | | **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) | diff --git a/docs/user_guide/diffusion/parallelism/vae_patch_parallel.md b/docs/user_guide/diffusion/parallelism/vae_parallelism.md similarity index 71% rename from docs/user_guide/diffusion/parallelism/vae_patch_parallel.md rename to docs/user_guide/diffusion/parallelism/vae_parallelism.md index 4e8513eab..ca0c428e1 100644 --- a/docs/user_guide/diffusion/parallelism/vae_patch_parallel.md +++ b/docs/user_guide/diffusion/parallelism/vae_parallelism.md @@ -1,4 +1,4 @@ -# VAE Patch Parallelism Guide +# VAE Parallelism Guide ## Table of Content @@ -15,7 +15,7 @@ ## Overview -VAE Patch Parallelism distributes the VAE (Variational AutoEncoder) decode/encode computation across multiple GPUs by splitting the latent space into spatial tiles or patches. Each GPU processes a subset of tiles in parallel, significantly reducing peak memory consumption during the VAE decode stage while maintaining output quality. +VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width. This is particularly useful for: - **High-resolution image generation** where VAE decode can become a memory bottleneck @@ -25,7 +25,7 @@ This is particularly useful for: See supported models list in [Supported Models](../../diffusion_features.md#supported-models). -VAE Patch Parallelism uses two strategies based on image size: +VAE patch parallelism uses two strategies based on image size: | Strategy | Use Case | How It Works | Overlap Handling | Output Quality | |----------|----------|--------------|------------------|----------------| @@ -113,6 +113,7 @@ In `DiffusionParallelConfig`: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. | +| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). | Additional requirements: @@ -125,6 +126,51 @@ Additional requirements: --- +## Spatially-Sharded Decode (Wan) + +The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`. + +Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance. + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + vae_patch_parallel_size=2, # must match the DiT group size + vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height" + ), +) +``` + +Or from the CLI / serving entrypoint: + +```bash +vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \ + --tensor-parallel-size 2 \ + --vae-patch-parallel-size 2 \ + --vae-parallel-mode spatial_shard_width +``` + +**Constraints and behavior:** + +- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes. +- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime. +- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension). + +For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark: + +```bash +python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ + --endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \ + --height 480 --width 832 --num-frames 17 --max-concurrency 1 +``` + +--- + ## Best Practices ### When to Use diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 0f96aa2c6..fef039784 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -61,7 +61,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on |--------|-------------|----------| | **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs | | **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss | -| **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak | +| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak | ### Extensions @@ -270,7 +270,7 @@ Measured on NVIDIA H800: **Memory Optimization:** - **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage -- **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images +- **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos - **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models **Extensions:** diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index c8cc8a3f4..8f15347d1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -423,6 +423,7 @@ def test_sub_config_fields_match_rfc_scopes(): "ulysses_mode", "cfg_parallel_size", "vae_patch_parallel_size", + "vae_parallel_mode", "use_hsdp", "mask_sp_padding", "hsdp_shard_size", diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py new file mode 100644 index 000000000..6d668df3a --- /dev/null +++ b/tests/diffusion/distributed/test_wan_spatial_shard.py @@ -0,0 +1,417 @@ +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, +) +from vllm_omni.platforms import current_omni_platform + +# Module-level markers apply to the CPU unit tests below. The multi-GPU +# correctness test at the end of the file adds the nightly/distributed markers +# and a skip guard so it only runs where >= 2 CUDA devices are available. +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_split_for_parallel_decode_pads_uneven_height(): + x = torch.arange(1 * 1 * 1 * 5 * 2, dtype=torch.float32).reshape(1, 1, 1, 5, 2) + + local, expected_height = wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=2, + rank=2, + world_size=3, + ) + + assert expected_height == 20 + assert local.shape == (1, 1, 1, 2, 2) + assert torch.equal(local[..., 0, :], x[..., 4, :]) + assert torch.equal(local[..., 1, :], torch.zeros_like(local[..., 1, :])) + + +def test_split_for_parallel_decode_pads_uneven_width(): + x = torch.arange(1 * 1 * 1 * 2 * 5, dtype=torch.float32).reshape(1, 1, 1, 2, 5) + + local, expected_width = wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=2, + split_dim="width", + rank=2, + world_size=3, + ) + + assert expected_width == 20 + assert local.shape == (1, 1, 1, 2, 2) + assert torch.equal(local[..., :, 0], x[..., :, 4]) + assert torch.equal(local[..., :, 1], torch.zeros_like(local[..., :, 1])) + + +def test_split_for_parallel_decode_rejects_invalid_split_dim(): + x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32) + + with pytest.raises(ValueError, match="split_dim"): + wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=1, + split_dim="depth", + rank=0, + world_size=2, + ) + + +def test_split_for_parallel_decode_rejects_zero_world_size(): + x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32) + + with pytest.raises(ValueError, match="world_size"): + wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=1, + rank=0, + world_size=0, + ) + + +def test_split_for_parallel_decode_rejects_rank_out_of_range(): + x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32) + + with pytest.raises(ValueError, match="rank"): + wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=1, + rank=3, + world_size=3, + ) + + +def test_gather_and_trim_height(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3)) + + def fake_all_gather(gathered, x, group=None): + for idx, output in enumerate(gathered): + output.copy_(x + idx) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.zeros((1, 1, 1, 2, 1), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object()) + + assert out.shape == (1, 1, 1, 5, 1) + assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0])) + + +def test_gather_and_trim_width(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3)) + + def fake_all_gather(gathered, x, group=None): + for idx, output in enumerate(gathered): + output.copy_(x + idx) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="width", group=object()) + + assert out.shape == (1, 1, 1, 1, 5) + assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0])) + + +def test_gather_and_trim_rank0_only_assembles_on_rank0(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3)) + + def fake_all_gather(gathered, x, group=None): + for idx, output in enumerate(gathered): + output.copy_(x + idx) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.zeros((1, 1, 1, 2, 1), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object(), dst=0) + + assert out.shape == (1, 1, 1, 5, 1) + assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0])) + + +def test_gather_and_trim_rank0_only_returns_empty_on_non_zero_rank(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) + + gathered_sizes = [] + + def fake_all_gather(gathered, x, group=None): + # Every rank must still take part in the collective even when it discards the result. + gathered_sizes.append(len(gathered)) + for output in gathered: + output.copy_(x) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.ones((1, 1, 1, 2, 1), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object(), dst=0) + + assert gathered_sizes == [3] + assert out.numel() == 0 + + +def test_reshard_from_trimmed_height_pads_invalid_rows(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + + x = torch.arange(5, dtype=torch.float32).reshape(1, 1, 1, 5, 1) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=5, + local_input_extent=2, + split_dim="height", + rank=2, + world_size=3, + ) + ) + try: + out = wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=2, + split_dim="height", + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + assert out.shape == (1, 1, 1, 2, 1) + assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0])) + + +def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + + x = torch.arange(5, dtype=torch.float32).reshape(1, 1, 1, 1, 5) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=5, + local_input_extent=2, + split_dim="width", + rank=2, + world_size=3, + ) + ) + try: + out = wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=2, + split_dim="width", + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + assert out.shape == (1, 1, 1, 1, 2) + assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0])) + + +def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1)) + + x = torch.randn((1, 1, 1, 4, 2)) + out, recv_top, recv_bottom = wan_spatial_shard.halo_exchange( + x, + group=object(), + halo_size=1, + ) + + assert out is x + assert recv_top is None + assert recv_bottom is None + + +def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.MonkeyPatch): + x = torch.ones((1, 1, 2, 2)) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) + mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object()) + mid = mid_rank_pad(x) + assert mid.shape == (1, 1, 2, 3) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object()) + last = last_rank_pad(x) + assert last.shape == (1, 1, 3, 3) + + +def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.MonkeyPatch): + x = torch.ones((1, 1, 2, 2)) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) + mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width") + mid = mid_rank_pad(x) + assert mid.shape == (1, 1, 2, 2) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width") + last = last_rank_pad(x) + assert last.shape == (1, 1, 2, 3) + + +def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size", + lambda group=None: 4, + ) + + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + vae.use_tiling = True + vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="spatial_shard_height") + vae.is_distributed_enabled = lambda: True + + z = torch.zeros((1, 16, 1, 8, 8)) + + assert vae._spatial_shard_decode_split_dim() == "height" + assert vae._spatial_shard_decode_enabled(z) is False + + +def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size", + lambda group=None: 2, + ) + + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="spatial_shard_width") + vae.is_distributed_enabled = lambda: True + + z = torch.zeros((1, 16, 1, 8, 8)) + + assert vae._spatial_shard_decode_split_dim() == "width" + assert vae._spatial_shard_decode_enabled(z) is True + + +def test_tile_mode_disables_spatial_shard_decode(): + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="tile") + vae.is_distributed_enabled = lambda: True + + z = torch.zeros((1, 16, 1, 8, 8)) + + assert vae._spatial_shard_decode_split_dim() is None + assert vae._spatial_shard_decode_enabled(z) is False + + +# ============================================================================= +# Multi-GPU numerical-correctness test (nightly / full_model, distributed CUDA) +# +# Spawns a small process group and verifies that spatial_shard_height/spatial_shard_width decode match +# a single-process (non-distributed) reference decode of the same latent within +# tolerance. Requires >= 2 accelerator devices and downloads a Wan VAE. +# ============================================================================= + +_SPATIAL_SHARD_MODEL = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +_SPATIAL_SHARD_SUBFOLDER = "vae" +_SPATIAL_SHARD_WORLD_SIZE = 2 +_SPATIAL_SHARD_LATENT_FRAMES = 5 +_SPATIAL_SHARD_LATENT_HEIGHT = 60 +_SPATIAL_SHARD_LATENT_WIDTH = 104 +_SPATIAL_SHARD_TOLERANCE = 3e-2 + + +def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_port: str) -> None: + from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, + ) + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = master_port + device = current_omni_platform.get_torch_device(rank) + current_omni_platform.set_device(device) + dtype = torch.float32 + + backend = current_omni_platform.dist_backend + init_distributed_environment(world_size=_SPATIAL_SHARD_WORLD_SIZE, rank=rank, local_rank=rank, backend=backend) + initialize_model_parallel( + sequence_parallel_size=_SPATIAL_SHARD_WORLD_SIZE, ulysses_degree=_SPATIAL_SHARD_WORLD_SIZE, backend=backend + ) + + try: + vae = DistributedAutoencoderKLWan.from_pretrained( + _SPATIAL_SHARD_MODEL, subfolder=_SPATIAL_SHARD_SUBFOLDER, torch_dtype=dtype + ) + vae.to(device=device, dtype=dtype) + vae.eval() + + generator = torch.Generator(device=device).manual_seed(0) + latents = torch.randn( + ( + 1, + vae.config.z_dim, + _SPATIAL_SHARD_LATENT_FRAMES, + _SPATIAL_SHARD_LATENT_HEIGHT, + _SPATIAL_SHARD_LATENT_WIDTH, + ), + generator=generator, + device=device, + dtype=dtype, + ) + + with torch.inference_mode(): + # Ground-truth reference: standard non-parallel, untiled decode (computed identically on + # every rank). Tiling must be OFF so neither the tile-parallel nor the single-GPU tiled + # path is exercised; otherwise we would be comparing SP against tiled decode. + vae.use_tiling = False + vae.set_parallel_size(1, mode="tile") + reference = vae.decode(latents, return_dict=False)[0].float() + + # Spatially-sharded decode across the full group (requires tiling to enter tiled_decode). + vae.use_tiling = True + vae.set_parallel_size(_SPATIAL_SHARD_WORLD_SIZE, mode=f"spatial_shard_{split_dim}") + sharded = vae.decode(latents, return_dict=False)[0].float() + + # Only rank 0 assembles the full decoded sample (matching broadcast_result=False); + # non-zero ranks return an empty placeholder, so the comparison runs on rank 0 only. + if rank == 0: + diff = (sharded - reference).abs() + return_dict["max_abs_diff"] = diff.max().item() + return_dict["mean_abs_diff"] = diff.mean().item() + return_dict["shape"] = tuple(sharded.shape) + finally: + destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.distributed_cuda +@pytest.mark.skipif( + current_omni_platform.get_device_count() < _SPATIAL_SHARD_WORLD_SIZE, + reason="Requires >= 2 accelerator devices", +) +@pytest.mark.parametrize("split_dim", ["height", "width"]) +def test_spatial_shard_decode_matches_reference(split_dim: str): + manager = mp.get_context("spawn").Manager() + return_dict = manager.dict() + # Use a per-split-dim port to avoid collisions across parametrized runs. + master_port = str(29500 + (1 if split_dim == "width" else 0)) + + mp.spawn( + _spatial_shard_decode_worker, + args=(split_dim, return_dict, master_port), + nprocs=_SPATIAL_SHARD_WORLD_SIZE, + join=True, + ) + + assert "max_abs_diff" in return_dict, "rank 0 did not report a result" + max_abs_diff = return_dict["max_abs_diff"] + mean_abs_diff = return_dict["mean_abs_diff"] + print( + f"spatial_shard_{split_dim} vs reference: max_abs_diff={max_abs_diff:.6e} " + f"mean_abs_diff={mean_abs_diff:.6e} shape={return_dict.get('shape')}" + ) + assert max_abs_diff <= _SPATIAL_SHARD_TOLERANCE, ( + f"spatial_shard_{split_dim} max_abs_diff {max_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}" + ) + assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, ( + f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}" + ) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 4ab8858fa..02b413064 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -353,6 +353,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): ulysses_mode: str = "strict" cfg_parallel_size: int = Field(default=1, ge=1, le=3) vae_patch_parallel_size: int = Field(default=1, ge=1) + vae_parallel_mode: str = "tile" use_hsdp: bool = False mask_sp_padding: bool = False hsdp_shard_size: int = -1 @@ -362,6 +363,11 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): self.sequence_parallel_size = self.ulysses_degree * self.ring_degree if self.ulysses_mode not in {"strict", "advanced_uaa"}: raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") + if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}: + raise ValueError( + "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + f"but got {self.vae_parallel_mode!r}." + ) other_parallel_world_size = ( self.pipeline_parallel_size diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 30d9081fd..d1c728732 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -338,6 +338,7 @@ class StageDeployConfig: sequence_parallel_size: int | None = None cfg_parallel_size: int | None = None vae_patch_parallel_size: int | None = None + vae_parallel_mode: str | None = None use_hsdp: bool | None = None hsdp_shard_size: int | None = None hsdp_replicate_size: int | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index bac5f90da..4c031bc15 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -183,6 +183,21 @@ class DiffusionParallelConfig: vae_patch_parallel_size: int = 1 """Number of ranks used for VAE patch/tile parallelism (decode/encode).""" + vae_parallel_mode: str = "tile" + """VAE parallel decode strategy. + + - "tile": Patch/tile parallel decode (default). Each rank decodes a subset + of spatial tiles and the results are stitched on rank 0. + - "spatial_shard_height": Spatially-sharded decode that splits decoder + feature maps along height and exchanges halo rows around spatial + convolutions. + - "spatial_shard_width": Same as "spatial_shard_height" but sharded along width. + + The "spatial_shard_*" modes are decode-only and currently require + ``vae_patch_parallel_size`` to match the DiT group size; otherwise the VAE + falls back to tile-parallel decode at runtime. + """ + use_hsdp: bool = False """Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding.""" @@ -214,6 +229,10 @@ class DiffusionParallelConfig: f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}" ) assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0" + assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, ( + "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + f"but got {self.vae_parallel_mode!r}." + ) assert self.sequence_parallel_size == self.ulysses_degree * self.ring_degree, ( "Sequence parallel size must be equal to the product of ulysses degree and ring degree," f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}" diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 35c9434d0..44c6efe66 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -5,11 +5,13 @@ from contextlib import nullcontext from typing import Any import torch +import torch.distributed as dist from diffusers.models.autoencoders import AutoencoderKLWan from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify from diffusers.models.autoencoders.vae import DecoderOutput from vllm.logger import init_logger +from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedOperator, DistributedVaeMixin, @@ -268,10 +270,54 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): dec = torch.clamp(dec, min=-1.0, max=1.0) return dec + def _spatial_shard_decode_split_dim(self) -> str | None: + mode = self.distributed_executor.parallel_mode + if mode == "spatial_shard_width": + return "width" + if mode == "spatial_shard_height": + return "height" + return None + + def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool: + split_dim = self._spatial_shard_decode_split_dim() + if split_dim is None: + return False + if z.ndim != 5: + logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.") + return False + if not self.is_distributed_enabled(): + return False + + group = self.distributed_executor.group + world_size = dist.get_world_size(group=group) + requested_size = int(self.distributed_executor.parallel_size) + if requested_size != world_size: + logger.warning( + "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size " + "to match the DIT group size; falling back to tiled decode. " + "requested=%s dit_group=%s split_dim=%s", + requested_size, + world_size, + split_dim, + ) + return False + return True + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): if not self.is_distributed_enabled(): return super().tiled_decode(z, return_dict=return_dict) + if self._spatial_shard_decode_enabled(z): + split_dim = self._spatial_shard_decode_split_dim() + logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim) + return wan_spatial_shard.spatial_shard_decode( + self, + z, + group=self.distributed_executor.group, + return_dict=return_dict, + split_dim=split_dim, + ) + logger.debug("Decode running with distributed executor") result = self.distributed_executor.execute( z, diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py index 8fb4ebeb3..d804740f9 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -49,9 +49,11 @@ class DistributedVaeExecutor: self.world_size = dist.get_world_size(self.group) self.rank = dist.get_rank(self.group) self.parallel_size = 1 + self.parallel_mode = "tile" - def set_parallel_size(self, parallel_size: int): + def set_parallel_size(self, parallel_size: int, mode: str = "tile"): self.parallel_size = parallel_size + self.parallel_mode = mode def gather_tensors(self, tensor: torch.Tensor): gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] @@ -170,8 +172,8 @@ class DistributedVaeMixin: def init_distributed(self): self.distributed_executor = DistributedVaeExecutor() - def set_parallel_size(self, parallel_size: int) -> None: - self.distributed_executor.set_parallel_size(parallel_size) + def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None: + self.distributed_executor.set_parallel_size(parallel_size, mode=mode) def is_distributed_enabled(self) -> bool: if ( diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py new file mode 100644 index 000000000..bda728540 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -0,0 +1,816 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# The halo-exchange spatial-parallel decode here is adapted from SGLang's +# spatial-parallel VAE decode +# https://github.com/sgl-project/sglang +# (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py) +# which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo). +# This version generalizes the height-only sharding to shard along height or +# width and adds the Wan causal-conv ``feat_cache`` handling. +"""Spatially-sharded Wan VAE decode. + +The existing distributed Wan VAE path shards *tiles*. This module adds an +opt-in decode backend that shards decoder feature maps along height or width and +exchanges boundary rows/columns before spatial convolutions. It is +intentionally decode-only and keeps checkpoint loading unchanged by patching the +already-loaded decoder. +""" + +from __future__ import annotations + +import math +from contextlib import nullcontext +from contextvars import ContextVar +from dataclasses import dataclass +from types import MethodType +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify +from diffusers.models.autoencoders.vae import DecoderOutput +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class SpatialShardContext: + input_extent: int + local_input_extent: int + split_dim: str + rank: int + world_size: int + + +_SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar( + "wan_vae_spatial_shard_context", + default=None, +) + + +def _spatial_dim(split_dim: str) -> int: + if split_dim == "height": + return -2 + if split_dim == "width": + return -1 + raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.") + + +def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor: + if dim < 0: + dim += x.dim() + return x.narrow(dim, start, length) + + +def _global_rank(group: dist.ProcessGroup, group_rank: int) -> int: + try: + return dist.get_global_rank(group, group_rank) + except Exception: + return group_rank + + +def _rank_world(group: dist.ProcessGroup) -> tuple[int, int]: + return dist.get_rank(group), dist.get_world_size(group) + + +def _pad_along_dim(x: torch.Tensor, pad: int, dim: int, value: float = 0.0) -> torch.Tensor: + if pad <= 0: + return x + shape = list(x.shape) + shape[dim] = pad + padding = torch.full(shape, value, dtype=x.dtype, device=x.device) + return torch.cat([x, padding], dim=dim) + + +def _maybe_contiguous_for_shard_gather(x: torch.Tensor) -> torch.Tensor: + if ( + x.dim() == 5 + and hasattr(torch, "channels_last_3d") + and x.is_contiguous(memory_format=torch.channels_last_3d) + and not x.is_contiguous() + ): + return x.contiguous() + return x + + +def _halo_memory_format(reference: torch.Tensor) -> torch.memory_format: + if reference.dim() > 1 and reference.stride(1) == 1: + if reference.dim() == 5 and hasattr(torch, "channels_last_3d"): + return torch.channels_last_3d + if reference.dim() == 4: + return torch.channels_last + return torch.contiguous_format + + +def _current_full_extent(local_extent: int) -> int | None: + ctx = _SPATIAL_SHARD_CONTEXT.get() + if ctx is None: + return None + if ctx.local_input_extent <= 0: + return None + scale = local_extent / ctx.local_input_extent + rounded_scale = round(scale) + if not math.isclose(scale, rounded_scale, rel_tol=0.0, abs_tol=1e-6): + return None + return ctx.input_extent * rounded_scale + + +def _local_valid_extent(local_extent: int) -> int: + ctx = _SPATIAL_SHARD_CONTEXT.get() + full_extent = _current_full_extent(local_extent) + if ctx is None or full_extent is None: + return local_extent + start = ctx.rank * local_extent + return max(0, min(local_extent, full_extent - start)) + + +def _zero_invalid_extent(x: torch.Tensor, *, split_dim: str) -> torch.Tensor: + dim = _spatial_dim(split_dim) + dim_size = x.shape[dim] + valid_extent = _local_valid_extent(dim_size) + if valid_extent >= dim_size: + return x + x = x.clone() + invalid = _narrow_along_dim(x, dim, valid_extent, dim_size - valid_extent) + invalid.zero_() + return x + + +def split_for_parallel_decode( + x: torch.Tensor, + *, + upsample_count: int, + split_dim: str = "height", + group: dist.ProcessGroup | None = None, + rank: int | None = None, + world_size: int | None = None, +) -> tuple[torch.Tensor, int]: + """Shard latent/feature spatial extent and return expected full output extent.""" + if group is not None: + rank, world_size = _rank_world(group) + rank = 0 if rank is None else int(rank) + world_size = 1 if world_size is None else int(world_size) + if world_size < 1: + raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.") + if not 0 <= rank < world_size: + raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.") + + dim = _spatial_dim(split_dim) + expected_extent = x.shape[dim] * (2**upsample_count) + if world_size <= 1: + return x, expected_extent + + pad = (world_size - (x.shape[dim] % world_size)) % world_size + if pad: + x = _pad_along_dim(x, pad, dim=dim) + chunk_size = x.shape[dim] // world_size + return _narrow_along_dim(x, dim, rank * chunk_size, chunk_size).contiguous(), expected_extent + + +def all_gather_along_dim( + x: torch.Tensor, + *, + group: dist.ProcessGroup, + dim: int, + dst: int | None = None, +) -> torch.Tensor: + rank, world_size = _rank_world(group) + if world_size <= 1: + return x + x = _maybe_contiguous_for_shard_gather(x) + gathered = [torch.empty_like(x) for _ in range(world_size)] + # NCCL has no rank-local gather, so every rank joins the collective; only ``dst`` + # keeps the assembled tensor while the rest drop their copies. + dist.all_gather(gathered, x.contiguous(), group=group) + if dst is not None and rank != dst: + return x.new_zeros(0) + return torch.cat(gathered, dim=dim) + + +def reshard_from_trimmed_extent( + x: torch.Tensor, + *, + local_extent: int, + split_dim: str, + group: dist.ProcessGroup, +) -> torch.Tensor: + rank, world_size = _rank_world(group) + if world_size <= 1: + return x + + dim = _spatial_dim(split_dim) + valid_extent = _local_valid_extent(local_extent) + start = rank * local_extent + local = _narrow_along_dim(x, dim, start, valid_extent).contiguous() + if valid_extent < local_extent: + local = _pad_along_dim(local, local_extent - valid_extent, dim=dim) + return local + + +def gather_and_trim_extent( + x: torch.Tensor, + *, + expected_extent: int | None, + split_dim: str, + group: dist.ProcessGroup, + dst: int | None = None, +) -> torch.Tensor: + dim = _spatial_dim(split_dim) + rank, _ = _rank_world(group) + out = all_gather_along_dim(x, group=group, dim=dim, dst=dst) + if dst is not None and rank != dst: + return out + if expected_extent is not None and out.shape[dim] != expected_extent: + out = _narrow_along_dim(out, dim, 0, expected_extent).contiguous() + return out + + +def _ensure_recv_buf(recv_buf: torch.Tensor | None, reference: torch.Tensor) -> torch.Tensor: + memory_format = _halo_memory_format(reference) + if ( + recv_buf is None + or recv_buf.shape != reference.shape + or recv_buf.dtype != reference.dtype + or recv_buf.device != reference.device + or not recv_buf.is_contiguous(memory_format=memory_format) + ): + return torch.empty( + reference.shape, + dtype=reference.dtype, + device=reference.device, + memory_format=memory_format, + ) + return recv_buf + + +def _halo_exchange_p2p( + *, + rank: int, + world_size: int, + group: dist.ProcessGroup, + top_row_ref: torch.Tensor, + bottom_row_ref: torch.Tensor, + recv_top_buf: torch.Tensor, + recv_bottom_buf: torch.Tensor, +) -> None: + p2p_ops = [] + if rank > 0: + prev_rank = _global_rank(group, rank - 1) + top_row = top_row_ref.contiguous(memory_format=_halo_memory_format(top_row_ref)) + p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group)) + p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group)) + else: + recv_top_buf.zero_() + + if rank < world_size - 1: + next_rank = _global_rank(group, rank + 1) + bottom_row = bottom_row_ref.contiguous(memory_format=_halo_memory_format(bottom_row_ref)) + p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group)) + p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group)) + else: + recv_bottom_buf.zero_() + + if p2p_ops: + reqs = dist.batch_isend_irecv(p2p_ops) + for req in reqs: + req.wait() + + +def halo_exchange( + x: torch.Tensor, + *, + group: dist.ProcessGroup, + halo_size: int, + split_dim: str = "height", + recv_top_buf: torch.Tensor | None = None, + recv_bottom_buf: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + if halo_size <= 0: + return x, recv_top_buf, recv_bottom_buf + + rank, world_size = _rank_world(group) + if world_size <= 1: + return x, recv_top_buf, recv_bottom_buf + + dim = _spatial_dim(split_dim) + top_row_ref = _narrow_along_dim(x, dim, 0, halo_size) + bottom_row_ref = _narrow_along_dim(x, dim, x.shape[dim] - halo_size, halo_size) + recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row_ref) + recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row_ref) + + _halo_exchange_p2p( + rank=rank, + world_size=world_size, + group=group, + top_row_ref=top_row_ref, + bottom_row_ref=bottom_row_ref, + recv_top_buf=recv_top_buf, + recv_bottom_buf=recv_bottom_buf, + ) + + return torch.cat([recv_top_buf, x, recv_bottom_buf], dim=dim), recv_top_buf, recv_bottom_buf + + +class WanDistZeroPad2d(nn.Module): + """Apply ZeroPad2d only at global split-dimension boundaries.""" + + def __init__( + self, + padding: tuple[int, int, int, int], + group: dist.ProcessGroup, + *, + split_dim: str = "height", + split_padding: tuple[int, int] | None = None, + ) -> None: + super().__init__() + self.padding = padding + self.split_dim = split_dim + default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]) + self.split_padding = split_padding or default_split_padding + self.group = group + + def forward(self, x: torch.Tensor) -> torch.Tensor: + rank, world_size = _rank_world(self.group) + left, right, top, bottom = self.padding + if world_size > 1: + if self.split_dim == "height": + top = top if rank == 0 else 0 + bottom = bottom if rank == world_size - 1 else 0 + else: + left = left if rank == 0 else 0 + right = right if rank == world_size - 1 else 0 + return F.pad(x, (left, right, top, bottom)) + + +class WanDistConv2d(nn.Conv2d): + def __init__( + self, + source: nn.Conv2d, + group: dist.ProcessGroup, + split_dim: str = "height", + split_padding: tuple[int, int] | None = None, + ): + super().__init__( + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=0, + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self.load_state_dict(source.state_dict()) + self.group = group + self.split_dim = split_dim + self.split_tensor_dim = _spatial_dim(split_dim) + kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1] + self.halo_size = (kernel_extent - 1) // 2 + pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding + pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding + if split_dim == "height": + if split_padding is None: + split_padding = (pad_h, pad_h) + self._non_split_padding = (pad_w, pad_w, 0, 0) + self.kernel_extent = self.kernel_size[-2] + self.stride_extent = self.stride[-2] + else: + if split_padding is None: + split_padding = (pad_w, pad_w) + self._non_split_padding = (0, 0, pad_h, pad_h) + self.kernel_extent = self.kernel_size[-1] + self.stride_extent = self.stride[-1] + self.split_pad_left, self.split_pad_right = split_padding + self._halo_recv_top_buf: torch.Tensor | None = None + self._halo_recv_bottom_buf: torch.Tensor | None = None + self._trim_cache: dict[int, tuple[int, int, int]] = {} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, self._non_split_padding) + x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( + x, + group=self.group, + halo_size=self.halo_size, + split_dim=self.split_dim, + recv_top_buf=self._halo_recv_top_buf, + recv_bottom_buf=self._halo_recv_bottom_buf, + ) + shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim]) + if shift: + x_padded = _narrow_along_dim( + x_padded, + self.split_tensor_dim, + shift, + x_padded.shape[self.split_tensor_dim] - shift, + ) + out = super().forward(x_padded) + out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim) + return _zero_invalid_extent(out, split_dim=self.split_dim) + + def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]: + trim_params = self._trim_cache.get(local_extent) + if trim_params is None: + rank, world_size = _rank_world(self.group) + trim_params = _compute_conv_trim_params( + local_extent=local_extent, + rank=rank, + world_size=world_size, + halo_size=self.halo_size, + pad_before=self.split_pad_left, + pad_after=self.split_pad_right, + kernel_extent=self.kernel_extent, + stride_extent=self.stride_extent, + ) + self._trim_cache[local_extent] = trim_params + return trim_params + + +class WanDistCausalConv3d(nn.Conv3d): + def __init__( + self, + source: nn.Conv3d, + group: dist.ProcessGroup, + split_dim: str = "height", + ): + super().__init__( + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=0, + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self.load_state_dict(source.state_dict()) + self.group = group + self.split_dim = split_dim + self.split_tensor_dim = _spatial_dim(split_dim) + source_padding = getattr(source, "_padding", None) + if source_padding is None: + p_t, p_h, p_w = source.padding + source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0) + self._source_padding = tuple(source_padding) + if split_dim == "height": + self.split_pad_left = int(self._source_padding[2]) + self.split_pad_right = int(self._source_padding[3]) + self.kernel_extent = self.kernel_size[-2] + self.stride_extent = self.stride[-2] + padding = ( + self._source_padding[0], + self._source_padding[1], + 0, + 0, + self._source_padding[4], + self._source_padding[5], + ) + else: + self.split_pad_left = int(self._source_padding[0]) + self.split_pad_right = int(self._source_padding[1]) + self.kernel_extent = self.kernel_size[-1] + self.stride_extent = self.stride[-1] + padding = ( + 0, + 0, + self._source_padding[2], + self._source_padding[3], + self._source_padding[4], + self._source_padding[5], + ) + self.halo_size = (self.kernel_extent - 1) // 2 + self._padding = padding if self.halo_size > 0 else self._source_padding + self._halo_recv_top_buf: torch.Tensor | None = None + self._halo_recv_bottom_buf: torch.Tensor | None = None + self._trim_cache: dict[int, tuple[int, int, int]] = {} + + def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + padding = list(self._padding) + if cache_x is not None and padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + + x = F.pad(x, padding) + x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( + x, + group=self.group, + halo_size=self.halo_size, + split_dim=self.split_dim, + recv_top_buf=self._halo_recv_top_buf, + recv_bottom_buf=self._halo_recv_bottom_buf, + ) + shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim]) + if shift: + x_padded = _narrow_along_dim( + x_padded, + self.split_tensor_dim, + shift, + x_padded.shape[self.split_tensor_dim] - shift, + ) + out = super().forward(x_padded) + out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim) + return _zero_invalid_extent(out, split_dim=self.split_dim) + + def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]: + trim_params = self._trim_cache.get(local_extent) + if trim_params is None: + rank, world_size = _rank_world(self.group) + trim_params = _compute_conv_trim_params( + local_extent=local_extent, + rank=rank, + world_size=world_size, + halo_size=self.halo_size, + pad_before=self.split_pad_left, + pad_after=self.split_pad_right, + kernel_extent=self.kernel_extent, + stride_extent=self.stride_extent, + ) + self._trim_cache[local_extent] = trim_params + return trim_params + + +def _compute_conv_trim_params( + *, + local_extent: int, + rank: int, + world_size: int, + halo_size: int, + pad_before: int, + pad_after: int, + kernel_extent: int, + stride_extent: int, +) -> tuple[int, int, int]: + global_start = rank * local_extent + shift = 0 + if halo_size > 0 and stride_extent > 1: + shift = (global_start - halo_size + pad_before) % stride_extent + if shift: + global_start += shift + + global_extent = local_extent * world_size + min_i = math.ceil(((-pad_before) - (global_start - halo_size)) / stride_extent) + max_i = math.floor( + ((global_extent - 1 + pad_after) - (kernel_extent - 1) - (global_start - halo_size)) / stride_extent + ) + return shift, max(min_i, 0), max_i + 1 + + +def _trim_local_conv_output( + out: torch.Tensor, + halo_size: int, + start: int, + upper_bound: int, + *, + split_dim: str, +) -> torch.Tensor: + if halo_size <= 0: + return out + dim = _spatial_dim(split_dim) + end = min(upper_bound, out.shape[dim]) + if start != 0 or end != out.shape[dim]: + out = _narrow_along_dim(out, dim, start, end - start) + return out + + +def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None: + if getattr(module, "_vllm_omni_spatial_shard_attention", False): + return + orig_forward = module.forward + + def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + _, world_size = _rank_world(group) + if world_size <= 1: + return orig_forward(x, *args, **kwargs) + dim = _spatial_dim(split_dim) + local_extent = x.shape[dim] + gathered = all_gather_along_dim(x, group=group, dim=dim).contiguous() + full_extent = _current_full_extent(local_extent) + if full_extent is not None: + gathered = _narrow_along_dim(gathered, dim, 0, full_extent).contiguous() + out = orig_forward(gathered, *args, **kwargs) + return reshard_from_trimmed_extent(out, local_extent=local_extent, split_dim=split_dim, group=group) + + module.forward = MethodType(_forward, module) + module._vllm_omni_spatial_shard_attention = True # type: ignore[attr-defined] + + +def _replace_child( + parent: nn.Module, + name: str, + child: nn.Module, + group: dist.ProcessGroup, + split_dim: str, +) -> None: + if child.__class__.__name__ == "WanCausalConv3d": + setattr( + parent, + name, + WanDistCausalConv3d( + child, + group, + split_dim=split_dim, + ), + ) + return + if isinstance(child, nn.ZeroPad2d): + padding = tuple(int(p) for p in child.padding) + module_padding = padding + if parent.__class__.__name__ == "Sequential": + # Let the following WanDistConv2d account for global after-edge + # padding; this module handles the non-split dimension and before edge. + if split_dim == "height": + module_padding = (padding[0], padding[1], padding[2], 0) + else: + module_padding = (padding[0], 0, padding[2], padding[3]) + setattr( + parent, + name, + WanDistZeroPad2d( + module_padding, + group, + split_dim=split_dim, + split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]), + ), + ) + return + if isinstance(child, nn.Conv2d): + split_padding = None + if name == "1" and parent.__class__.__name__ == "Sequential": + # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a + # stride-2 conv with padding=0. Only the last rank should see the + # bottom/right global padding, which is approximated by split_padding. + prev = getattr(parent, "0", None) + if isinstance(prev, WanDistZeroPad2d): + split_padding = prev.split_padding + elif isinstance(prev, nn.ZeroPad2d): + pad = prev.padding + split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1])) + setattr( + parent, + name, + WanDistConv2d( + child, + group, + split_dim=split_dim, + split_padding=split_padding, + ), + ) + + +def _patch_decoder_modules( + module: nn.Module, + group: dist.ProcessGroup, + split_dim: str, + inside_attention: bool = False, +) -> None: + if module.__class__.__name__ == "WanAttentionBlock": + _patch_attention_block(module, group, split_dim) + inside_attention = True + + for name, child in list(module.named_children()): + if child.__class__.__name__ == "WanCausalConv3d": + _replace_child(module, name, child, group, split_dim) + continue + if isinstance(child, nn.Conv2d) and not inside_attention: + _replace_child(module, name, child, group, split_dim) + continue + _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention) + + +def _decoder_upsample_count(decoder: nn.Module) -> int: + count = 0 + for block in getattr(decoder, "up_blocks", []): + if getattr(block, "upsampler", None) is not None or getattr(block, "upsamplers", None) is not None: + count += 1 + return count + + +def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None: + """Patch ``vae.decoder`` once for spatially-sharded decode. + + This mutates the already-loaded decoder in place by swapping its spatial + convolutions/padding for halo-exchanging variants and wrapping + ``decoder.forward``. The patch is permanent for the lifetime of the VAE + instance and is applied only once (subsequent calls are no-ops). A given + instance is bound to a single ``split_dim``; switching between + ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here + otherwise. + + Only group-relative rank 0 assembles the final decoded frame, mirroring the + distributed tiled-decode ``broadcast_result=False`` contract; the other ranks + take part in the collectives but return an empty placeholder. + """ + _spatial_dim(split_dim) + if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False): + installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height") + if installed_split_dim != split_dim: + raise ValueError( + "Wan spatial-shard VAE decoder was already patched for " + f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split." + ) + return + decoder = getattr(vae, "decoder", None) + if decoder is None: + raise ValueError("Wan spatial-shard VAE decode requires a decoder module.") + + _patch_decoder_modules(decoder, group, split_dim) + upsample_count = _decoder_upsample_count(decoder) + orig_forward = decoder.forward + + def _forward( + self: nn.Module, + x: torch.Tensor, + feat_cache: list[torch.Tensor] | None = None, + feat_idx: list[int] | None = None, + first_chunk: bool = False, + ) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + tensor_dim = _spatial_dim(split_dim) + input_extent = x.shape[tensor_dim] + x, expected_extent = split_for_parallel_decode( + x, + upsample_count=upsample_count, + split_dim=split_dim, + group=group, + ) + rank, world_size = _rank_world(group) + token = _SPATIAL_SHARD_CONTEXT.set( + SpatialShardContext( + input_extent=input_extent, + local_input_extent=x.shape[tensor_dim], + split_dim=split_dim, + rank=rank, + world_size=world_size, + ) + ) + try: + out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + finally: + _SPATIAL_SHARD_CONTEXT.reset(token) + return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0) + + decoder.forward = MethodType(_forward, decoder) + vae._vllm_omni_wan_spatial_shard_installed = True + vae._vllm_omni_wan_spatial_shard_split_dim = split_dim + logger.info("Installed Wan VAE %s-sharded decode.", split_dim) + + +def spatial_shard_decode( + vae: Any, + z: torch.Tensor, + *, + group: dist.ProcessGroup, + return_dict: bool = True, + split_dim: str = "height", +) -> DecoderOutput | tuple[torch.Tensor]: + install_wan_spatial_shard_decode(vae, group, split_dim=split_dim) + + if z.shape[2] == 0: + raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.") + + # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with + # the halo/all-gather collectives; they just skip keeping/assembling the output. + rank, world_size = _rank_world(group) + produce_output = world_size <= 1 or rank == 0 + + vae.clear_cache() + try: + context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext() + with context: + x = vae.post_quant_conv(z) + decoded_chunks = [] + for i in range(z.shape[2]): + vae._conv_idx = [0] + chunk = vae.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=vae._feat_map, + feat_idx=vae._conv_idx, + first_chunk=(i == 0), + ) + if produce_output: + decoded_chunks.append(chunk) + + if produce_output: + out = torch.cat(decoded_chunks, dim=2) + if vae.config.patch_size is not None: + out = unpatchify(out, patch_size=vae.config.patch_size) + out = torch.clamp(out, min=-1.0, max=1.0) + else: + out = z.new_zeros(0) + finally: + vae.clear_cache() + + if not return_dict: + return (out,) + return DecoderOutput(sample=out) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 93d7f5d75..341e62ede 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -390,7 +390,7 @@ def initialize_model( model.vae.use_tiling = od_config.vae_use_tiling if is_distributed_vae: - model.vae.set_parallel_size(vae_pp_size) + model.vae.set_parallel_size(vae_pp_size, mode=od_config.parallel_config.vae_parallel_mode) # Apply sequence parallelism if enabled # This follows diffusers' pattern where enable_parallelism() is called diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d97fb26e5..acfc0a83f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -477,6 +477,7 @@ class OrchestratorArgs: diffusion_kv_cache_skip_layers: str | None = None cfg_parallel_size: int = 1 vae_patch_parallel_size: int = 1 + vae_parallel_mode: str = "tile" default_sampling_params: str | None = None max_generated_image_size: int | None = None tts_max_instructions_length: int | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 3cee3ccd8..5b8d2430f 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -953,6 +953,7 @@ class AsyncOmniEngine: cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 + vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile" enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False use_hsdp = normalized_kwargs.get("use_hsdp", False) hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) @@ -971,6 +972,7 @@ class AsyncOmniEngine: ulysses_mode=ulysses_mode, cfg_parallel_size=cfg_parallel_size, vae_patch_parallel_size=vae_patch_parallel_size, + vae_parallel_mode=vae_parallel_mode, use_hsdp=use_hsdp, hsdp_shard_size=hsdp_shard_size, hsdp_replicate_size=hsdp_replicate_size, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 97075c434..d4c0b462d 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -644,6 +644,18 @@ class OmniServeCommand(CLISubcommand): "Distributes VAE decode workload across multiple ranks by splitting the latent spatially. " "Equivalent to setting DiffusionParallelConfig.vae_patch_parallel_size.", ) + omni_config_group.add_argument( + "--vae-parallel-mode", + type=str, + default="tile", + choices=["tile", "spatial_shard_height", "spatial_shard_width"], + help="VAE parallel decode strategy for diffusion models. " + "'tile' (default) uses patch/tile parallel decode; " + "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits " + "decoder feature maps along height/width and exchanges halo regions. The " + "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. " + "Equivalent to setting DiffusionParallelConfig.vae_parallel_mode.", + ) # Default sampling parameters omni_config_group.add_argument(