From d275c0c993e09bef5e6dfd560c67f6db69ff5f95 Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Tue, 3 Mar 2026 15:48:51 +0800 Subject: [PATCH] [FEATURE] Tile/Patch parallelism refactor for easily support other models (#1366) Signed-off-by: dengyunyang <584797741@qq.com> --- docs/design/feature/vae_parallel.md | 289 ++++++++++++++++++ .../diffusion/parallelism_acceleration.md | 10 +- docs/user_guide/diffusion_acceleration.md | 12 +- .../image_to_video/image_to_video.py | 19 +- .../text_to_video/text_to_video.py | 11 +- .../test_distributed_vae_executor.py | 165 ++++++++++ .../test_vae_decode_parallelism.py | 154 ++++++++++ .../distributed/autoencoders/__init__.py | 0 .../autoencoders/autoencoder_kl.py | 210 +++++++++++++ .../autoencoders/autoencoder_kl_qwenimage.py | 120 ++++++++ .../autoencoders/autoencoder_kl_wan.py | 142 +++++++++ .../autoencoders/distributed_vae_executor.py | 193 ++++++++++++ .../models/qwen_image/pipeline_qwen_image.py | 10 +- .../diffusion/models/sd3/pipeline_sd3.py | 8 +- .../models/wan2_2/pipeline_wan2_2.py | 4 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 4 +- .../models/z_image/pipeline_z_image.py | 8 +- vllm_omni/diffusion/registry.py | 37 +-- 18 files changed, 1336 insertions(+), 60 deletions(-) create mode 100644 docs/design/feature/vae_parallel.md create mode 100644 tests/diffusion/distributed/test_distributed_vae_executor.py create mode 100644 tests/e2e/offline_inference/test_vae_decode_parallelism.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/__init__.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md new file mode 100644 index 000000000..9009ece72 --- /dev/null +++ b/docs/design/feature/vae_parallel.md @@ -0,0 +1,289 @@ +# VAE Patch Parallelism + +This document describes how to add **VAE Patch Parallelism** support to a diffusion model. +We use **Qwen-Image** as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Testing](#testing) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is Vae Patch parallel? + +**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is: + ++ Split into multiple spatial tiles + ++ Distributed across multiple ranks + ++ Decoded in parallel + ++ Merged to reconstruct the final output + +This approach: + ++ Distributes computation across multiple devices + ++ Reduces peak memory usage per device + ++ Accelerates decoding latency + +### Architecture +We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding. + +The executor is model-agnostic and accepts three function parameters: + ++ split – Partition the latent into tiles + ++ exec – Decode a single tile + ++ merge – Combine decoded tiles into the final output + +#### Execution Flow + ++ Call split(z) to generate a list of TileTask and a GridSpec + ++ Dispatch tasks across ranks using workload-based balancing + ++ Each rank executes exec(task) on its assigned tiles + ++ Gather decoded tile results to rank 0 + ++ Rank 0 performs merge(...) + ++ (Optional) Broadcast final result to all ranks + +This design separates: + ++ Distributed execution logic + ++ Model-specific tiling and merging logic + +#### Why split / exec / merge is necessary? + +The latent tensor cannot be arbitrarily partitioned. + +During decoding: + ++ Each output pixel may depend on neighboring pixels + ++ The receptive field is model-dependent + +Therefore: + ++ Tiles must include overlap + ++ Merge must perform blending to avoid seams + +## Step-by-Step Implementation + +### Step 1: Implement DistributedAutoencoderKLQwenImage +`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version: + + +``` +class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model +``` +**Key points**: ++ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin ++ Call init_distributed() after loading weights + +### Step 2: Implement split/exec/merge +Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`: +``` +class GridSpec: + split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W)) + grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols) + tile_spec: dict = field(default_factory=dict) # Metadata required for merging + output_dtype: torch.dtype | None = None # Final output dtype +``` +``` +class TileTask: + tile_id: int # task id + grid_coord: tuple[int, ...] # Tile position in grid + tensor: torch.Tensor | list[torch.Tensor] # The tile tensor + workload: int | float = 1 # Used for load balancing (e.g., tile area) +``` +And tiled base split/exec/merge as follow: +``` +def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + # mostly copy from AutoencoderKL + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + tiletask_list = [] + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + time_list = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + time_list.append(tile) + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // tile_latent_stride_height, j // tile_latent_stride_width), + time_list, + workload=time_list[0].shape[3] * time_list[0].shape[4], + ) + ) + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + } + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + +def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + self.clear_cache() + time = [] + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx) + time.append(decoded) + result = torch.cat(time, dim=2) + return result + +def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + self.clear_cache() + + result_rows = [] + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] + ] + return dec +``` + +### Step 3: Override tiled_decode +We need to override tiled_decode, the main logic is: ++ check distributed is enabled ++ select split/exec/merge ++ Invoke self.distributed_decoder.execute to decode +``` +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) + + logger.info("Decode run with distributed executor") + result = self.distributed_decoder.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=True, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) +``` +`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0. + +### Step 4: Modify Pipeline +Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage +``` +class YourModelPipeline(nn.Module): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + ... +- self.vae = AutoencoderKL.from_pretrained( +- model, subfolder="vae", local_files_only=local_files_only).to(self.device) ++ self.vae = DistributedAutoencoderKL.from_pretrained( ++ model, subfolder="vae", local_files_only=local_files_only ++ ).to(self.device) +``` + +## Testing +Verify numerical consistency between: ++ vae_patch_parallel_size = 1 + ++ vae_patch_parallel_size = N + +Example: +torch.allclose(output_1, output_n, atol=1e-5) + +Testing requirements: ++ Fix random seed ++ Use identical tiling strategy + +```python +m = Omni( + model=model_name, + vae_use_tiling=True, + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + vae_patch_parallel_size=1, # or 2 + ), + ) +``` +When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead. + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Notes | +|-------|------|-------| +| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL | +| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan | +| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage | + +--- + +## Summary + +Adding Vae Patch Parallel support to diffusion model: + +1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge +2. **Change vae model in pipeline to Distributed Vae** +3. **Test** - Verify with `tensor_parallel_size=N` quality diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index b0225e2ff..4f4a209d0 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -27,7 +27,7 @@ The following table shows which models are currently supported by parallelism me | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | @@ -51,9 +51,9 @@ The following table shows which models are currently supported by parallelism me ### VideoGen -| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | -|-------|------------------|:----------:|:--------------:|:---------------:|:----:| -| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | +| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel | +|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:| +| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ | ### Tensor Parallelism @@ -85,7 +85,7 @@ outputs = omni.generate( VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP). !!! note "Enablement and feature gate" - - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`). + - VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information). - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.) #### Offline Inference diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 9fd14a2b0..2c36c966c 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -58,8 +58,8 @@ The following table shows which models are currently supported by each accelerat | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | -| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | @@ -67,13 +67,13 @@ The following table shows which models are currently supported by each accelerat | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | -| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | +| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ### VideoGen -| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | -|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:| -| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | +| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel | +|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:| +| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ### Quantization diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index a491ddfc3..4e42b8b29 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -20,6 +20,7 @@ Usage: import argparse import os +import time from pathlib import Path import numpy as np @@ -105,6 +106,12 @@ def parse_args() -> argparse.Namespace: choices=[1, 2], help="Number of GPUs used for classifier free guidance parallel size.", ) + parser.add_argument( + "--vae-patch-parallel-size", + type=int, + default=1, + help="Number of GPUs used for VAE patch/tile parallelism (decode).", + ) parser.add_argument( "--enforce-eager", action="store_true", @@ -173,6 +180,7 @@ def main(): ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, use_hsdp=args.use_hsdp, hsdp_shard_size=args.hsdp_shard_size, hsdp_replicate_size=args.hsdp_replicate_size, @@ -199,10 +207,14 @@ def main(): print(f" Model: {args.model}") print(f" Inference steps: {args.num_inference_steps}") print(f" Frames: {args.num_frames}") - print(f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}") + print( + f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}," + f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}" + ) print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") + generation_start = time.perf_counter() # omni.generate() returns Generator[OmniRequestOutput, None, None] frames = omni.generate( { @@ -220,6 +232,11 @@ def main(): num_frames=args.num_frames, ), ) + generation_end = time.perf_counter() + generation_time = generation_end - generation_start + + # Print profiling results + print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") # Extract video frames from OmniRequestOutput if isinstance(frames, list) and len(frames) > 0: diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 40fafa100..5db1b73dc 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -109,6 +109,12 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", ) + parser.add_argument( + "--vae-patch-parallel-size", + type=int, + default=1, + help="Number of GPUs used for VAE patch/tile parallelism (decode).", + ) return parser.parse_args() @@ -141,6 +147,7 @@ def main(): ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, ) # Check if profiling is requested via environment variable @@ -172,7 +179,9 @@ def main(): print(f" Inference steps: {args.num_inference_steps}") print(f" Frames: {args.num_frames}") print( - f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}" + f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}," + f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}," + f" vae_patch_parallel_size={args.vae_patch_parallel_size}" ) print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py new file mode 100644 index 000000000..b53014f39 --- /dev/null +++ b/tests/diffusion/distributed/test_distributed_vae_executor.py @@ -0,0 +1,165 @@ +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedVaeExecutor, + DistributedVaeMixin, + GridSpec, + TileTask, +) + + +class E2EOperator: + """tiles with (2, 3) -- (H,W)""" + + def split(self, z): + rows_num = 2 + cols_num = 3 + + h_size = z.shape[0] + w_size = z.shape[1] + + tasks = [] + for i in range(rows_num): + for j in range(cols_num): + tasks.append( + TileTask( + tile_id=len(tasks), + grid_coord=(i, j), + tensor=z[ + ((i * h_size) // rows_num) : (((i + 1) * h_size) // rows_num), + ((j * w_size) // cols_num) : (((j + 1) * w_size) // cols_num), + ], + ) + ) + + grid_spec = GridSpec( + split_dims=(0, 1), + grid_shape=(2, 3), + ) + return tasks, grid_spec + + def exec(self, task: TileTask): + return torch.full_like(task.tensor, fill_value=task.tile_id) + + def merge(self, coord_tensor_map, grid_spec): + tiles = [] + for r in range(grid_spec.grid_shape[0]): + row_tiles = [] + for c in range(grid_spec.grid_shape[1]): + coord = (r, c) + row_tiles.append(coord_tensor_map[coord]) + tiles.append(torch.cat(row_tiles, dim=1)) + return torch.cat(tiles, dim=0) + + +class DummyMixin(DistributedVaeMixin): + def __init__(self): + self.use_tiling = True + self.distributed_decoder = MagicMock() + self.distributed_decoder.parallel_size = 2 + self.distributed_decoder.group = None + + +@pytest.fixture(autouse=True) +def mock_dist(): + with ( + patch.object(dist, "get_world_size", return_value=2), + patch.object(dist, "get_rank", return_value=0), + patch.object(dist, "is_initialized", return_value=True), + patch.object(dist, "gather", return_value=None), + patch.object(dist, "broadcast", return_value=None), + ): + yield + + +@pytest.fixture(autouse=True) +def mock_dit_group(): + with patch( + "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group", + new=MagicMock(return_value=None), + ): + yield + + +@pytest.fixture(autouse=True) +def mock_dist_vae_executor(): + with ( + patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]), + patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x), + ): + yield + + +# ============================ +# Unitest +# ============================ + + +def test_balance_tasks(): + executor = DistributedVaeExecutor() + workloads = [2, 5, 13, 8, 2, 4] # each is 17 + tasks = [TileTask(0, (i,), torch.tensor([i]), workload=load) for i, load in enumerate(workloads)] + assigned = executor._balance_tasks(tasks, num_rank=2) + assert len(assigned) == 2 + total_work = [sum(t.workload for t in group) for group in assigned] + assert total_work[0] == total_work[1] + + +def test_compute_global_padding_shape(): + executor = DistributedVaeExecutor() + executor.rank = 0 + + local_results = [(0, torch.zeros((2, 3))), (1, torch.zeros((4, 2)))] + shape = executor._compute_global_padding_shape(local_results, 2, "cpu") + + assert shape == [2, 4, 3] + + +def test_pack_and_unpack(): + executor = DistributedVaeExecutor() + executor.world_size = 1 + + grid_spec = GridSpec(split_dims=(0, 1), grid_shape=(2, 2)) + + # ====================== + # pack + # ====================== + local_results = [(0, torch.tensor([[1, 2], [3, 4]]))] + + global_shape = [3, 3, 3] # (tiles, H, W) + + tile_tensor, meta_tensor = executor._pack_local_tiles( + local_results, global_shape, grid_spec, device="cpu", dtype=torch.int64 + ) + + # check pack + assert tile_tensor.shape == torch.Size(global_shape) + assert meta_tensor.shape == (global_shape[0], len(grid_spec.split_dims) + 1) + assert meta_tensor[0, 0] == 0 + assert meta_tensor[0, 1] == 2 + assert meta_tensor[0, 2] == 2 + + # ====================== + # unpack + # ====================== + meta_gather = [meta_tensor] + tile_gather = [tile_tensor] + + tid_coord_map = {0: (0, 0)} + + coord_tensor_map = executor._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map) + + # check unpack + assert torch.equal(coord_tensor_map[(0, 0)], torch.tensor([[1, 2], [3, 4]])) + + +def test_is_distributed_enabled(): + mixin = DummyMixin() + assert mixin.is_distributed_enabled() is True + + mixin.use_tiling = False + assert mixin.is_distributed_enabled() is False diff --git a/tests/e2e/offline_inference/test_vae_decode_parallelism.py b/tests/e2e/offline_inference/test_vae_decode_parallelism.py new file mode 100644 index 000000000..0bd3e9218 --- /dev/null +++ b/tests/e2e/offline_inference/test_vae_decode_parallelism.py @@ -0,0 +1,154 @@ +import os +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from vllm.distributed.parallel_state import cleanup_dist_env_and_memory + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import time + +from vllm_omni import Omni +from vllm_omni.platforms import current_omni_platform + +# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +all_model_configs = [ + { + "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + "out_width": 1280, + "out_height": 720, + "out_frames": 16, + }, + { + "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + "out_width": 1280, + "out_height": 720, + "out_frames": 16, + }, + { + "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "out_width": 1280, + "out_height": 720, + "out_frames": 16, + }, + { + "model_name": "Qwen/Qwen-Image", + "out_width": 3072, + "out_height": 3072, + "out_frames": 1, + }, + # { + # "model_name": "stepfun-ai/NextStep-1.1", + # "out_width": 3072, + # "out_height": 3072, + # "out_frames": 1, + # }, +] + + +def is_nextstep_model(model_name: str) -> bool: + """Check if the model is a NextStep model by reading its config.""" + from vllm.transformers_utils.config import get_hf_file_to_dict + + try: + cfg = get_hf_file_to_dict("config.json", model_name) + if cfg and cfg.get("model_type") == "nextstep": + return True + except Exception: + pass + return False + + +def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1): + try: + parallel_config = DiffusionParallelConfig( + tensor_parallel_size=tp, + vae_patch_parallel_size=vae_patch_parallel_size, + ) + + omni_kwargs = { + "model": model_configs["model_name"], + "vae_use_tiling": using_tile, + "parallel_config": parallel_config, + } + use_nextstep = is_nextstep_model(model_configs["model_name"]) + if use_nextstep: + # NextStep-1.1 requires explicit pipeline class + omni_kwargs["model_class_name"] = "NextStep11Pipeline" + m = Omni(**omni_kwargs) + image = Image.new("RGB", (out_width, out_height), (0, 0, 0)) + start = time.perf_counter() + outputs = m.generate( + { + "prompt": "A cat sitting on a table", + "multi_modal_data": {"image": image}, + }, + sampling_params_list=OmniDiffusionSamplingParams( + height=model_configs["out_height"], + width=model_configs["out_width"], + num_frames=model_configs["out_frames"], + num_inference_steps=2, + generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), + ), + ) + end = time.perf_counter() + first_output = outputs[0] + req_out = first_output.request_output[0] + frames = req_out.images[0] + if isinstance(frames, torch.Tensor): + frames = frames.detach().cpu().numpy() + elif isinstance(frames, Image.Image): + frames = np.array(frames) + # frames shape: (batch, num_frames, height, width, channels) + cost = (end - start) * 1000 + return frames, cost + finally: + m.close() + cleanup_dist_env_and_memory() + + +@pytest.mark.parametrize("model_configs", all_model_configs) +def test_vae_parallel_model(model_configs: str): + parallel_result, parallel_time = model_run( + model_configs=model_configs, + tp=2, + out_width=1280, + out_height=704, + out_frames=5, + using_tile=True, + vae_patch_parallel_size=2, + ) + non_parallel_result, non_parallel_time = model_run( + model_configs=model_configs, + tp=2, + out_width=1280, + out_height=704, + out_frames=5, + using_tile=True, + vae_patch_parallel_size=1, + ) + result_diff = np.abs(non_parallel_result - parallel_result) + + mean_threshold = 3e-2 + max_threshold = 3e-2 # they should be totally same + print( + f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): " + f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " + f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms" + ) + # assert non_parallel_time > parallel_time + assert result_diff.mean() < mean_threshold + assert result_diff.max() < max_threshold diff --git a/vllm_omni/diffusion/distributed/autoencoders/__init__.py b/vllm_omni/diffusion/distributed/autoencoders/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py new file mode 100644 index 000000000..7df2d6a8a --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from typing import Any + +import torch +from diffusers.models.autoencoders import AutoencoderKL as Diffusers_AutoencoderKL +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +# from vllm_omni.diffusion.models.nextstep_1_1.modeling_flux_vae import AutoencoderKL as Next_Step_AutoencoderKL + +logger = init_logger(__name__) + + +# We use base class because some model re-implement AutoencoderKL, but split is share. +class DistributedAutoencoderKL_base(DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + # mostly copy from AutoencoderKL + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping 64x64 tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + tiletask_list = [] + for i in range(0, z.shape[2], overlap_size): + for j in range(0, z.shape[3], overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // overlap_size, j // overlap_size), + tile, + workload=tile.shape[2] * tile.shape[3], + ) + ) + + tile_spec = { + "blend_extent": blend_extent, + "row_limit": row_limit, + } + grid_spec = GridSpec( + split_dims=(2, 3), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + tile = task.tensor + if self.config.use_post_quant_conv: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + return decoded + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_extent"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_extent"]) + result_row.append(tile[:, :, : grid_spec.tile_spec["row_limit"], : grid_spec.tile_spec["row_limit"]]) + result_rows.append(torch.cat(result_row, dim=3)) + + dec = torch.cat(result_rows, dim=2) + return dec + + def patch_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + overlap_latent = int(self.tile_latent_min_size * float(self.tile_overlap_factor)) + halo_base = max(0, overlap_latent // 2) + + _, _, latent_h, latent_w = z.shape + scale = int(2 ** (len(self.config.block_out_channels) - 1)) + max_parallel_size = self.distributed_decoder.parallel_size + + root = int(math.sqrt(max_parallel_size)) + for rows in range(root, 0, -1): + if max_parallel_size % rows == 0: + grid_rows, grid_cols = rows, max_parallel_size // rows + break + tiletask_list = [] + halo_size = dict() + for i in range(grid_rows): + for j in range(grid_cols): + h0 = (i * latent_h) // grid_rows + h1 = ((i + 1) * latent_h) // grid_rows + w0 = (j * latent_w) // grid_cols + w1 = ((j + 1) * latent_w) // grid_cols + + core_h = max(0, h1 - h0) + core_w = max(0, w1 - w0) + halo = max(halo_base, min(core_h, core_w) // 2) + ph0 = max(0, h0 - halo) + ph1 = min(latent_h, h1 + halo) + pw0 = max(0, w0 - halo) + pw1 = min(latent_w, w1 + halo) + tile = z[:, :, ph0:ph1, pw0:pw1] + tiletask_list.append(TileTask(len(tiletask_list), (i, j), tile, tile.shape[2] * tile.shape[3])) + halo_size[(i, j)] = { + "up": h0 - ph0, + "down": ph1 - h1, + "left": w0 - pw0, + "right": pw1 - w1, + } + + tile_spec = { + "halo_size": halo_size, + "scale": scale, + } + grid_spec = GridSpec( + split_dims=(2, 3), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def patch_exec(self, task: TileTask) -> torch.Tensor: + return self.tile_exec(task) + + def patch_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + halo = grid_spec.tile_spec["halo_size"][(i, j)] + scale = grid_spec.tile_spec["scale"] + + tile = coord_tensor_map[(i, j)] + halo_up = halo["up"] * scale + halo_down = halo["down"] * scale + halo_left = halo["left"] * scale + halo_right = halo["right"] * scale + + core_tile = tile[ + :, + :, + halo_up : (None if halo_down == 0 else -halo_down), + halo_left : (None if halo_right == 0 else -halo_right), + ] + result_row.append(core_tile) + result_rows.append(torch.cat(result_row, dim=3)) + dec = torch.cat(result_rows, dim=2) + return dec + + def _strategy_select(self, z: torch.Tensor): + tile_latent_min_size = getattr(self, "tile_latent_min_size", None) + tile_overlap_factor = getattr(self, "tile_overlap_factor", None) + if tile_latent_min_size is None or tile_overlap_factor is None: + return None, None, None + if z.shape[-1] > tile_latent_min_size or z.shape[-2] > tile_latent_min_size: + return self.tile_split, self.tile_exec, self.tile_merge + + return self.patch_split, self.patch_exec, self.patch_merge + + # Normally, we should override tiled_decode. However, since we also need to + # support the patch split strategy, we override decode instead. + def decode(self, z: torch.Tensor, return_dict: bool = True, *args: Any, **kwargs: Any): + if not self.is_distributed_enabled(): + return super().decode(z, return_dict=return_dict, *args, **kwargs) + + split, exec, merge = self._strategy_select(z) + + if split is not None: + strategy = "tile" if split == self.tile_split else "patch" + logger.info(f"Decode run with distributed executor, split strategy is {strategy}") + result = self.distributed_decoder.execute( + z, DistributedOperator(split=split, exec=exec, merge=merge), broadcast_result=False + ) + if not return_dict: + return (result,) + + from diffusers.models.autoencoders.vae import DecoderOutput + + return DecoderOutput(sample=result) + else: + return super().decode(z, return_dict=return_dict, *args, **kwargs) + + +class DistributedAutoencoderKL(DistributedAutoencoderKL_base, Diffusers_AutoencoderKL): + pass + + +# Next_Step_AutoencoderKL not support tiling now, so we currently disable it. + +# class DistributedAutoencoderKL_NextStep(DistributedAutoencoderKL_base, Next_Step_AutoencoderKL): +# pass diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py new file mode 100644 index 000000000..7549bbd3d --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +from diffusers.models.autoencoders import AutoencoderKLQwenImage +from diffusers.models.autoencoders.vae import DecoderOutput +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +logger = init_logger(__name__) + + +class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + # mostly copy from AutoencoderKL + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + tiletask_list = [] + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + time_list = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + time_list.append(tile) + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // tile_latent_stride_height, j // tile_latent_stride_width), + time_list, + workload=time_list[0].shape[3] * time_list[0].shape[4], + ) + ) + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + } + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + self.clear_cache() + time = [] + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx) + time.append(decoded) + result = torch.cat(time, dim=2) + return result + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + grid_h, grid_w = grid_spec.grid_shape + self.clear_cache() + + result_rows = [] + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] + ] + return dec + + 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) + + logger.info("Decode run with distributed executor") + result = self.distributed_decoder.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=True, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py new file mode 100644 index 000000000..7defbae79 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +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.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +logger = init_logger(__name__) + + +class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_sample_stride_height = self.tile_sample_stride_height + tile_sample_stride_width = self.tile_sample_stride_width + if self.config.patch_size is not None: + sample_height = sample_height // self.config.patch_size + sample_width = sample_width // self.config.patch_size + tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size + tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size + blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height + blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width + else: + blend_height = self.tile_sample_min_height - tile_sample_stride_height + blend_width = self.tile_sample_min_width - tile_sample_stride_width + + tiletask_list = [] + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + time_list = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + time_list.append(tile) + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // tile_latent_stride_height, j // tile_latent_stride_width), + time_list, + workload=time_list[0].shape[3] * time_list[0].shape[4], + ) + ) + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + "tile_sample_stride_height": tile_sample_stride_height, + "tile_sample_stride_width": tile_sample_stride_width, + } + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + self.clear_cache() + time = [] + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0)) + time.append(decoded) + result = torch.cat(time, dim=2) + return result + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + self.clear_cache() + for i in range(grid_h): + result_row = [] + for j in range(grid_w): + tile = coord_tensor_map[(i, j)] + if i > 0: + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append( + tile[ + :, + :, + :, + : grid_spec.tile_spec["tile_sample_stride_height"], + : grid_spec.tile_spec["tile_sample_stride_width"], + ] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] + ] + + if self.config.patch_size is not None: + dec = unpatchify(dec, patch_size=self.config.patch_size) + + dec = torch.clamp(dec, min=-1.0, max=1.0) + return dec + + 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) + + logger.info("Decode run with distributed executor") + result = self.distributed_decoder.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=False, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py new file mode 100644 index 000000000..bdf664741 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +import torch.distributed as dist +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.parallel_state import get_dit_group + +logger = init_logger(__name__) + + +@dataclass +class GridSpec: + """The Grid shape split""" + + split_dims: tuple[int, ...] # For example: (3,4) for (B, C, T, H, W), (2,3) for (B, C, H, W) + grid_shape: tuple[int, ...] # For example: (nh, nw) for (B, C, T, H, W), (nh, nw) for (B, C, H, W) + tile_spec: dict = field(default_factory=dict) + output_dtype: torch.dtype | None = None + + +@dataclass +class TileTask: + tile_id: int + grid_coord: tuple[int, ...] # The coordinate of the tile in GridSpec.grid_shape + tensor: torch.Tensor | list[torch.Tensor] # The tile tensor + workload: int | float = 1 + + +@dataclass +class DistributedOperator: + split: callable + exec: callable + merge: callable + + +class DistributedVaeExecutor: + """ + Abstract util class for distributed patch/tile parallel VAE decoding. + """ + + def __init__(self): + self.group = get_dit_group() + self.world_size = dist.get_world_size(self.group) + self.rank = dist.get_rank(self.group) + self.parallel_size = 1 + + def set_parallel_size(self, parallel_size: int): + self.parallel_size = parallel_size + + def gather_tensors(self, tensor: torch.Tensor): + gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] if self.rank == 0 else None + dist.gather(tensor, gather_list=gather_list, dst=0, group=self.group) + return gather_list + + def broadcast_tensor(self, tensor: torch.Tensor): + dist.broadcast(tensor, src=0, group=self.group) + return tensor + + def _compute_global_padding_shape(self, local_results, output_ndim: int, device): + local_tile_max_dims = [0] * output_ndim + for _, tile_tensor in local_results: + for dim_idx, dim_size in enumerate(tile_tensor.shape): + local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size) + local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device) + dist.all_reduce( + local_shape_stat, + op=dist.ReduceOp.MAX, + group=self.group, + ) + return local_shape_stat.tolist() + + def _pack_local_tiles(self, local_results, global_padding_shape, grid_spec, device, dtype): + tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype) + meta_tensor = torch.full( + (global_padding_shape[0], len(grid_spec.split_dims) + 1), -1, device=device, dtype=torch.int64 + ) + for idx, (tid, t_tensor) in enumerate(local_results): + meta_tensor[idx, 0] = tid + for i, dim in enumerate(grid_spec.split_dims): + meta_tensor[idx, i + 1] = t_tensor.shape[dim] + slices = tuple(slice(0, s) for s in t_tensor.shape) + tile_tensor[idx][slices] = t_tensor + return tile_tensor, meta_tensor + + def _unpack_tiles(self, meta_gather, tile_gather, grid_spec, tid_coord_map): + coord_tensor_map = {} + for r in range(self.world_size): + meta_src = meta_gather[r] + tiles_src = tile_gather[r] + for idx in range(meta_src.shape[0]): + tid = int(meta_src[idx, 0]) + if tid < 0: + continue + slices = [slice(None)] * tiles_src[idx].ndim + for i, dim in enumerate(grid_spec.split_dims): + slices[dim] = slice(0, int(meta_src[idx, i + 1])) + slices = tuple(slices) + coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices] + + return coord_tensor_map + + def _balance_tasks(self, task_list, num_rank): + workloads = [0] * num_rank + assigned = [[] for _ in range(num_rank)] + + for task in sorted(task_list, key=lambda t: t.workload, reverse=True): + r = workloads.index(min(workloads)) + assigned[r].append(task) + workloads[r] += task.workload + + return assigned + + def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True): + pp_size = min(self.parallel_size, self.world_size) + + # 1. Split into tiles + tiletask_list, grid_spec = operator.split(z) + tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list} + + # 2. local decode + assigned = self._balance_tasks(tiletask_list, pp_size) + local_tasks = assigned[self.rank] if pp_size <= self.world_size else [] + local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks] + + # 3. compute shape per rank + global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device) + + # 5. prepare tile tensors + output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype + local_tile_tensor, local_meta_tensor = self._pack_local_tiles( + local_results, global_padding_shape, grid_spec, z.device, output_dtype + ) + + # 6. gather tiles & meta + meta_gather = self.gather_tensors(local_meta_tensor) + tile_gather = self.gather_tensors(local_tile_tensor) + + if self.rank != 0: + result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks + else: + # 7. reconstruct full tensor (rank 0) + coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map) + result = operator.merge(coord_tensor_map, grid_spec) + + if broadcast_result: + result = self._sync_final_result(result, z.ndim, z.device, output_dtype) + return result + + def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype): + shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64) + if self.rank == 0: + shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64)) + dist.broadcast(shape_tensor, src=0, group=self.group) + + if self.rank != 0: + sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype) + else: + sync_result = rank0_result + dist.broadcast(sync_result, src=0, group=self.group) + return sync_result + + +class DistributedVaeMixin: + def init_distributed(self): + self.distributed_decoder = DistributedVaeExecutor() + + def set_parallel_size(self, parallel_size: int) -> bool: + return self.distributed_decoder.set_parallel_size(parallel_size) + + def is_distributed_enabled(self) -> bool: + if ( + self.distributed_decoder.parallel_size <= 1 + or not dist.is_initialized() + or not getattr(self, "use_tiling", False) + ): + return False + world_size = dist.get_world_size(group=self.distributed_decoder.group) + pp_size = min(int(self.distributed_decoder.parallel_size), int(world_size)) + if pp_size <= 1: + return False + if self.distributed_decoder.parallel_size > pp_size: + logger.warning( + f"vae_patch_parallel_size={self.distributed_decoder.parallel_size} " + f"is greater than dit_group={world_size};" + f" using dit_group size={world_size}" + ) + return True diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index a61044304..02758ee19 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -13,9 +13,6 @@ import numpy as np import torch import torch.distributed from diffusers.image_processor import VaeImageProcessor -from diffusers.models.autoencoders.autoencoder_kl_qwenimage import ( - AutoencoderKLQwenImage, -) from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) @@ -25,6 +22,7 @@ from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( @@ -269,9 +267,9 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin): self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ) - self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( - self.device - ) + self.vae = DistributedAutoencoderKLQwenImage.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only + ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config) self.transformer = QwenImageTransformer2DModel( diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py index 3668c132f..9c9722c36 100644 --- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py +++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py @@ -6,7 +6,6 @@ from collections.abc import Iterable import torch from diffusers.image_processor import VaeImageProcessor -from diffusers.models.autoencoders import AutoencoderKL from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) @@ -16,6 +15,7 @@ from transformers import CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderMo from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -174,9 +174,9 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin): ) self.transformer = SD3Transformer2DModel(od_config=od_config) - self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( - self.device - ) + self.vae = DistributedAutoencoderKL.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only + ).to(self.device) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 23eebc7fd..63de7395f 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -11,13 +11,13 @@ from typing import Any, cast import PIL.Image import torch -from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -267,7 +267,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin): self.text_encoder = UMT5EncoderModel.from_pretrained( model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) - self.vae = AutoencoderKLWan.from_pretrained( + self.vae = DistributedAutoencoderKLWan.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 4b32d5a5f..2d750c762 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -11,13 +11,13 @@ from typing import Any, cast import numpy as np import PIL.Image import torch -from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -210,7 +210,7 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin): self.image_encoder = None # VAE - self.vae = AutoencoderKLWan.from_pretrained( + self.vae = DistributedAutoencoderKLWan.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device) diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 2c6ef4c86..b16d24bf7 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -24,7 +24,6 @@ from typing import Any import torch import torch.nn as nn from diffusers.image_processor import VaeImageProcessor -from diffusers.models.autoencoders import AutoencoderKL from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils import logging from diffusers.utils.torch_utils import randn_tensor @@ -32,6 +31,7 @@ from transformers import AutoModel, AutoTokenizer from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.z_image.z_image_transformer import ( @@ -171,9 +171,9 @@ class ZImagePipeline(nn.Module): self.text_encoder = AutoModel.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ) - self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( - self._execution_device - ) + self.vae = DistributedAutoencoderKL.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only + ).to(self._execution_device) # Get vLLM quantization config for linear layers quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config) self.transformer = ZImageTransformer2DModel(quant_config=quant_config) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index ed99b1473..c806dbae9 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -8,6 +8,7 @@ from vllm.logger import init_logger from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel @@ -123,13 +124,6 @@ DiffusionModelRegistry = _ModelRegistry( } ) -_VAE_PATCH_PARALLEL_ALLOWLIST = { - # Only enable for models we have validated end-to-end. - "StableDiffusion3Pipeline", - "ZImagePipeline", - "NextStep11Pipeline", -} - _NO_CACHE_ACCELERATION = { # Pipelines that do not support cache acceleration (cache_dit / tea_cache). "NextStep11Pipeline", @@ -161,17 +155,14 @@ def initialize_model( model = model_class(od_config=od_config) vae_pp_size = od_config.parallel_config.vae_patch_parallel_size - if vae_pp_size > 1 and od_config.model_class_name not in _VAE_PATCH_PARALLEL_ALLOWLIST: + is_distributed_vae = hasattr(model, "vae") and isinstance(model.vae, DistributedVaeMixin) + if vae_pp_size > 1 and not is_distributed_vae: logger.warning( - "vae_patch_parallel_size=%d is set but VAE patch parallelism is only enabled for %s; ignoring.", + "vae_patch_parallel_size=%d is set but VAE patch parallelism is NOT enabled for %s; ignoring.", vae_pp_size, - sorted(_VAE_PATCH_PARALLEL_ALLOWLIST), + od_config.model_class_name, ) - if ( - vae_pp_size > 1 - and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST - and not od_config.vae_use_tiling - ): + if vae_pp_size > 1 and is_distributed_vae and not od_config.vae_use_tiling: logger.info( "vae_patch_parallel_size=%d requires vae_use_tiling; automatically enabling it.", vae_pp_size, @@ -184,20 +175,8 @@ def initialize_model( if hasattr(model.vae, "use_tiling"): model.vae.use_tiling = od_config.vae_use_tiling - if ( - vae_pp_size > 1 - and hasattr(model, "vae") - and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST - and od_config.vae_use_tiling - ): - from vllm_omni.diffusion.distributed.parallel_state import get_dit_group - from vllm_omni.diffusion.distributed.vae_patch_parallel import maybe_wrap_vae_decode_with_patch_parallelism - - maybe_wrap_vae_decode_with_patch_parallelism( - model, - vae_patch_parallel_size=vae_pp_size, - group_getter=get_dit_group, - ) + if is_distributed_vae: + model.vae.set_parallel_size(vae_pp_size) # Apply sequence parallelism if enabled # This follows diffusers' pattern where enable_parallelism() is called