289 lines
9.4 KiB
Markdown
289 lines
9.4 KiB
Markdown
|
# 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
|