[Feature] Add support for Pipeline Parallel and integrate it into Wan 2.2 (#2322)
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
This commit is contained in:
Rustam Khadipash authored and GitHub committed
1 parent
475a4002b0
commit
89f8819525
20 files changed
+1981
-1009
No files matched your search
@@ -4,13 +4,14 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
|
||||
|
||||
## Supported Methods
|
||||
|
||||
| Method | Description |
|
||||
|--------|-------------|
|
||||
| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
|
||||
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
|
||||
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
|
||||
| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
|
||||
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
|
||||
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
|
||||
| Method | Description |
|
||||
|----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------|
|
||||
| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
|
||||
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
|
||||
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
|
||||
| **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory |
|
||||
| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
|
||||
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
|
||||
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
|
||||
|
||||
See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.
|
||||
@@ -0,0 +1,205 @@
|
||||
# Pipeline Parallelism Guide
|
||||
|
||||
## Table of Content
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Example Script](#example-script)
|
||||
- [Configuration Parameters](#configuration-parameters)
|
||||
- [Best Practices](#best-practices)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Summary](#summary)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
Pipeline Parallelism splits the denoising transformer block-wise into sequential stages across GPUs. Each rank owns only
|
||||
part of the transformer, which reduces per-GPU model memory and enables larger diffusion models to run across multiple
|
||||
devices.
|
||||
|
||||
It can also be combined with other distributed methods such as CFG-Parallel, Tensor Parallelism, and Sequence
|
||||
Parallelism.
|
||||
|
||||
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
|
||||
|
||||
---
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Basic Usage
|
||||
|
||||
Simplest working example:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
pipeline_parallel_size=2,
|
||||
),
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
{"prompt": "A cinematic drone shot over snowy mountains"},
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=40,
|
||||
num_frames=81,
|
||||
height=704,
|
||||
width=1280,
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Example Script
|
||||
|
||||
### Offline Inference
|
||||
|
||||
Use python scripts under:
|
||||
|
||||
- `examples/offline_inference/text_to_video/text_to_video.py`
|
||||
- `examples/offline_inference/image_to_video/image_to_video.py`
|
||||
|
||||
Text-to-video example:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--width=1280 \
|
||||
--height=704 \
|
||||
--guidance-scale=5.0 \
|
||||
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
|
||||
--output=t2v_5B_pp2.mp4 \
|
||||
--pipeline-parallel-size=2
|
||||
```
|
||||
|
||||
Pipeline Parallelism can also be combined with CFG-Parallel:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--width=1280 \
|
||||
--height=704 \
|
||||
--guidance-scale=5.0 \
|
||||
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
|
||||
--output=t2v_5B_pp2_cfg2.mp4 \
|
||||
--pipeline-parallel-size=2 \
|
||||
--cfg-parallel-size=2
|
||||
```
|
||||
|
||||
### Online Serving
|
||||
|
||||
Enable Pipeline Parallelism in online serving:
|
||||
|
||||
```bash
|
||||
# Default PP configuration
|
||||
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 --pipeline-parallel-size 2
|
||||
|
||||
# PP + CFG-Parallel
|
||||
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 \
|
||||
--pipeline-parallel-size 2 \
|
||||
--cfg-parallel-size 2
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Configuration Parameters
|
||||
|
||||
In `DiffusionParallelConfig`
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|--------------------------|------|---------|------------------------------------------------------------------------------------------------------------------|
|
||||
| `pipeline_parallel_size` | int | 1 | Number of pipeline-parallel stages. Set to a value greater than 1 to split the denoising transformer across GPUs |
|
||||
|
||||
> [!NOTE]
|
||||
> Total GPU count is the product of all enabled distributed dimensions, for example
|
||||
`pipeline_parallel_size * cfg_parallel_size * tensor_parallel_size * ulysses_degree * ring_degree`.
|
||||
|
||||
### Manual Layer Partitioning
|
||||
|
||||
By default, transformer layers are distributed evenly across PP ranks. You can override this with the
|
||||
`VLLM_PP_LAYER_PARTITION` environment variable to assign a specific number of layers to each rank:
|
||||
|
||||
```bash
|
||||
# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
|
||||
export VLLM_PP_LAYER_PARTITION=8,12,12,8
|
||||
```
|
||||
|
||||
The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
|
||||
the total number of transformer layers. This is useful when you want to balance memory or compute asymmetrically across
|
||||
ranks.
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Best Practices
|
||||
|
||||
### When to Use
|
||||
|
||||
**Good for:**
|
||||
|
||||
- Large diffusion transformers that do not fit comfortably on one GPU
|
||||
- Multi-GPU setups where reducing per-GPU model memory is more important than minimizing communication
|
||||
- Combining with CFG-Parallel or other distributed methods on supported models
|
||||
|
||||
**Not for:**
|
||||
|
||||
- Single GPU setups
|
||||
- Models that do not support Pipeline Parallelism
|
||||
(check [supported models](../../diffusion_features.md#supported-models))
|
||||
- Very small models where inter-stage communication overhead may outweigh the benefit
|
||||
|
||||
### Expected Behavior
|
||||
|
||||
Pipeline Parallelism primarily reduces per-GPU memory usage by splitting the transformer into stage-local blocks.
|
||||
Depending on the model, topology, and resolution, it may also help execution fit into available hardware, but it is not
|
||||
primarily a latency optimization.
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issue 1: No benefit from Pipeline Parallelism
|
||||
|
||||
**Symptoms**: PP is enabled, but latency does not improve or becomes slightly worse.
|
||||
|
||||
**Solutions**:
|
||||
|
||||
1. **Check your goal:**
|
||||
|
||||
```python
|
||||
# PP is mainly for memory scaling, not guaranteed latency speedup
|
||||
parallel_config = DiffusionParallelConfig(pipeline_parallel_size=2)
|
||||
```
|
||||
|
||||
2. **Check model support:**
|
||||
- Verify your model in [supported models](../../diffusion_features.md#supported-models)
|
||||
- PP is currently validated only on selected pipelines
|
||||
|
||||
3. **Combine with other methods when appropriate:**
|
||||
- PP can be combined with CFG-Parallel, Tensor Parallelism, or Sequence Parallelism on supported models
|
||||
|
||||
### Common Issue 2: PP pipeline fails at import
|
||||
|
||||
**Symptoms**: Importing a custom pipeline raises a `TypeError` about `CFGParallelMixin` or mixin order.
|
||||
|
||||
**Solutions**:
|
||||
|
||||
1. Inherit both mixins.
|
||||
2. Put `PipelineParallelMixin` before `CFGParallelMixin` in the class base list.
|
||||
3. Use `class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): ...` as the reference pattern.
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
1. ✅ **Enable Pipeline Parallelism** - Set `pipeline_parallel_size > 1` in `DiffusionParallelConfig`
|
||||
2. ✅ **Use Supported Models** - Verify your model supports PP in
|
||||
[supported models](../../diffusion_features.md#supported-models)
|
||||
3. ✅ **Combine When Needed** - PP can be combined with CFG-Parallel and other distributed methods on supported pipelines
|
||||
4. ✅ **Scale for Memory** - Use PP primarily to reduce per-GPU model memory and fit larger transformers
|
||||
Reference in new issue
Block a user