[skip ci][Doc] Refine the Diffusion Features User Guide (#1928)
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
This commit is contained in:
1 parent
c9370f493e
commit
911be0d426
20 files changed
+1912
-1244
No files matched your search
@@ -0,0 +1,151 @@
|
||||
# Tensor 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
|
||||
|
||||
Tensor Parallelism (TP) shards some model weights across multiple GPUs, usually the Linear layers. This enables running large models that don't fit on a single GPU. It's essential for memory-constrained setups or very large models.
|
||||
|
||||
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
|
||||
|
||||
!!! note "TP Limitations for Diffusion Models"
|
||||
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
|
||||
|
||||
- Good news: The text_encoder typically has minimal impact on overall inference performance.
|
||||
- Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
|
||||
|
||||
We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
|
||||
|
||||
---
|
||||
|
||||
## Quick Start
|
||||
|
||||
|
||||
### Basic Usage
|
||||
|
||||
Simplest working example:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
omni = Omni(
|
||||
model="Tongyi-MAI/Z-Image-Turbo",
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2), # Enable TP
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"a cat reading a book",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=9),
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Example Script
|
||||
|
||||
### Offline Inference
|
||||
|
||||
Use Python script under `examples/offline_inference`, and enable TP:
|
||||
|
||||
```bash
|
||||
# Text-to-Image with Qwen-Image
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--tensor-parallel-size 2
|
||||
|
||||
# Image Editing with Qwen-Image-Edit
|
||||
python examples/offline_inference/image_to_image/image_edit.py \
|
||||
--model Qwen/Qwen-Image-Edit \
|
||||
--image input.png \
|
||||
--prompt "Edit description" \
|
||||
--tensor-parallel-size 2
|
||||
```
|
||||
|
||||
### Online Serving
|
||||
|
||||
You can enable tensor parallelism in online serving via `--tensor-parallel-size`:
|
||||
|
||||
```bash
|
||||
# Text-to-Image with Qwen-Image on 2 GPUs
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 \
|
||||
--tensor-parallel-size 2
|
||||
|
||||
# Text-to-Image with Z-Image (TP=2 only)
|
||||
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
|
||||
--tensor-parallel-size 2
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Configuration Parameters
|
||||
|
||||
In `DiffusionParallelConfig`:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `tensor_parallel_size` | int | 1 | Number of GPUs to shard model weights across. Must divide number of heads. |
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Best Practices
|
||||
|
||||
### When to Use
|
||||
|
||||
**Good for:**
|
||||
|
||||
- Large models that don't fit on a single GPU, especially for models with large DiT blocks (transformer layers)
|
||||
- Memory-constrained environments
|
||||
|
||||
**Not for:**
|
||||
|
||||
- When maximum throughput is needed and memory is sufficient
|
||||
- Models with incompatible dimensions (e.g., Z-Image `num_heads=30`, which now supports `tensor_parallel_size=2`)
|
||||
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issue 1: Out of Memory (OOM)
|
||||
|
||||
**Symptoms**: CUDA OOM errors during model loading or inference, process crashes with memory errors
|
||||
|
||||
**Solution**:
|
||||
```python
|
||||
# Step 1: Enable TP with smallest degree
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2)
|
||||
|
||||
# Step 2: If still OOM, increase TP degree
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=4)
|
||||
|
||||
```
|
||||
|
||||
### Common Issue 2: Divisibility Error
|
||||
|
||||
**Symptoms**: Error like "Model dimension X not divisible by tensor_parallel_size Y"
|
||||
|
||||
**Solutions**:
|
||||
1. Check model-specific constraints (e.g., Z-Image only supports TP=2)
|
||||
2. Use a smaller TP size that divides model dimensions
|
||||
3. Consult [Supported Models](../../diffusion_features.md#supported-models) for compatible TP sizes
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
1. ✅ **Enable TP** - Set `--tensor-parallel-size` to reduce per-GPU memory
|
||||
2. ✅ **Increase TP size** - Only increase if OOM persists
|
||||
3. ⚠️ **Text encoder not sharded** - Known limitation
|
||||
Reference in new issue
Block a user