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

[Doc] Adding models/pipelines/features Tutorial (#1196)

Browse Source 
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: dongbo910220 <32610838+dongbo910220@users.noreply.github.com>
This commit is contained in:
Didan DengCopilotdongbo910220
 authored and GitHub committed 2026-02-09 15:53:01 +08:00
1 parent c6b526bc2e
commit 29e6a29e54
8 files changed
+2518 -561
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion/parallelism_acceleration.md
-276
View File
@@ -233,143 +233,6 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
| Hybrid Ulysses + Ring | 2 | 2 | 24.3s | 1.87x |
##### How to parallelize a new model
NOTE: "Terminology: SP vs CP"
Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py).
We use "Sequence Parallelism" to align with vLLM-Omni's terminology.
---
###### Non-intrusive `_sp_plan` (Recommended)
The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
**Requirements for `forward()` function:**
- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations)
- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule**
**When to create a submodule:**
```python
# ❌ BAD: Inline operations - hooks cannot intercept
def forward(self, x, cap_feats):
unified = torch.cat([x, cap_feats], dim=1) # Cannot be sharded via _sp_plan
...
# ✅ GOOD: Extract into a submodule
class UnifiedPrepare(nn.Module):
def forward(self, x, cap_feats):
return torch.cat([x, cap_feats], dim=1) # Now can be sharded via _sp_plan
class MyModel(nn.Module):
def __init__(self):
self.unified_prepare = UnifiedPrepare() # Submodule
def forward(self, x, cap_feats):
unified = self.unified_prepare(x, cap_feats) # Hook can intercept here
```
---
###### Defining `_sp_plan`
**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference):
```python
from vllm_omni.diffusion.distributed.sp_plan import (
SequenceParallelInput, # Corresponds to diffusers' ContextParallelInput
SequenceParallelOutput, # Corresponds to diffusers' ContextParallelOutput
)
```
| Parameter | Description |
|-----------|-------------|
| `split_dim` | Dimension to split/gather (usually `1` for sequence) |
| `expected_dims` | Expected tensor rank for validation (optional) |
| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors |
| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) |
**Key naming convention:**
| Key | Meaning | Python equivalent |
|-----|---------|-------------------|
| `""` | Root model | `model` |
| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` |
| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` |
| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` |
**Dictionary key types:**
| Key type | `split_output` | Description |
|----------|----------------|-------------|
| `"param_name"` (str) | `False` | Shard **input parameter** by name |
| `0`, `1` (int) | `True` | Shard **output tuple** by index |
**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)):
```python
class MyTransformer(nn.Module):
_sp_plan = {
# Shard rope module OUTPUTS (returns tuple)
"rope": {
0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # cos
1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # sin
},
# Shard transformer block INPUT parameter
"blocks.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
},
# Gather at final projection
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
}
```
---
###### Hook flow
```
Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ...
↓
... → [SequenceParallelGatherHook: post_forward] → Output
```
1. **SplitHook** shards tensors before/after the target module
2. **Attention layers** handle Ulysses/Ring communication internally
3. **GatherHook** collects sharded outputs
The framework automatically applies these hooks when `sequence_parallel_size > 1`.
---
###### Method 2: Intrusive modification (For complex cases)
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`:
```python
from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
def forward(self, hidden_states, ...):
if self.parallel_config.sequence_parallel_size > 1:
hidden_states = sp_shard(hidden_states, dim=1)
# ... computation ...
output = sp_gather(output, dim=1)
return output
```
---
###### Choosing the right approach
| Scenario | Approach |
|----------|----------|
| Standard transformer | `_sp_plan` |
| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` |
| Dynamic/conditional sharding | Intrusive modification |
### CFG-Parallel
#### Offline Inference
@@ -426,142 +289,3 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa
```bash
vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
```
#### How to parallelize a pipeline
This section describes how to add CFG-Parallel to a diffusion **pipeline**. We use the Qwen-Image pipeline (`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`) as the reference implementation.
In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes sequentially:
- positive (prompt-conditioned)
- negative (negative-prompt-conditioned)
CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results.
vLLM-omni provides `CFGParallelMixin` base class that encapsulates the CFG parallel logic. By inheriting from this mixin and calling its methods, pipelines can easily implement CFG parallel without writing repetitive code.
**Key Methods in CFGParallelMixin:**
- `predict_noise_maybe_with_cfg()`: Automatically handles CFG parallel noise prediction
- `scheduler_step_maybe_with_cfg()`: Scheduler step with automatic CFG rank synchronization
**Example Implementation:**
```python
class QwenImageCFGParallelMixin(CFGParallelMixin):
"""
Base Mixin class for Qwen Image pipelines providing shared CFG methods.
"""
def diffuse(
self,
prompt_embeds: torch.Tensor,
prompt_embeds_mask: torch.Tensor,
negative_prompt_embeds: torch.Tensor,
negative_prompt_embeds_mask: torch.Tensor,
latents: torch.Tensor,
img_shapes: torch.Tensor,
txt_seq_lens: torch.Tensor,
negative_txt_seq_lens: torch.Tensor,
timesteps: torch.Tensor,
do_true_cfg: bool,
guidance: torch.Tensor,
true_cfg_scale: float,
image_latents: torch.Tensor | None = None,
cfg_normalize: bool = True,
additional_transformer_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor:
self.transformer.do_true_cfg = do_true_cfg
for i, t in enumerate(timesteps):
timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
# Prepare kwargs for positive (conditional) prediction
positive_kwargs = {
"hidden_states": latents,
"timestep": timestep / 1000,
"guidance": guidance,
"encoder_hidden_states_mask": prompt_embeds_mask,
"encoder_hidden_states": prompt_embeds,
"img_shapes": img_shapes,
"txt_seq_lens": txt_seq_lens,
}
# Prepare kwargs for negative (unconditional) prediction
if do_true_cfg:
negative_kwargs = {
"hidden_states": latents,
"timestep": timestep / 1000,
"guidance": guidance,
"encoder_hidden_states_mask": negative_prompt_embeds_mask,
"encoder_hidden_states": negative_prompt_embeds,
"img_shapes": img_shapes,
"txt_seq_lens": negative_txt_seq_lens,
}
else:
negative_kwargs = None
# Predict noise with automatic CFG parallel handling
# - In CFG parallel mode: rank0 computes positive, rank1 computes negative
# - Automatically gathers results and combines them on rank0
noise_pred = self.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=true_cfg_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=cfg_normalize,
)
# Step scheduler with automatic CFG synchronization
# - Only rank0 computes the scheduler step
# - Automatically broadcasts updated latents to all ranks
latents = self.scheduler_step_maybe_with_cfg(
noise_pred, t, latents, do_true_cfg
)
return latents
```
**How it works:**
1. Prepare separate `positive_kwargs` and `negative_kwargs` for conditional and unconditional predictions
2. Call `predict_noise_maybe_with_cfg()` which:
- Detects if CFG parallel is enabled (`get_classifier_free_guidance_world_size() > 1`)
- Distributes computation: rank0 processes positive, rank1 processes negative
- Gathers predictions and combines them using `combine_cfg_noise()` on rank0
- Returns combined noise prediction (only valid on rank0)
3. Call `scheduler_step_maybe_with_cfg()` which:
- Only rank0 computes the scheduler step
- Broadcasts the updated latents to all ranks for synchronization
**How to customize**
Some pipelines may need to customize the following functions in `CFGParallelMixin`:
1. You may need to edit `predict_noise` function for custom behaviors.
```python
def predict_noise(self, *args, **kwargs):
"""
Forward pass through transformer to predict noise.
Subclasses should override this if they need custom behavior,
but the default implementation calls self.transformer.
"""
return self.transformer(*args, **kwargs)[0]
```
2. The default normalization function after combining the noise predictions from both branches is as follows. You may need to customize it.
```python
def cfg_normalize_function(self, noise_pred, comb_pred):
"""
Normalize the combined noise prediction.
Args:
noise_pred: positive noise prediction
comb_pred: combined noise prediction after CFG
Returns:
Normalized noise prediction tensor
"""
cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
noise_pred = comb_pred * (cond_norm / noise_norm)
return noise_pred
```