[FEATURE] Tile/Patch parallelism refactor for easily support other models (#1366)
Signed-off-by: dengyunyang <584797741@qq.com>
This commit is contained in:
dengyunyang authored and GitHub committed
1 parent
e7f2c88701
commit
d275c0c993
18 files changed
+1336
-60
No files matched your search
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
Whitespace-only changes.
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user