anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature] Support HSDP for diffusion models (#1339)

Browse Source 
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Canlin GuoandHongsheng Liu authored and GitHub committed 2026-02-26 21:31:10 +08:00
1 parent 66c3e04263
commit f53152d340
17 files changed
+1023 -46
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion_acceleration.md
+30 -3
View File
@@ -30,6 +30,8 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
## Quick Comparison
### Cache Methods
@@ -69,9 +71,9 @@ The following table shows which models are currently supported by each accelerat
### VideoGen
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
### Quantization
@@ -265,6 +267,30 @@ outputs = omni.generate(
)
```
### Using HSDP
HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
use_hsdp=True,
hsdp_replicate_size=1, # Number of replica groups
hsdp_shard_size=8, # Number of GPUs to shard across (or -1 for auto)
),
)
outputs = omni.generate(
"A cat playing piano",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Using CFG-Parallel
Run image-to-image:
@@ -319,3 +345,4 @@ For detailed information on each acceleration method:
- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
- **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.