[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
@@ -0,0 +1,356 @@
|
||||
# Pipeline Parallel
|
||||
|
||||
This section describes how to add Pipeline Parallelism (PP) to a diffusion pipeline. We use the Wan2.2 text-to-video and
|
||||
image-to-video pipelines as the reference implementations.
|
||||
|
||||
## Implementation Checklist
|
||||
|
||||
Adding Pipeline Parallel support requires:
|
||||
|
||||
1. ✅ **Inherit the mixin** - add `PipelineParallelMixin` before `CFGParallelMixin` in the pipeline class
|
||||
2. ✅ **Make stage forward resumable** - support `intermediate_tensors` in the transformer forward path
|
||||
3. ✅ **Return the right object type** - `IntermediateTensors` on non-last PP ranks, final model output on the last rank
|
||||
4. ✅ **Use the standard denoising loop** - call `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
|
||||
5. ✅ **Test parity** - compare PP results against the single-GPU baseline
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
- [Testing](#testing)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Reference Implementations](#reference-implementations)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
### What is Pipeline Parallelism?
|
||||
|
||||
Pipeline Parallelism splits the denoising transformer into multiple sequential stages and places each stage on a
|
||||
different rank. Instead of every rank holding the full DiT, each PP rank owns only a slice of the layers.
|
||||
|
||||
For each denoising step:
|
||||
|
||||
1. Rank 0 starts the forward pass using the current latents.
|
||||
2. Each intermediate rank receives hidden states from the previous rank, runs its local layer slice, and forwards the
|
||||
intermediate tensors downstream.
|
||||
3. The last PP rank produces the final noise prediction.
|
||||
4. The last PP rank applies the scheduler step and sends the updated latents back to rank 0 for the next timestep.
|
||||
|
||||
This reduces per-rank model memory and enables larger diffusion transformers to run across multiple GPUs. It can also be
|
||||
combined with CFG-Parallel, where each PP pipeline carries one CFG branch.
|
||||
|
||||
### Architecture
|
||||
|
||||
vLLM-Omni provides `PipelineParallelMixin` to encapsulate the PP communication pattern for diffusion pipelines.
|
||||
|
||||
| Method | Purpose | Automatic Behavior |
|
||||
|-----------------------------------|-------------------------------|---------------------------------------------------------------------------------|
|
||||
| `diffuse()` | Denoising loop boundary | Wrapped by `PipelineParallelMixin` to flush pending async PP sends on exit |
|
||||
| `predict_noise_maybe_with_cfg()` | Predict noise with PP support | Runs partial forwards on non-last PP ranks, combines with CFG logic when needed |
|
||||
| `scheduler_step_maybe_with_cfg()` | Step scheduler with PP sync | Runs scheduler on the last PP rank, returns updated latents to rank 0 |
|
||||
| `_sync_pp_send()` | Flush pending async sends | Waits for outstanding `isend` handles before later collectives or decode |
|
||||
|
||||
`PipelineParallelMixin` is intentionally a pipeline-level abstraction. Your model-specific `predict_noise()` still
|
||||
defines how a local stage executes.
|
||||
|
||||
### How It Works
|
||||
|
||||
`predict_noise_maybe_with_cfg()` automatically switches between these modes:
|
||||
|
||||
- **PP disabled** (`pipeline_parallel_size == 1`):
|
||||
- Falls back to `CFGParallelMixin.predict_noise_maybe_with_cfg()`
|
||||
- **PP only** (`pipeline_parallel_size > 1`, `cfg_parallel_size == 1`):
|
||||
- Rank 0 starts with the input latents
|
||||
- Middle ranks receive `intermediate_tensors`, run their local layer range, and asynchronously send downstream
|
||||
- The last rank returns the final noise prediction
|
||||
- When CFG is enabled in this mode (sequential CFG), both the positive and negative branches run through the same
|
||||
PP pipeline, doubling the communication volume per denoising step. Prefer PP + CFG-Parallel when
|
||||
`cfg_parallel_size > 1` is available to avoid this overhead.
|
||||
- **PP + CFG-Parallel** (`pipeline_parallel_size > 1`, `cfg_parallel_size > 1`):
|
||||
- Each PP pipeline carries one CFG branch
|
||||
- The last PP rank all-gathers across the CFG group
|
||||
- CFG combination happens on every last PP rank in the CFG group, matching the non-PP CFG-parallel behavior
|
||||
|
||||
`scheduler_step_maybe_with_cfg()` keeps the denoising loop consistent:
|
||||
|
||||
- **PP disabled**:
|
||||
- Falls back to `scheduler_step_maybe_with_cfg()`
|
||||
- **PP enabled**:
|
||||
- Only the last PP rank has `noise_pred` and runs the scheduler step
|
||||
- The resulting latents are sent back to rank 0
|
||||
- Rank 0 receives an `AsyncLatents` wrapper that resolves only when the tensor is next consumed
|
||||
|
||||
This asynchronous design avoids unnecessary blocking between denoising steps.
|
||||
|
||||
When a pipeline class is defined, `PipelineParallelMixin` wraps its `diffuse()` method and calls `_sync_pp_send()` in a
|
||||
`finally` block after `diffuse()` returns or raises. This keeps model code free of explicit PP cleanup while still
|
||||
ensuring that the final non-blocking PP send completes before decode or later collectives.
|
||||
|
||||
---
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Inherit `PipelineParallelMixin` and `CFGParallelMixin`
|
||||
|
||||
`PipelineParallelMixin` **requires** `CFGParallelMixin` and must appear before it in the class MRO. This is enforced at
|
||||
class definition time via `__init_subclass__`: defining a pipeline that inherits `PipelineParallelMixin` without
|
||||
`CFGParallelMixin`, or after `CFGParallelMixin`, raises a `TypeError` immediately on import.
|
||||
|
||||
`PipelineParallelMixin` delegates noise prediction, CFG combination, and scheduler stepping to `CFGParallelMixin`, which
|
||||
supplies `predict_noise()`, `predict_noise_maybe_with_cfg()`, `scheduler_step_maybe_with_cfg()`, and
|
||||
`combine_cfg_noise()`.
|
||||
|
||||
**Example:**
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
|
||||
...
|
||||
```
|
||||
|
||||
The order matters: `PipelineParallelMixin` must be listed before `CFGParallelMixin` so calls to
|
||||
`predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` resolve to the PP-aware wrappers, while their
|
||||
`super()` calls delegate to the CFG implementation when PP is disabled or after the last PP stage.
|
||||
|
||||
### Step 2: Make model forward and `predict_noise()` PP-aware
|
||||
|
||||
The PP mixin injects `intermediate_tensors` into the normal `predict_noise()` call. Your model forward path must support
|
||||
two inputs:
|
||||
|
||||
- Normal input from rank 0, typically passed as `hidden_states` or `x`
|
||||
- `intermediate_tensors` from upstream PP ranks
|
||||
|
||||
The standard model forward pattern is:
|
||||
|
||||
1. If `intermediate_tensors` is present, read the local hidden state from it.
|
||||
2. Run only this rank's layer slice.
|
||||
3. Return `IntermediateTensors(...)` on non-last PP ranks.
|
||||
4. Return the final model output on the last PP rank. `CFGParallelMixin.predict_noise()` already follows this contract
|
||||
for common pipelines.
|
||||
|
||||
**Minimal example:**
|
||||
|
||||
```python
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
|
||||
|
||||
|
||||
def forward(self, hidden_states=None, intermediate_tensors=None, **kwargs):
|
||||
if intermediate_tensors is not None:
|
||||
hidden_states = intermediate_tensors["hidden_states"]
|
||||
|
||||
for i in range(self.start_layer, self.end_layer):
|
||||
hidden_states = self.layers[i](hidden_states)
|
||||
|
||||
pp_group = get_pp_group()
|
||||
if not pp_group.is_last_rank:
|
||||
return IntermediateTensors({"hidden_states": hidden_states})
|
||||
return (hidden_states,)
|
||||
```
|
||||
|
||||
### Step 3: Partition the transformer layers
|
||||
|
||||
The local module on each PP rank must expose only that rank's layer slice. In the reference tests and vLLM model
|
||||
implementations, this is typically done with vLLM utilities such as `make_layers(...)`, which fill missing ranges with
|
||||
`PPMissingLayer`. Besides partitioning the layers themselves, model authors should also wire
|
||||
`make_empty_intermediate_tensors_factory(...)` for intermediate tensor allocation and `is_pp_missing_parameter(...)` for
|
||||
PP-aware weight loading.
|
||||
|
||||
To prepare a transformer for PP, implement the following pieces in order.
|
||||
|
||||
#### 3.1 Split the transformer layers across PP ranks
|
||||
|
||||
Each PP rank should own only its local layer range, typically exposed as `[start_layer, end_layer)`. In practice this is
|
||||
usually done with `make_layers(...)`, which constructs local layers and fills missing ranges with `PPMissingLayer`.
|
||||
|
||||
The goal is:
|
||||
|
||||
- every PP rank knows its `[start_layer, end_layer)` range
|
||||
- non-local layers do not execute on this rank
|
||||
- the forward can resume from incoming `intermediate_tensors`
|
||||
|
||||
By default, layers are distributed evenly across PP ranks using `get_pp_indices()`. When the number of layers is not
|
||||
divisible by the PP size, remaining layers are assigned to middle partitions to balance compute and memory. You can
|
||||
override this with the `VLLM_PP_LAYER_PARTITION` environment variable to specify an exact per-rank layer count:
|
||||
|
||||
```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.
|
||||
|
||||
#### 3.2 Expose `make_empty_intermediate_tensors`
|
||||
|
||||
Your transformer module should expose `self.make_empty_intermediate_tensors`, typically created with
|
||||
`make_empty_intermediate_tensors_factory(...)`.
|
||||
|
||||
This is important because PP ranks need a consistent way to allocate placeholder `IntermediateTensors` with the expected
|
||||
keys and hidden dimension.
|
||||
|
||||
**Example:**
|
||||
|
||||
```python
|
||||
from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
|
||||
|
||||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
||||
["hidden_states"],
|
||||
inner_dim,
|
||||
)
|
||||
```
|
||||
|
||||
For Wan2.2, the intermediate payload between PP stages is the token sequence stored under `"hidden_states"`, so the
|
||||
factory is created with that key and the transformer hidden size.
|
||||
|
||||
#### 3.3 Return `IntermediateTensors` on non-last PP ranks
|
||||
|
||||
Your model `forward()` or custom `predict_noise()` implementation should consume `intermediate_tensors` on non-first
|
||||
ranks and return `IntermediateTensors(...)` on non-last ranks.
|
||||
|
||||
This lets each PP stage resume from the upstream hidden states and pass the local result to the next stage.
|
||||
|
||||
#### 3.4 Skip non-local weights during `load_weights()`
|
||||
|
||||
When a model is PP-partitioned, many parameters in the checkpoint belong to layers that do not exist on the current
|
||||
rank. `load_weights()` must therefore skip parameters for missing PP stages using `is_pp_missing_parameter(...)`.
|
||||
|
||||
If this is not done, weight loading will either fail or incorrectly try to load tensors into `PPMissingLayer`
|
||||
placeholders.
|
||||
|
||||
The Wan2.2 transformer is the best reference here: it uses `is_pp_missing_parameter(...)` before loading both remapped
|
||||
and fused parameters.
|
||||
|
||||
If your model has multiple transformer variants, PP still works as long as each selected transformer obeys the same
|
||||
contract.
|
||||
|
||||
### Step 4: Use the standard denoising contract
|
||||
|
||||
vLLM-Omni diffusion pipelines already route denoising through `diffuse()`, `predict_noise_maybe_with_cfg()`, and
|
||||
`scheduler_step_maybe_with_cfg()`. `PipelineParallelMixin` overrides those standard helpers when
|
||||
`pipeline_parallel_size > 1`, so model integrations should not add separate PP-specific helper names or manual post-loop
|
||||
synchronization.
|
||||
|
||||
In PP mode:
|
||||
|
||||
- Non-last PP ranks return `None` from `predict_noise_maybe_with_cfg()` after sending `IntermediateTensors` downstream.
|
||||
- The last PP rank returns the final prediction, and with CFG-Parallel every last PP rank in the CFG group receives the
|
||||
combined result.
|
||||
- Rank 0 receives `AsyncLatents` from `scheduler_step_maybe_with_cfg()`, which resolves only when the latents are
|
||||
consumed.
|
||||
- Pending non-blocking sends are flushed automatically when `diffuse()` exits.
|
||||
|
||||
---
|
||||
|
||||
## Testing
|
||||
|
||||
Use an offline inference script with `pipeline_parallel_size > 1`.
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
For PP + CFG-Parallel together:
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
**Verify:**
|
||||
|
||||
1. The run completes without hangs at the PP stage boundary
|
||||
2. Output quality matches the non-PP baseline within normal numerical variance
|
||||
3. Peak memory per GPU drops compared with the single-rank model
|
||||
4. There are no outstanding communication errors before decode
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Issue: `TypeError` on import — `CFGParallelMixin` missing
|
||||
|
||||
**Symptoms:** Importing a pipeline that inherits `PipelineParallelMixin` raises:
|
||||
|
||||
```
|
||||
TypeError: YourPipeline inherits PipelineParallelMixin but not CFGParallelMixin.
|
||||
```
|
||||
|
||||
or:
|
||||
|
||||
```
|
||||
TypeError: YourPipeline must inherit PipelineParallelMixin before CFGParallelMixin ...
|
||||
```
|
||||
|
||||
**Cause:** `PipelineParallelMixin` enforces via `__init_subclass__` that the subclass also inherits `CFGParallelMixin`
|
||||
and lists `PipelineParallelMixin` first in the MRO.
|
||||
|
||||
**Solution:** Add `CFGParallelMixin` to the pipeline's base classes after `PipelineParallelMixin`:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
|
||||
|
||||
|
||||
class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
|
||||
...
|
||||
```
|
||||
|
||||
### Issue: Non-last PP ranks crash when calling `predict_noise`
|
||||
|
||||
**Symptoms:** Shape or missing-input errors on ranks other than the first or last PP rank.
|
||||
|
||||
**Cause:** The model forward path assumes direct input tensors and ignores `intermediate_tensors`.
|
||||
|
||||
**Solution:** Update the transformer `forward()` or custom `predict_noise()` path to load hidden states from
|
||||
`intermediate_tensors` when present.
|
||||
|
||||
### Issue: PP outputs differ from single-GPU baseline
|
||||
|
||||
**Symptoms:** The PP run finishes but produces numerically inconsistent outputs.
|
||||
|
||||
**Causes & Solutions:**
|
||||
|
||||
- **Local layer partitioning is wrong**
|
||||
- Verify each rank runs only its `[start_layer, end_layer)` slice
|
||||
- **Non-last ranks return a plain tensor instead of `IntermediateTensors`**
|
||||
- Return `IntermediateTensors({...})` until the last PP stage
|
||||
- **CFG branch wiring is incorrect**
|
||||
- When CFG is enabled, confirm positive and negative kwargs are passed exactly as in the non-PP path
|
||||
|
||||
---
|
||||
|
||||
## Reference Implementations
|
||||
|
||||
Complete examples in the codebase:
|
||||
|
||||
| Component | Path | Notes |
|
||||
|-------------------------|------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|
|
||||
| `PipelineParallelMixin` | `vllm_omni/diffusion/distributed/pipeline_parallel.py` | Core PP communication and scheduler helpers |
|
||||
| `CFGParallelMixin` | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Default `predict_noise()` tuple normalization and CFG helper fallback |
|
||||
| Wan2.2 transformer | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference for layer partitioning, `IntermediateTensors`, `make_empty_intermediate_tensors`, and PP-aware weight loading |
|
||||
| Wan2.2 T2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Reference PP + CFG integration for text-to-video |
|
||||
| Wan2.2 I2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py` | Reference PP + CFG integration for image-to-video |
|
||||
| PP tests | `tests/diffusion/distributed/test_pipeline_parallel.py` | Baseline parity and async communication tests |
|
||||
@@ -28,7 +28,6 @@ Frame interpolation is currently supported for:
|
||||
|
||||
- `WanPipeline` (Wan2.2 text-to-video)
|
||||
- `WanImageToVideoPipeline`
|
||||
- `Wan22TI2VPipeline`
|
||||
|
||||
## Request Parameters
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -35,14 +35,15 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
|
||||
|
||||
Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
|
||||
|
||||
| Method | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
|
||||
| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
|
||||
| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
|
||||
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
|
||||
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
|
||||
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
|
||||
| Method | Description | Best For |
|
||||
|------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------------------------------------------------------------|
|
||||
| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
|
||||
| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
|
||||
| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
|
||||
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
|
||||
| **[Pipeline Parallelism](diffusion/parallelism/pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages | Large diffusion transformers that need lower per-GPU model memory |
|
||||
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
|
||||
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
|
||||
|
||||
#### Startup Optimization
|
||||
|
||||
@@ -105,32 +106,32 @@ The following tables show which models support each feature:
|
||||
|
||||
### ImageGen
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
|
||||
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
|
||||
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
> Notes:
|
||||
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
|
||||
@@ -138,15 +139,15 @@ The following tables show which models support each feature:
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
**Frame Interpolation Support**
|
||||
|
||||
@@ -155,9 +156,9 @@ The following tables show which models support each feature:
|
||||
|
||||
### AudioGen
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
|
||||
|
||||
## Feature Compatibility
|
||||
@@ -257,7 +258,7 @@ Measured on NVIDIA H800:
|
||||
|
||||
**Parallelism Methods:**
|
||||
|
||||
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
|
||||
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
|
||||
|
||||
**Memory Optimization:**
|
||||
|
||||
|
||||
@@ -142,6 +142,42 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable torch.compile and force eager execution.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["cache_dit", "tea_cache"],
|
||||
help=(
|
||||
"Cache backend to use for acceleration. "
|
||||
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
|
||||
"Default: None (no cache acceleration)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "int8", "gguf"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
|
||||
# Distributed and parallel execution
|
||||
parser.add_argument(
|
||||
"--ulysses-degree",
|
||||
type=int,
|
||||
@@ -173,28 +209,6 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable torch.compile and force eager execution.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["cache_dit", "tea_cache"],
|
||||
help=(
|
||||
"Cache backend to use for acceleration. "
|
||||
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
|
||||
"Default: None (no cache acceleration)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-hsdp",
|
||||
action="store_true",
|
||||
@@ -219,16 +233,10 @@ def parse_args() -> argparse.Namespace:
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "int8", "gguf"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
"--pipeline-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of pipeline parallel stages.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-config",
|
||||
@@ -326,6 +334,7 @@ def main():
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
)
|
||||
omni_kwargs = dict(
|
||||
model=args.model,
|
||||
@@ -366,7 +375,8 @@ def main():
|
||||
print(f" diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}")
|
||||
print(
|
||||
f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
|
||||
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
|
||||
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
|
||||
f" pipeline_parallel_size={args.pipeline_parallel_size}"
|
||||
)
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
@@ -142,6 +142,32 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-config",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "int8", "gguf"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
|
||||
# Distributed and parallel execution
|
||||
parser.add_argument(
|
||||
"--ulysses-mode",
|
||||
type=str,
|
||||
@@ -174,41 +200,23 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pipeline-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of pipeline parallel stages.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-expert-parallel",
|
||||
action="store_true",
|
||||
help="Enable expert parallelism for MoE layers.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-config",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "int8", "gguf"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-hsdp",
|
||||
action="store_true",
|
||||
@@ -262,6 +270,7 @@ def main():
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
)
|
||||
|
||||
@@ -307,7 +316,8 @@ def main():
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
|
||||
f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
|
||||
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
|
||||
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, pipeline_parallel_size={args.pipeline_parallel_size},"
|
||||
f" enable_expert_parallel={args.enable_expert_parallel}"
|
||||
)
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
@@ -0,0 +1,802 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
import socket
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
|
||||
|
||||
import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
destroy_distributed_env,
|
||||
get_classifier_free_guidance_rank,
|
||||
get_pp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.parallel]
|
||||
|
||||
|
||||
def _find_free_port() -> str:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return str(s.getsockname()[1])
|
||||
|
||||
|
||||
def update_environment_variables(envs_dict: dict[str, str]) -> None:
|
||||
for k, v in envs_dict.items():
|
||||
os.environ[k] = v
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared stubs used by both unit and distributed tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FakeWork:
|
||||
"""Drop-in for torch.distributed.Work that records whether wait() was called."""
|
||||
|
||||
def __init__(self):
|
||||
self.waited = False
|
||||
|
||||
def wait(self):
|
||||
self.waited = True
|
||||
|
||||
|
||||
class SimpleScheduler:
|
||||
"""Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred."""
|
||||
|
||||
def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False):
|
||||
return (latents - 0.1 * noise_pred,)
|
||||
|
||||
|
||||
class FakeVAE:
|
||||
def __init__(self, distributed_enabled: bool = False):
|
||||
self.calls = 0
|
||||
self.distributed_enabled = distributed_enabled
|
||||
|
||||
def decode(self, z: torch.Tensor):
|
||||
"""Original decode docstring."""
|
||||
self.calls += 1
|
||||
return (z + 1,)
|
||||
|
||||
def is_distributed_enabled(self) -> bool:
|
||||
return self.distributed_enabled
|
||||
|
||||
|
||||
class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
|
||||
"""Minimal pipeline used to exercise PipelineParallelMixin.
|
||||
|
||||
Uses vLLM's ``make_layers`` for layer partitioning — the same utility used
|
||||
by real DiT models — so the PP layer-split logic is exercised faithfully.
|
||||
|
||||
Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i``
|
||||
is initialized identically on every rank regardless of which ranks are active,
|
||||
allowing the distributed output to be compared against the single-GPU baseline.
|
||||
|
||||
Args:
|
||||
num_layers: Total number of Linear layers.
|
||||
dim: Input / hidden dimension.
|
||||
seed: Base RNG seed; layer ``i`` uses ``seed + i``.
|
||||
device: Target device for layer weights (default: CPU).
|
||||
dtype: Target dtype for layer weights (default: float32).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_layers: int = 4,
|
||||
dim: int = 64,
|
||||
seed: int = 42,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
):
|
||||
self.start_layer, self.end_layer, self.layers = make_layers(
|
||||
num_layers,
|
||||
lambda prefix: torch.nn.Linear(dim, dim, bias=False),
|
||||
prefix="layers",
|
||||
)
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
if not isinstance(layer, PPMissingLayer):
|
||||
torch.manual_seed(seed + i)
|
||||
torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02)
|
||||
|
||||
self.layers.to(device=device, dtype=dtype)
|
||||
|
||||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim)
|
||||
self.scheduler = SimpleScheduler()
|
||||
|
||||
def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors:
|
||||
"""Layer-split forward pass.
|
||||
|
||||
* First PP rank: uses ``x`` from caller kwargs.
|
||||
* Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]``
|
||||
(which transparently waits for the async receive).
|
||||
* Non-last PP ranks return ``IntermediateTensors``; the last rank
|
||||
returns the plain noise-prediction tensor.
|
||||
"""
|
||||
if intermediate_tensors is not None:
|
||||
x = intermediate_tensors["hidden_states"]
|
||||
|
||||
for i in range(self.start_layer, self.end_layer):
|
||||
x = self.layers[i](x)
|
||||
|
||||
pp_group = get_pp_group()
|
||||
if not pp_group.is_last_rank:
|
||||
return IntermediateTensors({"hidden_states": x})
|
||||
return x
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. AsyncLatents – unit tests (no distributed env required)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAsyncLatents:
|
||||
"""Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
|
||||
|
||||
pytestmark = [pytest.mark.cpu]
|
||||
|
||||
def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
|
||||
return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
|
||||
|
||||
def test_resolve_returns_wrapped_tensor(self):
|
||||
t = torch.randn(2, 4)
|
||||
al = self._make(t)
|
||||
assert al._resolve() is t
|
||||
|
||||
def test_attribute_access_resolves(self):
|
||||
t = torch.randn(2, 4)
|
||||
al = self._make(t)
|
||||
assert al.shape == t.shape
|
||||
assert al.dtype == t.dtype
|
||||
|
||||
def test_torch_function_protocol(self):
|
||||
"""torch ops that receive an AsyncLatents should see the underlying tensor."""
|
||||
t = torch.randn(2, 4)
|
||||
al = self._make(t)
|
||||
mask = torch.ones_like(t)
|
||||
result = mask * al # triggers __torch_function__
|
||||
torch.testing.assert_close(result, mask * t)
|
||||
|
||||
def test_torch_function_with_list_arg(self):
|
||||
"""__torch_function__ must unwrap AsyncLatents inside list/tuple args."""
|
||||
t = torch.randn(2, 4)
|
||||
al = self._make(t)
|
||||
result = torch.cat([al, al], dim=0)
|
||||
torch.testing.assert_close(result, torch.cat([t, t], dim=0))
|
||||
|
||||
def test_torch_tensor_conversion(self):
|
||||
"""torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy)."""
|
||||
t = torch.randn(2, 4)
|
||||
al = self._make(t)
|
||||
result = torch.as_tensor(al)
|
||||
assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage"
|
||||
|
||||
def test_handles_are_waited_before_resolve(self):
|
||||
t = torch.randn(2, 4)
|
||||
h1, h2 = FakeWork(), FakeWork()
|
||||
al = self._make(t, handles=[h1, h2])
|
||||
_ = al.shape # trigger resolution
|
||||
assert h1.waited and h2.waited, "Not all handles were waited on"
|
||||
|
||||
def test_postproc_callbacks_invoked_on_resolve(self):
|
||||
t = torch.randn(2, 4)
|
||||
log: list[int] = []
|
||||
al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)])
|
||||
_ = al.shape
|
||||
assert log == [1, 2], f"postproc not called in order: {log}"
|
||||
|
||||
def test_idempotent_resolve(self):
|
||||
"""handle.wait() must not be called twice if _resolve() is called twice."""
|
||||
t = torch.randn(2, 4)
|
||||
h = FakeWork()
|
||||
al = self._make(t, handles=[h])
|
||||
_ = al.shape # first resolve
|
||||
h.waited = False # reset sentinel
|
||||
_ = al.dtype # second resolve
|
||||
assert not h.waited, "handle.wait() was called a second time"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. _sync_pp_send / diffuse wrapper – unit tests (no distributed env required)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSyncPPSend:
|
||||
"""Verifies PipelineParallelMixin's internal PP-send flush."""
|
||||
|
||||
pytestmark = [pytest.mark.cpu]
|
||||
|
||||
@staticmethod
|
||||
def _make_pipeline() -> PipelineParallelMixin:
|
||||
# Instantiate a bare mixin — no layers, no distributed env needed.
|
||||
# _sync_pp_send only touches _pp_send_work, so this is sufficient.
|
||||
class _BarePP(PipelineParallelMixin, CFGParallelMixin):
|
||||
pass
|
||||
|
||||
return _BarePP()
|
||||
|
||||
def test_noop_when_work_list_empty(self):
|
||||
pipeline = self._make_pipeline()
|
||||
pipeline._sync_pp_send()
|
||||
assert pipeline._pp_send_work == []
|
||||
|
||||
def test_waits_all_pending_handles(self):
|
||||
pipeline = self._make_pipeline()
|
||||
works = [FakeWork(), FakeWork(), FakeWork()]
|
||||
pipeline._pp_send_work = works
|
||||
pipeline._sync_pp_send()
|
||||
assert all(w.waited for w in works), "Some handles were not waited on"
|
||||
|
||||
def test_clears_work_list_after_sync(self):
|
||||
pipeline = self._make_pipeline()
|
||||
pipeline._pp_send_work = [FakeWork()]
|
||||
pipeline._sync_pp_send()
|
||||
assert pipeline._pp_send_work == []
|
||||
|
||||
|
||||
class TestDiffuseWrapper:
|
||||
"""Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
|
||||
|
||||
pytestmark = [pytest.mark.cpu]
|
||||
|
||||
def test_diffuse_flushes_pending_sends_on_success(self):
|
||||
work = FakeWork()
|
||||
|
||||
class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
|
||||
def diffuse(self):
|
||||
self._pp_send_work = [work]
|
||||
return "done"
|
||||
|
||||
pipeline = _DiffusePP()
|
||||
|
||||
assert pipeline.diffuse() == "done"
|
||||
assert work.waited
|
||||
assert pipeline._pp_send_work == []
|
||||
|
||||
def test_diffuse_flushes_pending_sends_on_exception(self):
|
||||
work = FakeWork()
|
||||
|
||||
class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
|
||||
def diffuse(self):
|
||||
self._pp_send_work = [work]
|
||||
raise RuntimeError("boom")
|
||||
|
||||
pipeline = _DiffusePP()
|
||||
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
pipeline.diffuse()
|
||||
assert work.waited
|
||||
assert pipeline._pp_send_work == []
|
||||
|
||||
def test_diffuse_wrapper_preserves_metadata(self):
|
||||
class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
|
||||
def diffuse(self):
|
||||
"""Original diffuse docstring."""
|
||||
return "done"
|
||||
|
||||
assert _DiffusePP.diffuse.__name__ == "diffuse"
|
||||
assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring."
|
||||
|
||||
|
||||
class TestVaeDecodeGuard:
|
||||
pytestmark = [pytest.mark.cpu]
|
||||
|
||||
@staticmethod
|
||||
def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
|
||||
class _DecodePP(PipelineParallelMixin, CFGParallelMixin):
|
||||
def __init__(self):
|
||||
self.vae = FakeVAE(distributed_enabled=distributed_enabled)
|
||||
|
||||
return _DecodePP()
|
||||
|
||||
@staticmethod
|
||||
def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None:
|
||||
monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size)
|
||||
monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage)
|
||||
|
||||
def test_calls_original_decode_when_pp_disabled(self, monkeypatch):
|
||||
self._set_rank(monkeypatch, world_size=1, first_stage=True)
|
||||
pipeline = self._make_pipeline()
|
||||
z = torch.ones(2, 3)
|
||||
|
||||
output = pipeline.vae.decode(z)[0]
|
||||
|
||||
assert pipeline.vae.calls == 1
|
||||
torch.testing.assert_close(output, z + 1)
|
||||
|
||||
def test_calls_original_decode_on_first_stage(self, monkeypatch):
|
||||
self._set_rank(monkeypatch, world_size=2, first_stage=True)
|
||||
pipeline = self._make_pipeline()
|
||||
z = torch.ones(2, 3)
|
||||
|
||||
output = pipeline.vae.decode(z)[0]
|
||||
|
||||
assert pipeline.vae.calls == 1
|
||||
torch.testing.assert_close(output, z + 1)
|
||||
|
||||
def test_skips_decode_on_non_first_stage(self, monkeypatch):
|
||||
self._set_rank(monkeypatch, world_size=2, first_stage=False)
|
||||
pipeline = self._make_pipeline()
|
||||
z = torch.ones(2, 3)
|
||||
|
||||
output = pipeline.vae.decode(z)
|
||||
|
||||
assert pipeline.vae.calls == 0
|
||||
assert output == (None,)
|
||||
|
||||
def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch):
|
||||
self._set_rank(monkeypatch, world_size=2, first_stage=False)
|
||||
pipeline = self._make_pipeline(distributed_enabled=True)
|
||||
z = torch.ones(2, 3)
|
||||
|
||||
output = pipeline.vae.decode(z)[0]
|
||||
|
||||
assert pipeline.vae.calls == 1
|
||||
torch.testing.assert_close(output, z + 1)
|
||||
|
||||
def test_decode_wrapper_preserves_metadata(self):
|
||||
pipeline = self._make_pipeline()
|
||||
|
||||
assert pipeline.vae.decode.__name__ == "decode"
|
||||
assert pipeline.vae.decode.__doc__ == "Original decode docstring."
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_pipeline_parallel_requires_cfg_mixin():
|
||||
with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
|
||||
|
||||
class _MissingCFG(PipelineParallelMixin):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_pipeline_parallel_requires_mro_before_cfg_mixin():
|
||||
with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
|
||||
|
||||
class _WrongOrder(CFGParallelMixin, PipelineParallelMixin):
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Distributed test helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
|
||||
"""Initialise the distributed environment for a spawned worker."""
|
||||
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
|
||||
current_omni_platform.set_device(device)
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": master_port,
|
||||
}
|
||||
)
|
||||
init_distributed_environment()
|
||||
return device
|
||||
|
||||
|
||||
def make_pipeline_and_inputs(
|
||||
test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
|
||||
) -> tuple["MockPipelineParallel", dict, dict | None]:
|
||||
"""Create a MockPipelineParallel and seeded inputs from a test_config dict.
|
||||
|
||||
Must be called after ``initialize_model_parallel`` so that ``make_layers``
|
||||
can read the PP group to determine this rank's layer slice.
|
||||
|
||||
Returns ``(pipeline, positive_kwargs, negative_kwargs)``.
|
||||
``negative_kwargs`` is ``None`` when ``do_true_cfg=False``.
|
||||
"""
|
||||
pipeline = MockPipelineParallel(
|
||||
num_layers=test_config["num_layers"],
|
||||
dim=test_config["dim"],
|
||||
seed=test_config["model_seed"],
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
torch.manual_seed(test_config["input_seed"])
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(test_config["input_seed"])
|
||||
pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
|
||||
|
||||
negative_kwargs = None
|
||||
if do_true_cfg:
|
||||
torch.manual_seed(test_config["input_seed"] + 1)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
|
||||
neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
|
||||
negative_kwargs = {"x": neg_x}
|
||||
|
||||
return pipeline, pos_x, negative_kwargs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. isend_tensor_dict / irecv_tensor_dict (2 GPUs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
|
||||
device = init_dist(local_rank, world_size, master_port)
|
||||
initialize_model_parallel(pipeline_parallel_size=world_size)
|
||||
pp_group = get_pp_group()
|
||||
|
||||
if pp_group.is_first_rank:
|
||||
torch.manual_seed(77)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(77)
|
||||
tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
|
||||
handles = pp_group.isend_tensor_dict({"t": tensor})
|
||||
for h in handles:
|
||||
h.wait()
|
||||
result_queue.put(("sent", tensor.cpu()))
|
||||
else:
|
||||
received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
|
||||
result_queue.put(("received", received["t"].cpu()))
|
||||
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
destroy_distributed_env()
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
|
||||
@pytest.mark.parametrize("pp_size", [2])
|
||||
def test_isend_irecv_tensor_dict(pp_size: int):
|
||||
"""isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
q = manager.Queue()
|
||||
|
||||
port = _find_free_port()
|
||||
torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
|
||||
|
||||
results = {label: tensor for label, tensor in [q.get(), q.get()]}
|
||||
torch.testing.assert_close(
|
||||
results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. predict_noise_maybe_with_cfg
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_baseline_cache: dict[tuple, torch.Tensor] = {}
|
||||
|
||||
|
||||
def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
|
||||
"""Compute expected single-GPU output using the same MockPipelineParallel.
|
||||
|
||||
Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
|
||||
Results are cached so identical configs are only computed once.
|
||||
"""
|
||||
key = (
|
||||
test_config["num_layers"],
|
||||
test_config["dim"],
|
||||
test_config["batch_size"],
|
||||
test_config["model_seed"],
|
||||
test_config["input_seed"],
|
||||
test_config["cfg_scale"],
|
||||
dtype,
|
||||
do_true_cfg,
|
||||
)
|
||||
if key in _baseline_cache:
|
||||
return _baseline_cache[key]
|
||||
|
||||
device = init_dist(0, 1, _find_free_port())
|
||||
initialize_model_parallel(pipeline_parallel_size=1)
|
||||
|
||||
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
|
||||
test_config, dtype, device, do_true_cfg=do_true_cfg
|
||||
)
|
||||
|
||||
with torch.inference_mode():
|
||||
noise_pred = pipeline.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=test_config["cfg_scale"],
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
|
||||
destroy_distributed_env()
|
||||
|
||||
_baseline_cache[key] = noise_pred.cpu()
|
||||
return _baseline_cache[key]
|
||||
|
||||
|
||||
def predict_noise_worker(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
master_port: str,
|
||||
pp_size: int,
|
||||
cfg_size: int,
|
||||
do_true_cfg: bool,
|
||||
dtype: torch.dtype,
|
||||
test_config: dict,
|
||||
result_queue,
|
||||
):
|
||||
"""Generic predict-noise worker parameterized by PP and CFG topology."""
|
||||
device = init_dist(local_rank, world_size, master_port)
|
||||
initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
|
||||
|
||||
pp_group = get_pp_group()
|
||||
cfg_rank = get_classifier_free_guidance_rank()
|
||||
|
||||
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
|
||||
test_config, dtype, device, do_true_cfg=do_true_cfg
|
||||
)
|
||||
|
||||
with torch.inference_mode():
|
||||
noise_pred = pipeline.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=test_config["cfg_scale"],
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
# This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
|
||||
# Flush the non-last PP rank's async send before barrier / process teardown.
|
||||
pipeline._sync_pp_send()
|
||||
|
||||
if pp_group.is_last_rank:
|
||||
assert noise_pred is not None
|
||||
if cfg_rank == 0:
|
||||
result_queue.put(noise_pred.cpu())
|
||||
else:
|
||||
assert noise_pred is None
|
||||
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
destroy_distributed_env()
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.parametrize(
|
||||
"pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
|
||||
[
|
||||
pytest.param(
|
||||
2,
|
||||
1,
|
||||
False,
|
||||
torch.float32,
|
||||
4,
|
||||
100,
|
||||
1e-5,
|
||||
1e-5,
|
||||
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
|
||||
id="pp2-no_cfg-float32",
|
||||
),
|
||||
pytest.param(
|
||||
2,
|
||||
1,
|
||||
False,
|
||||
torch.bfloat16,
|
||||
4,
|
||||
100,
|
||||
1e-2,
|
||||
1e-2,
|
||||
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
|
||||
id="pp2-no_cfg-bfloat16",
|
||||
),
|
||||
pytest.param(
|
||||
2,
|
||||
1,
|
||||
True,
|
||||
torch.bfloat16,
|
||||
4,
|
||||
100,
|
||||
1e-2,
|
||||
1e-2,
|
||||
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
|
||||
id="pp2-seq_cfg-bfloat16",
|
||||
),
|
||||
pytest.param(
|
||||
2,
|
||||
2,
|
||||
True,
|
||||
torch.bfloat16,
|
||||
4,
|
||||
100,
|
||||
1e-2,
|
||||
1e-2,
|
||||
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
|
||||
id="pp2-cfg2-bfloat16",
|
||||
),
|
||||
pytest.param(
|
||||
3,
|
||||
1,
|
||||
False,
|
||||
torch.bfloat16,
|
||||
6,
|
||||
100,
|
||||
1e-2,
|
||||
1e-2,
|
||||
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
|
||||
id="pp3-no_cfg-bfloat16",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
|
||||
"""predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
|
||||
test_config = {
|
||||
"num_layers": num_layers,
|
||||
"dim": 64,
|
||||
"batch_size": 2,
|
||||
"cfg_scale": 7.5,
|
||||
"model_seed": 42,
|
||||
"input_seed": input_seed,
|
||||
}
|
||||
|
||||
baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
|
||||
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
pp_q = manager.Queue()
|
||||
|
||||
world_size = pp_size * cfg_size
|
||||
port = _find_free_port()
|
||||
torch.multiprocessing.spawn(
|
||||
predict_noise_worker,
|
||||
args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
|
||||
nprocs=world_size,
|
||||
)
|
||||
|
||||
pp_out = pp_q.get()
|
||||
|
||||
assert baseline_out.shape == pp_out.shape
|
||||
torch.testing.assert_close(
|
||||
pp_out,
|
||||
baseline_out,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. scheduler_step_maybe_with_cfg
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
|
||||
"""Single-GPU reference: predict_noise + scheduler_step."""
|
||||
device = init_dist(0, 1, _find_free_port())
|
||||
initialize_model_parallel(pipeline_parallel_size=1)
|
||||
|
||||
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
|
||||
test_config, torch.float32, device, do_true_cfg=do_true_cfg
|
||||
)
|
||||
latents = positive_kwargs["x"]
|
||||
t = torch.tensor(500, device=device)
|
||||
|
||||
with torch.inference_mode():
|
||||
noise_pred = pipeline.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=test_config["cfg_scale"],
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
result = pipeline.scheduler_step_maybe_with_cfg(
|
||||
noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
|
||||
)
|
||||
|
||||
destroy_distributed_env()
|
||||
return result.cpu()
|
||||
|
||||
|
||||
def scheduler_step_worker(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
master_port: str,
|
||||
pp_size: int,
|
||||
cfg_size: int,
|
||||
do_true_cfg: bool,
|
||||
test_config: dict,
|
||||
result_queue,
|
||||
):
|
||||
device = init_dist(local_rank, world_size, master_port)
|
||||
initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
|
||||
|
||||
pp_group = get_pp_group()
|
||||
cfg_rank = get_classifier_free_guidance_rank()
|
||||
|
||||
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
|
||||
test_config, torch.float32, device, do_true_cfg=do_true_cfg
|
||||
)
|
||||
latents = positive_kwargs["x"]
|
||||
t = torch.tensor(500, device=device)
|
||||
|
||||
with torch.inference_mode():
|
||||
noise_pred = pipeline.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=test_config["cfg_scale"],
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
latents = pipeline.scheduler_step_maybe_with_cfg(
|
||||
noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
|
||||
)
|
||||
# This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
|
||||
# Flush the last PP rank's async latent send before barrier / process teardown.
|
||||
pipeline._sync_pp_send()
|
||||
|
||||
if pp_group.is_first_rank and cfg_rank == 0:
|
||||
resolved = latents.contiguous()
|
||||
result_queue.put(resolved.cpu())
|
||||
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.barrier()
|
||||
destroy_distributed_env()
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.parametrize(
|
||||
"pp_size, cfg_size, do_true_cfg, input_seed",
|
||||
[
|
||||
pytest.param(
|
||||
2,
|
||||
1,
|
||||
False,
|
||||
300,
|
||||
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
|
||||
id="pp2-no_cfg",
|
||||
),
|
||||
pytest.param(
|
||||
2,
|
||||
2,
|
||||
True,
|
||||
600,
|
||||
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
|
||||
id="pp2-cfg2-true_cfg",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
|
||||
"""Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
|
||||
test_config = {
|
||||
"num_layers": 4,
|
||||
"dim": 64,
|
||||
"batch_size": 2,
|
||||
"cfg_scale": 7.5,
|
||||
"model_seed": 42,
|
||||
"input_seed": input_seed,
|
||||
}
|
||||
|
||||
baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
|
||||
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
q = manager.Queue()
|
||||
|
||||
port = _find_free_port()
|
||||
world_size = pp_size * cfg_size
|
||||
torch.multiprocessing.spawn(
|
||||
scheduler_step_worker,
|
||||
args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
|
||||
nprocs=world_size,
|
||||
)
|
||||
|
||||
result = q.get()
|
||||
torch.testing.assert_close(
|
||||
result,
|
||||
baseline,
|
||||
rtol=0,
|
||||
atol=0,
|
||||
msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
|
||||
)
|
||||
@@ -1,97 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
|
||||
from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
|
||||
Wan22TI2VPipeline,
|
||||
get_wan22_ti2v_pre_process_func,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _make_ti2v_pipeline() -> Wan22TI2VPipeline:
|
||||
pipeline = object.__new__(Wan22TI2VPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
|
||||
pipeline.vae = StubVAE(z_dim=4)
|
||||
pipeline.vae_scale_factor_temporal = 4
|
||||
pipeline.vae_scale_factor_spatial = 8
|
||||
pipeline.progress_bar = noop_progress_bar
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None:
|
||||
preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace())
|
||||
request = SimpleNamespace(
|
||||
prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}],
|
||||
sampling_params=SimpleNamespace(height=None, width=None),
|
||||
)
|
||||
|
||||
result = preprocess(request)
|
||||
|
||||
assert result.sampling_params.height == 672
|
||||
assert result.sampling_params.width == 1344
|
||||
assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672)
|
||||
|
||||
|
||||
def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None:
|
||||
pipeline = _make_ti2v_pipeline()
|
||||
latents = torch.zeros(1, 4, 2, 4, 4)
|
||||
calls = []
|
||||
|
||||
def fake_predict_noise_maybe_with_cfg(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return torch.ones_like(latents)
|
||||
|
||||
pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg # type: ignore[method-assign]
|
||||
pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise # type: ignore[method-assign]
|
||||
|
||||
result = pipeline.diffuse(
|
||||
latents=latents,
|
||||
timesteps=torch.tensor([7]),
|
||||
prompt_embeds=torch.zeros(1, 2, 3),
|
||||
negative_prompt_embeds=torch.zeros(1, 2, 3),
|
||||
guidance_scale=3.0,
|
||||
dtype=torch.float32,
|
||||
attention_kwargs={"a": "b"},
|
||||
num_latent_frames=2,
|
||||
latent_height=4,
|
||||
latent_width=4,
|
||||
)
|
||||
|
||||
positive = calls[0]["positive_kwargs"]
|
||||
assert calls[0]["do_true_cfg"] is True
|
||||
assert positive["timestep"].shape == (1, 8)
|
||||
torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype))
|
||||
torch.testing.assert_close(positive["hidden_states"], latents)
|
||||
torch.testing.assert_close(result, torch.ones_like(latents))
|
||||
|
||||
|
||||
def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None:
|
||||
pipeline = _make_ti2v_pipeline()
|
||||
latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents(
|
||||
image=torch.zeros(1, 3, 16, 16),
|
||||
batch_size=1,
|
||||
num_channels_latents=4,
|
||||
height=16,
|
||||
width=16,
|
||||
num_frames=5,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
generator=None,
|
||||
latents=torch.zeros(1, 4, 2, 2, 2),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2))
|
||||
assert latent_condition.shape == (1, 4, 1, 2, 2)
|
||||
assert first_frame_mask[:, :, 0].sum() == 0
|
||||
assert first_frame_mask[:, :, 1].sum() == 4
|
||||
@@ -1505,7 +1505,6 @@ CUSTOM_DIT_ENABLERS.update(
|
||||
{
|
||||
"Wan22Pipeline": enable_cache_for_wan22,
|
||||
"Wan22I2VPipeline": enable_cache_for_wan22,
|
||||
"Wan22TI2VPipeline": enable_cache_for_wan22,
|
||||
"HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
|
||||
"FluxPipeline": enable_cache_for_flux,
|
||||
"Flux2KleinPipeline": enable_cache_for_flux2_klein,
|
||||
|
||||
@@ -10,6 +10,7 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_cfg_group,
|
||||
@@ -369,7 +370,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
|
||||
results.append(comb)
|
||||
return _unwrap(tuple(results))
|
||||
|
||||
def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]:
|
||||
def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...] | IntermediateTensors:
|
||||
"""
|
||||
Forward pass through transformer to predict noise.
|
||||
|
||||
@@ -381,8 +382,11 @@ class CFGParallelMixin(metaclass=ABCMeta):
|
||||
multi-output models (e.g., video + audio). Multi-output models
|
||||
must also override combine_cfg_noise() and set self.scheduler
|
||||
to a composite scheduler that handles tuples.
|
||||
Non-last Pipeline Parallel stages return `IntermediateTensors`
|
||||
instead of final noise tensors wrapped in a tuple.
|
||||
"""
|
||||
return self.transformer(*args, **kwargs)[0]
|
||||
result = self.transformer(*args, **kwargs)
|
||||
return result if isinstance(result, IntermediateTensors) else result[0]
|
||||
|
||||
def diffuse(
|
||||
self,
|
||||
|
||||
@@ -429,6 +429,78 @@ class GroupCoordinator:
|
||||
async_handle.wait()
|
||||
return tensor_dict
|
||||
|
||||
def isend_tensor_dict(
|
||||
self,
|
||||
tensor_dict: dict[str, torch.Tensor | Any],
|
||||
dst: int | None = None,
|
||||
) -> list[torch.distributed.Work]:
|
||||
"""Non-blocking send of a tensor dictionary.
|
||||
|
||||
Sends metadata via the Gloo CPU group (blocking) then starts a
|
||||
non-blocking NCCL isend for each GPU tensor. Returns the list of
|
||||
Work handles; the caller must call handle.wait() before the tensors
|
||||
can be safely reused or freed.
|
||||
|
||||
NOTE: `dst` is the group rank of the destination.
|
||||
"""
|
||||
if not torch.distributed.is_initialized() or self.world_size == 1:
|
||||
return []
|
||||
|
||||
if dst is None:
|
||||
dst = self.group_next_rank
|
||||
assert dst < self.world_size, f"Invalid dst rank ({dst})"
|
||||
|
||||
metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
|
||||
self.send_object(metadata_list, dst=dst)
|
||||
|
||||
handles: list[torch.distributed.Work] = []
|
||||
for tensor in tensor_list:
|
||||
if tensor.numel() == 0:
|
||||
continue
|
||||
group = self.cpu_group if tensor.is_cpu else self.device_group
|
||||
handle = torch.distributed.isend(tensor, dst=self.ranks[dst], group=group)
|
||||
if tensor.is_cuda:
|
||||
# Keep allocator from reusing this CUDA buffer before the async send finishes.
|
||||
tensor.record_stream(torch.cuda.current_stream(tensor.device))
|
||||
handles.append(handle)
|
||||
return handles
|
||||
|
||||
def irecv_tensor_dict(
|
||||
self,
|
||||
src: int | None = None,
|
||||
) -> tuple[dict[str, torch.Tensor | Any], list[torch.distributed.Work], list]:
|
||||
"""Non-blocking receive of a tensor dictionary.
|
||||
|
||||
Receives metadata via the Gloo CPU group (blocking) then starts a
|
||||
non-blocking NCCL irecv for each GPU tensor. Returns
|
||||
``(tensor_dict, comm_handles, comm_postprocess)`` matching the
|
||||
interface expected by ``AsyncIntermediateTensors``.
|
||||
|
||||
NOTE: `src` is the group rank of the source.
|
||||
"""
|
||||
if not torch.distributed.is_initialized() or self.world_size == 1:
|
||||
return {}, [], []
|
||||
|
||||
if src is None:
|
||||
src = self.group_prev_rank
|
||||
assert src < self.world_size, f"Invalid src rank ({src})"
|
||||
|
||||
recv_metadata_list = self.recv_object(src=src)
|
||||
tensor_dict: dict[str, Any] = {}
|
||||
handles: list[torch.distributed.Work] = []
|
||||
|
||||
for key, value in recv_metadata_list:
|
||||
if isinstance(value, TensorMetadata):
|
||||
tensor = torch.empty(value.size, dtype=value.dtype, device=value.device)
|
||||
if tensor.numel() > 0:
|
||||
group = self.cpu_group if tensor.is_cpu else self.device_group
|
||||
handles.append(torch.distributed.irecv(tensor, src=self.ranks[src], group=group))
|
||||
_update_nested_dict(tensor_dict, key, tensor)
|
||||
else:
|
||||
_update_nested_dict(tensor_dict, key, value)
|
||||
|
||||
return tensor_dict, handles, []
|
||||
|
||||
def send_tensor_dict(
|
||||
self,
|
||||
tensor_dict: dict[str, torch.Tensor | Any],
|
||||
|
||||
@@ -746,11 +746,6 @@ def initialize_model_parallel(
|
||||
f" but got {sequence_parallel_size} != {ring_degree} * {ulysses_degree}"
|
||||
)
|
||||
|
||||
# FIXME: Since the async p2p communication operation of NPU is not same as cuda in torch,
|
||||
# the pipefusion is not ready for npu yet
|
||||
if current_omni_platform.is_npu():
|
||||
assert pipeline_parallel_size == 1, "Current pipefusion is not ready for NPU"
|
||||
|
||||
dit_parallel_size = (
|
||||
data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
|
||||
)
|
||||
|
||||
@@ -0,0 +1,291 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from functools import wraps
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
|
||||
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_cfg_group,
|
||||
get_classifier_free_guidance_rank,
|
||||
get_classifier_free_guidance_world_size,
|
||||
get_pipeline_parallel_world_size,
|
||||
get_pp_group,
|
||||
is_pipeline_first_stage,
|
||||
)
|
||||
|
||||
|
||||
class AsyncLatents:
|
||||
"""Transparent async wrapper returned by scheduler_step on rank 0.
|
||||
|
||||
Wraps a pending ``irecv_tensor_dict`` and defers ``handle.wait()`` until the
|
||||
underlying tensor is actually consumed — either via attribute access
|
||||
(e.g. ``latents.to(dtype)``, ``latents.shape``) or via a torch operation
|
||||
(e.g. ``mask * latents``). This keeps the first PP rank non-blocking after
|
||||
posting the receive, matching the async philosophy used everywhere else in
|
||||
the PP communication layer.
|
||||
"""
|
||||
|
||||
__slots__ = ("_tensor_dict", "_handles", "_postproc", "_tensor")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tensor_dict: dict[str, torch.Tensor],
|
||||
handles: list[torch.distributed.Work],
|
||||
postproc: list,
|
||||
):
|
||||
self._tensor_dict = tensor_dict
|
||||
self._handles = handles
|
||||
self._postproc = postproc
|
||||
self._tensor: torch.Tensor | None = None
|
||||
|
||||
def _resolve(self) -> torch.Tensor:
|
||||
if self._tensor is not None:
|
||||
return self._tensor
|
||||
for h in self._handles:
|
||||
h.wait()
|
||||
for fn in self._postproc:
|
||||
fn()
|
||||
self._tensor = self._tensor_dict["latents"]
|
||||
return self._tensor
|
||||
|
||||
# Attribute access (e.g. .shape, .to(), .dtype) delegates to the resolved tensor.
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(self._resolve(), name)
|
||||
|
||||
# Torch function protocol: any torch op involving an AsyncLatents resolves it first.
|
||||
@classmethod
|
||||
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
||||
kwargs = kwargs or {}
|
||||
|
||||
def _unwrap(x):
|
||||
if isinstance(x, AsyncLatents):
|
||||
return x._resolve()
|
||||
if isinstance(x, (list, tuple)):
|
||||
return type(x)(_unwrap(item) for item in x) # type(x) return the class of x to preserve its type
|
||||
return x
|
||||
|
||||
args = tuple(_unwrap(a) for a in args)
|
||||
kwargs = {k: _unwrap(v) for k, v in kwargs.items()}
|
||||
return func(*args, **kwargs)
|
||||
|
||||
|
||||
class PipelineParallelMixin:
|
||||
"""
|
||||
Mixin providing Pipeline Parallelism for diffusion pipelines.
|
||||
|
||||
All PP ranks run the full denoising loop in `forward()`.
|
||||
`predict_noise_maybe_with_cfg` and `scheduler_step_maybe_with_cfg` encapsulate all inter-rank communication.
|
||||
|
||||
Communication pattern per denoising step:
|
||||
Forward chain : rank 0 → 1 → … → N-1 via async isend/irecv (AsyncIntermediateTensors)
|
||||
Next timestep : last rank → rank 0 via async isend/irecv (AsyncLatents)
|
||||
|
||||
All communication is asynchronous using isend_tensor_dict/irecv_tensor_dict.
|
||||
Only rank 0 needs updated latents for the next forward pass start.
|
||||
|
||||
For sequential CFG (cfg_parallel_size=1) with PP, two full forward chains are
|
||||
executed — one for the positive pass and one for the negative pass — so that each
|
||||
PP stage operates on the correct encoder_hidden_states.
|
||||
"""
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
|
||||
if not issubclass(cls, CFGParallelMixin):
|
||||
raise TypeError(
|
||||
f"{cls.__name__} inherits PipelineParallelMixin but not CFGParallelMixin. "
|
||||
"Pipeline Parallelism requires CFGParallelMixin for predict_noise(), "
|
||||
"predict_noise_maybe_with_cfg(), scheduler_step_maybe_with_cfg(), and combine_cfg_noise(). "
|
||||
"Add CFGParallelMixin to the base classes of your pipeline."
|
||||
)
|
||||
mro = cls.mro()
|
||||
if mro.index(PipelineParallelMixin) > mro.index(CFGParallelMixin):
|
||||
raise TypeError(
|
||||
f"{cls.__name__} must inherit PipelineParallelMixin before CFGParallelMixin so MRO selects "
|
||||
f"PP-aware predict/scheduler wrappers and their `super()` calls delegate to CFGParallelMixin."
|
||||
)
|
||||
|
||||
init = cls.__dict__.get("__init__")
|
||||
if callable(init):
|
||||
|
||||
@wraps(init)
|
||||
def wrapped_init(self, *args: Any, **kwargs: Any) -> None:
|
||||
init(self, *args, **kwargs)
|
||||
vae = getattr(self, "vae", None)
|
||||
if vae is not None and hasattr(vae, "decode"):
|
||||
self._wrapped_vae_decode()
|
||||
|
||||
cls.__init__ = wrapped_init
|
||||
|
||||
diffuse = cls.__dict__.get("diffuse")
|
||||
if callable(diffuse):
|
||||
|
||||
@wraps(diffuse)
|
||||
def wrapped_diffuse(self, *args: Any, **kwargs: Any) -> Any:
|
||||
try:
|
||||
latents = diffuse(self, *args, **kwargs)
|
||||
if isinstance(latents, AsyncLatents):
|
||||
latents = torch.as_tensor(latents) # avoid copying
|
||||
return latents
|
||||
finally:
|
||||
self._sync_pp_send()
|
||||
|
||||
cls.diffuse = wrapped_diffuse
|
||||
|
||||
def _wrapped_vae_decode(self) -> None:
|
||||
vae, orig_decode = self.vae, self.vae.decode
|
||||
|
||||
@wraps(orig_decode)
|
||||
def wrapped_decode(z: torch.Tensor, *args: Any, **kwargs: Any):
|
||||
if hasattr(vae, "is_distributed_enabled") and vae.is_distributed_enabled():
|
||||
# Middle ranks (world size 3 or more) hold stale latents after the denoising loop.
|
||||
# Broadcast from rank 0 so every rank splits identical tiles.
|
||||
if get_pipeline_parallel_world_size() > 2:
|
||||
z = get_pp_group().broadcast(z, src=0)
|
||||
return orig_decode(z, *args, **kwargs)
|
||||
elif is_pipeline_first_stage():
|
||||
return orig_decode(z, *args, **kwargs)
|
||||
return (None,) # decoder returns a tuple
|
||||
|
||||
self.vae.decode = wrapped_decode
|
||||
|
||||
@property
|
||||
def _pp_send_work(self) -> list[torch.distributed.Work]:
|
||||
if not hasattr(self, "_pp_send_work_list"):
|
||||
self._pp_send_work_list: list[torch.distributed.Work] = []
|
||||
return self._pp_send_work_list
|
||||
|
||||
@_pp_send_work.setter
|
||||
def _pp_send_work(self, work: list[torch.distributed.Work]) -> None:
|
||||
self._pp_send_work_list = work
|
||||
|
||||
def _sync_pp_send(self) -> None:
|
||||
"""
|
||||
Wait on all pending non-blocking PP sends.
|
||||
|
||||
Must be called after the denoising loop so that the isend handles
|
||||
from the last iteration are completed before any subsequent
|
||||
collective (e.g. VAE decode broadcast) or tensor reuse.
|
||||
"""
|
||||
if self._pp_send_work:
|
||||
for handle in self._pp_send_work:
|
||||
handle.wait()
|
||||
self._pp_send_work = []
|
||||
|
||||
def predict_noise_maybe_with_cfg(
|
||||
self,
|
||||
do_true_cfg: bool,
|
||||
true_cfg_scale: float,
|
||||
positive_kwargs: dict[str, Any],
|
||||
negative_kwargs: dict[str, Any] | None,
|
||||
cfg_normalize: bool = True,
|
||||
output_slice: int | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, ...] | None:
|
||||
"""
|
||||
Drop-in replacement for predict_noise_maybe_with_cfg that also handles PP.
|
||||
|
||||
Supports three modes:
|
||||
- PP only, sequential CFG: both branches (cond and uncond) run through this PP pipeline.
|
||||
This doubles communication volume per denoising step compared to PP + CFG-parallel.
|
||||
- PP + CFG-parallel: each PP pipeline carries one branch. The last PP
|
||||
rank all-gathers across the CFG group and combines, mirroring
|
||||
CFGParallelMixin.predict_noise_maybe_with_cfg exactly.
|
||||
- PP only, no CFG: cond branch only.
|
||||
|
||||
Returns:
|
||||
noise_pred on the last PP rank (all CFG ranks when CFG-parallel is active).
|
||||
None on all other ranks.
|
||||
"""
|
||||
if get_pipeline_parallel_world_size() == 1:
|
||||
return super().predict_noise_maybe_with_cfg(
|
||||
do_true_cfg, true_cfg_scale, positive_kwargs, negative_kwargs, cfg_normalize, output_slice
|
||||
)
|
||||
|
||||
self._sync_pp_send()
|
||||
|
||||
pp_group = get_pp_group()
|
||||
|
||||
cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
|
||||
if cfg_parallel_ready:
|
||||
# Each PP pipeline carries exactly one CFG branch determined by cfg_rank.
|
||||
all_kwargs = [positive_kwargs if get_classifier_free_guidance_rank() == 0 else negative_kwargs]
|
||||
else:
|
||||
# Sequential CFG (or no CFG): this PP pipeline handles all branches.
|
||||
all_kwargs = [positive_kwargs] + ([negative_kwargs] if do_true_cfg else [])
|
||||
|
||||
# Non-first ranks receive intermediate tensors asynchronously
|
||||
n = len(all_kwargs)
|
||||
its: list[AsyncIntermediateTensors | None] = [None] * n
|
||||
if not pp_group.is_first_rank:
|
||||
for i in range(n):
|
||||
its[i] = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
|
||||
|
||||
if not pp_group.is_last_rank:
|
||||
# First / middle rank: run partial forwards and propagate ITs downstream.
|
||||
for kwargs, it in zip(all_kwargs, its):
|
||||
result = self.predict_noise(**kwargs, intermediate_tensors=it)
|
||||
self._pp_send_work.extend(pp_group.isend_tensor_dict(result.tensors))
|
||||
return None
|
||||
|
||||
# Last rank: run full forward
|
||||
noise_preds = [self.predict_noise(**kwargs, intermediate_tensors=it) for kwargs, it in zip(all_kwargs, its)]
|
||||
|
||||
if cfg_parallel_ready:
|
||||
# All-gather the single-branch prediction across the CFG group and combine
|
||||
# on all CFG ranks so every last PP rank has an identical noise_pred.
|
||||
local_pred = noise_preds[0]
|
||||
if output_slice is not None:
|
||||
local_pred = local_pred[:, :output_slice]
|
||||
gathered = get_cfg_group().all_gather(local_pred, separate_tensors=True)
|
||||
return self.combine_cfg_noise(gathered[0], gathered[1], true_cfg_scale, cfg_normalize)
|
||||
|
||||
# Sequential CFG or no-CFG path.
|
||||
if do_true_cfg:
|
||||
pos, neg = noise_preds[0], noise_preds[1]
|
||||
if output_slice is not None:
|
||||
pos = pos[:, :output_slice]
|
||||
neg = neg[:, :output_slice]
|
||||
return self.combine_cfg_noise(pos, neg, true_cfg_scale, cfg_normalize)
|
||||
pred = noise_preds[0]
|
||||
if output_slice is not None:
|
||||
pred = pred[:, :output_slice]
|
||||
return pred
|
||||
|
||||
def scheduler_step_maybe_with_cfg(
|
||||
self,
|
||||
noise_pred: torch.Tensor | tuple[torch.Tensor, ...] | None,
|
||||
t: torch.Tensor | tuple[torch.Tensor, ...],
|
||||
latents: torch.Tensor | tuple[torch.Tensor, ...],
|
||||
do_true_cfg: bool,
|
||||
per_request_scheduler: Any | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, ...] | AsyncLatents:
|
||||
"""
|
||||
Drop-in replacement for scheduler_step_maybe_with_cfg that also handles PP.
|
||||
|
||||
Only the last rank runs the scheduler (it already has noise_pred); the result
|
||||
is sent to rank 0 which needs it for the next forward pass.
|
||||
|
||||
Returns a ``AsyncLatents`` on rank 0 that transparently defers
|
||||
``handle.wait()`` until the tensor is actually consumed (via attribute
|
||||
access or a torch operation), keeping the rank non-blocking after the
|
||||
``irecv`` is posted.
|
||||
"""
|
||||
if get_pipeline_parallel_world_size() == 1:
|
||||
return super().scheduler_step_maybe_with_cfg(
|
||||
noise_pred, t, latents, do_true_cfg, per_request_scheduler, generator
|
||||
)
|
||||
|
||||
pp_group = get_pp_group()
|
||||
if pp_group.is_last_rank:
|
||||
latents = super().scheduler_step_maybe_with_cfg(
|
||||
noise_pred, t, latents, do_true_cfg, per_request_scheduler, generator
|
||||
)
|
||||
self._pp_send_work = pp_group.isend_tensor_dict({"latents": latents}, dst=0)
|
||||
elif pp_group.is_first_rank:
|
||||
latents = AsyncLatents(*pp_group.irecv_tensor_dict(src=pp_group.world_size - 1))
|
||||
return latents
|
||||
@@ -22,11 +22,6 @@ from .pipeline_wan2_2_s2v import (
|
||||
get_wan22_s2v_post_process_func,
|
||||
get_wan22_s2v_pre_process_func,
|
||||
)
|
||||
from .pipeline_wan2_2_ti2v import (
|
||||
Wan22TI2VPipeline,
|
||||
get_wan22_ti2v_post_process_func,
|
||||
get_wan22_ti2v_pre_process_func,
|
||||
)
|
||||
from .pipeline_wan2_2_vace import (
|
||||
Wan22VACEPipeline,
|
||||
get_wan22_vace_post_process_func,
|
||||
@@ -50,9 +45,6 @@ __all__ = [
|
||||
"Wan22S2VPipeline",
|
||||
"get_wan22_s2v_post_process_func",
|
||||
"get_wan22_s2v_pre_process_func",
|
||||
"Wan22TI2VPipeline",
|
||||
"get_wan22_ti2v_post_process_func",
|
||||
"get_wan22_ti2v_pre_process_func",
|
||||
"Wan22VACEPipeline",
|
||||
"get_wan22_vace_post_process_func",
|
||||
"get_wan22_vace_pre_process_func",
|
||||
|
||||
@@ -17,10 +17,12 @@ from torch import nn
|
||||
from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -297,7 +299,9 @@ def get_wan22_pre_process_func(
|
||||
return pre_process_func
|
||||
|
||||
|
||||
class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin):
|
||||
class Wan22Pipeline(
|
||||
nn.Module, PipelineParallelMixin, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
|
||||
):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
@@ -460,7 +464,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
|
||||
attention_kwargs: dict[str, Any],
|
||||
latent_condition: torch.Tensor | None = None,
|
||||
first_frame_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
) -> torch.Tensor | AsyncLatents:
|
||||
if attention_kwargs is None:
|
||||
attention_kwargs = {}
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
@@ -858,7 +862,11 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
|
||||
output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
|
||||
)
|
||||
|
||||
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
|
||||
def predict_noise(
|
||||
self,
|
||||
current_model: nn.Module | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
"""
|
||||
Forward pass through transformer to predict noise.
|
||||
|
||||
@@ -867,11 +875,12 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
|
||||
**kwargs: Arguments to pass to the transformer
|
||||
|
||||
Returns:
|
||||
Predicted noise tensor
|
||||
Predicted noise tensor or IntermediateTensors on non-last PP stages.
|
||||
"""
|
||||
if current_model is None:
|
||||
current_model = self.transformer
|
||||
return current_model(**kwargs)[0]
|
||||
result = current_model(**kwargs)
|
||||
return result if isinstance(result, IntermediateTensors) else result[0]
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
|
||||
@@ -17,10 +17,12 @@ from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
@@ -135,7 +137,12 @@ def get_wan22_i2v_pre_process_func(
|
||||
|
||||
|
||||
class Wan22I2VPipeline(
|
||||
nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
|
||||
nn.Module,
|
||||
SupportImageInput,
|
||||
PipelineParallelMixin,
|
||||
CFGParallelMixin,
|
||||
ProgressBarMixin,
|
||||
DiffusionPipelineProfilerMixin,
|
||||
):
|
||||
"""
|
||||
Wan2.2 Image-to-Video Pipeline.
|
||||
@@ -283,7 +290,7 @@ class Wan22I2VPipeline(
|
||||
attention_kwargs: dict[str, Any],
|
||||
condition: torch.Tensor,
|
||||
first_frame_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
) -> torch.Tensor | AsyncLatents:
|
||||
if attention_kwargs is None:
|
||||
attention_kwargs = {}
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
@@ -348,7 +355,6 @@ class Wan22I2VPipeline(
|
||||
|
||||
# Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync
|
||||
latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
|
||||
|
||||
pbar.update()
|
||||
|
||||
return latents
|
||||
@@ -667,7 +673,11 @@ class Wan22I2VPipeline(
|
||||
output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
|
||||
)
|
||||
|
||||
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
|
||||
def predict_noise(
|
||||
self,
|
||||
current_model: nn.Module | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
"""
|
||||
Forward pass through transformer to predict noise.
|
||||
|
||||
@@ -676,11 +686,12 @@ class Wan22I2VPipeline(
|
||||
**kwargs: Arguments to pass to the transformer
|
||||
|
||||
Returns:
|
||||
Predicted noise tensor
|
||||
Predicted noise tensor or IntermediateTensors on non-last PP stages.
|
||||
"""
|
||||
if current_model is None:
|
||||
current_model = self.transformer
|
||||
return current_model(**kwargs)[0]
|
||||
result = current_model(**kwargs)
|
||||
return result if isinstance(result, IntermediateTensors) else result[0]
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
|
||||
@@ -1,731 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Wan2.2 TI2V (Text-Image-to-Video) Pipeline.
|
||||
|
||||
This pipeline supports the unified TI2V-5B model that can generate videos from:
|
||||
- Text only (T2V mode)
|
||||
- Text + Image (I2V mode)
|
||||
|
||||
The key difference from the MoE-based I2V pipeline is:
|
||||
- Single transformer (not MoE with two transformers)
|
||||
- Uses expand_timesteps mode for image conditioning
|
||||
- No CLIP image encoder - only VAE encoding for image condition
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, cast
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
|
||||
from vllm_omni.diffusion.models.interface import SupportImageInput
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
||||
build_wan_scheduler,
|
||||
create_transformer_from_config,
|
||||
load_transformer_config,
|
||||
resolve_wan_flow_shift,
|
||||
resolve_wan_sample_solver,
|
||||
retrieve_latents,
|
||||
)
|
||||
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
|
||||
from vllm_omni.diffusion.postprocess import interpolate_video_tensor
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniTextPrompt
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_wan22_ti2v_post_process_func(
|
||||
od_config: OmniDiffusionConfig,
|
||||
):
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
|
||||
video_processor = VideoProcessor(vae_scale_factor=8)
|
||||
|
||||
def post_process_func(
|
||||
video: torch.Tensor,
|
||||
output_type: str = "np",
|
||||
sampling_params=None,
|
||||
):
|
||||
if output_type == "latent":
|
||||
return video
|
||||
custom_output = {}
|
||||
if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
|
||||
video, multiplier = interpolate_video_tensor(
|
||||
video,
|
||||
exp=sampling_params.frame_interpolation_exp,
|
||||
scale=sampling_params.frame_interpolation_scale,
|
||||
model_path=sampling_params.frame_interpolation_model_path,
|
||||
)
|
||||
custom_output["video_fps_multiplier"] = multiplier
|
||||
return {
|
||||
"video": video_processor.postprocess_video(video, output_type=output_type),
|
||||
"custom_output": custom_output,
|
||||
}
|
||||
|
||||
return post_process_func
|
||||
|
||||
|
||||
def get_wan22_ti2v_pre_process_func(
|
||||
od_config: OmniDiffusionConfig,
|
||||
):
|
||||
"""Pre-process function for TI2V: optionally load and resize input image."""
|
||||
|
||||
def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
|
||||
for i, prompt in enumerate(request.prompts):
|
||||
multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
|
||||
raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
|
||||
if isinstance(prompt, str):
|
||||
prompt = OmniTextPrompt(prompt=prompt)
|
||||
if "additional_information" not in prompt:
|
||||
prompt["additional_information"] = {}
|
||||
|
||||
if raw_image is None:
|
||||
raise ValueError(
|
||||
"""No image is provided. This model requires an image to run.""",
|
||||
"""Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
|
||||
)
|
||||
if not isinstance(raw_image, (str, PIL.Image.Image)):
|
||||
raise TypeError(
|
||||
f"""Unsupported image format {raw_image.__class__}.""",
|
||||
"""Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
|
||||
)
|
||||
image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image
|
||||
|
||||
# Calculate dimensions based on aspect ratio if not provided
|
||||
if request.sampling_params.height is None or request.sampling_params.width is None:
|
||||
# Default max area for 720P (TI2V-5B default)
|
||||
max_area = 720 * 1280
|
||||
aspect_ratio = image.height / image.width
|
||||
|
||||
# Calculate dimensions maintaining aspect ratio
|
||||
mod_value = 16 # Must be divisible by 16
|
||||
height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
|
||||
width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
|
||||
|
||||
if request.sampling_params.height is None:
|
||||
request.sampling_params.height = height
|
||||
if request.sampling_params.width is None:
|
||||
request.sampling_params.width = width
|
||||
|
||||
# Resize image to target dimensions
|
||||
image = image.resize(
|
||||
(request.sampling_params.width, request.sampling_params.height), # type: ignore # Above has ensured that width & height are not None
|
||||
PIL.Image.Resampling.LANCZOS,
|
||||
)
|
||||
prompt["multi_modal_data"]["image"] = image # type: ignore # key existence already checked above
|
||||
|
||||
request.prompts[i] = prompt
|
||||
return request
|
||||
|
||||
return pre_process_func
|
||||
|
||||
|
||||
class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin):
|
||||
"""
|
||||
Wan2.2 Text-Image-to-Video (TI2V) Pipeline.
|
||||
|
||||
This is a unified pipeline that supports both:
|
||||
- Text-to-Video (T2V): when no image is provided
|
||||
- Image-to-Video (I2V): when an image is provided
|
||||
|
||||
Uses expand_timesteps mode for I2V conditioning where the first frame
|
||||
is conditioned on the input image latent.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.od_config = od_config
|
||||
|
||||
self.device = get_local_device()
|
||||
dtype = getattr(od_config, "dtype", torch.bfloat16)
|
||||
|
||||
model = od_config.model
|
||||
local_files_only = os.path.exists(model)
|
||||
|
||||
# Set up weights sources for single transformer
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model,
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
fall_back_to_pt=True,
|
||||
),
|
||||
]
|
||||
|
||||
# See ``hub_prefetch.py`` for the transformers v5 subfolder race.
|
||||
prefetch_subfolders(
|
||||
model,
|
||||
["tokenizer", "text_encoder", "vae"],
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
|
||||
# Text encoder
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
|
||||
self.text_encoder = UMT5EncoderModel.from_pretrained(
|
||||
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
# VAE
|
||||
self.vae = OmniAutoencoderKLWan.from_pretrained(
|
||||
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
|
||||
).to(self.device)
|
||||
|
||||
# Single transformer (TI2V uses dense 5B model, not MoE)
|
||||
# Load config from model to get correct dimensions
|
||||
transformer_config = load_transformer_config(model, "transformer", local_files_only)
|
||||
self.transformer = self._create_transformer(transformer_config)
|
||||
|
||||
self._sample_solver = "unipc"
|
||||
self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0
|
||||
self.scheduler = build_wan_scheduler(self._sample_solver, self._flow_shift)
|
||||
|
||||
# VAE scale factors
|
||||
self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if hasattr(self.vae, "config") else 4
|
||||
self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if hasattr(self.vae, "config") else 8
|
||||
|
||||
# TI2V always uses expand_timesteps mode
|
||||
self.expand_timesteps = True
|
||||
|
||||
self._guidance_scale = None
|
||||
self._num_timesteps = None
|
||||
self._current_timestep = None
|
||||
|
||||
def _create_transformer(self, config: dict) -> WanTransformer3DModel:
|
||||
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
|
||||
quant_config = getattr(self.od_config, "quantization_config", None)
|
||||
return create_transformer_from_config(config, quant_config=quant_config)
|
||||
|
||||
@property
|
||||
def guidance_scale(self):
|
||||
return self._guidance_scale
|
||||
|
||||
@property
|
||||
def do_classifier_free_guidance(self):
|
||||
return self._guidance_scale is not None and self._guidance_scale > 1.0
|
||||
|
||||
@property
|
||||
def num_timesteps(self):
|
||||
return self._num_timesteps
|
||||
|
||||
@property
|
||||
def current_timestep(self):
|
||||
return self._current_timestep
|
||||
|
||||
def diffuse(
|
||||
self,
|
||||
latents: torch.Tensor,
|
||||
timesteps: torch.Tensor,
|
||||
prompt_embeds: torch.Tensor,
|
||||
negative_prompt_embeds: torch.Tensor | None,
|
||||
guidance_scale: float,
|
||||
dtype: torch.dtype,
|
||||
attention_kwargs: dict[str, Any],
|
||||
num_latent_frames: int,
|
||||
latent_height: int,
|
||||
latent_width: int,
|
||||
latent_condition: torch.Tensor | None = None,
|
||||
first_frame_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if attention_kwargs is None:
|
||||
attention_kwargs = {}
|
||||
with self.progress_bar(total=len(timesteps)) as pbar:
|
||||
for step_idx, t in enumerate(timesteps):
|
||||
self._current_timestep = t
|
||||
set_forward_context_denoise_step_idx(step_idx)
|
||||
|
||||
# Prepare latent input
|
||||
if latent_condition is not None:
|
||||
# I2V mode: blend condition with latents using mask
|
||||
latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents
|
||||
latent_model_input = latent_model_input.to(dtype)
|
||||
|
||||
# Expand timesteps for each patch (TI2V style)
|
||||
temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten()
|
||||
timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1)
|
||||
else:
|
||||
# T2V mode: use latents directly
|
||||
latent_model_input = latents.to(dtype)
|
||||
|
||||
# Expand timesteps for TI2V model architecture
|
||||
mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, device=latents.device)
|
||||
temp_ts = (mask[0][0][:, ::2, ::2] * t).flatten()
|
||||
timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1)
|
||||
|
||||
do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None
|
||||
positive_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"timestep": timestep,
|
||||
"encoder_hidden_states": prompt_embeds,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
"current_model": self.transformer,
|
||||
}
|
||||
if do_true_cfg:
|
||||
negative_kwargs = {
|
||||
"hidden_states": latent_model_input,
|
||||
"timestep": timestep,
|
||||
"encoder_hidden_states": negative_prompt_embeds,
|
||||
"attention_kwargs": attention_kwargs,
|
||||
"return_dict": False,
|
||||
"current_model": self.transformer,
|
||||
}
|
||||
else:
|
||||
negative_kwargs = None
|
||||
|
||||
noise_pred = self.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=do_true_cfg,
|
||||
true_cfg_scale=guidance_scale,
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
|
||||
latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
|
||||
pbar.update()
|
||||
|
||||
return latents
|
||||
|
||||
def forward(
|
||||
self,
|
||||
req: OmniDiffusionRequest,
|
||||
prompt: str | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
image: PIL.Image.Image | torch.Tensor | None = None,
|
||||
height: int = 704,
|
||||
width: int = 1280,
|
||||
num_inference_steps: int = 40,
|
||||
guidance_scale: float = 5.0,
|
||||
frame_num: int = 81,
|
||||
output_type: str | None = "np",
|
||||
generator: torch.Generator | list[torch.Generator] | None = None,
|
||||
prompt_embeds: torch.Tensor | None = None,
|
||||
negative_prompt_embeds: torch.Tensor | None = None,
|
||||
attention_kwargs: dict | None = None,
|
||||
**kwargs,
|
||||
) -> DiffusionOutput:
|
||||
# Get parameters from request or arguments
|
||||
if len(req.prompts) > 1:
|
||||
raise ValueError(
|
||||
"""This model only supports a single prompt, not a batched request.""",
|
||||
"""Please pass in a single prompt object or string, or a single-item list.""",
|
||||
)
|
||||
if len(req.prompts) == 1: # If req.prompt is empty, default to prompt & neg_prompt in param list
|
||||
prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
|
||||
negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
|
||||
if prompt is None and prompt_embeds is None:
|
||||
raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
|
||||
|
||||
# Get image from request (optional for TI2V)
|
||||
if image is None:
|
||||
multi_modal_data = (
|
||||
req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
|
||||
)
|
||||
raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
|
||||
if isinstance(raw_image, list):
|
||||
if len(raw_image) > 1:
|
||||
logger.warning(
|
||||
"""Received a list of image. Only a single image is supported by this model."""
|
||||
"""Taking only the first image for now."""
|
||||
)
|
||||
raw_image = raw_image[0]
|
||||
if raw_image is None:
|
||||
image = None
|
||||
elif isinstance(raw_image, str):
|
||||
image = PIL.Image.open(raw_image)
|
||||
else:
|
||||
image = cast(PIL.Image.Image | torch.Tensor, raw_image)
|
||||
|
||||
# Default dimensions for TI2V-5B (720P)
|
||||
height = req.sampling_params.height or height
|
||||
width = req.sampling_params.width or width
|
||||
num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
|
||||
num_steps = req.sampling_params.num_inference_steps or num_inference_steps
|
||||
|
||||
# Respect per-request guidance_scale when explicitly provided.
|
||||
if req.sampling_params.guidance_scale_provided:
|
||||
guidance_scale = req.sampling_params.guidance_scale
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
|
||||
# Validate inputs
|
||||
self.check_inputs(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
image=image,
|
||||
height=height,
|
||||
width=width,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
)
|
||||
|
||||
# Adjust num_frames to be compatible with VAE temporal scaling
|
||||
if num_frames % self.vae_scale_factor_temporal != 1:
|
||||
num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
|
||||
num_frames = max(num_frames, 1)
|
||||
|
||||
device = self.device
|
||||
dtype = self.transformer.dtype
|
||||
|
||||
# Generator setup
|
||||
if generator is None:
|
||||
generator = req.sampling_params.generator
|
||||
if generator is None and req.sampling_params.seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
|
||||
|
||||
# Encode prompts
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
do_classifier_free_guidance=guidance_scale > 1.0,
|
||||
num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
|
||||
max_sequence_length=req.sampling_params.max_sequence_length or 512,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
else:
|
||||
prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
|
||||
if negative_prompt_embeds is not None:
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
|
||||
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver)
|
||||
flow_shift = resolve_wan_flow_shift(req, self.od_config)
|
||||
if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
|
||||
self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
|
||||
self._sample_solver = sample_solver
|
||||
self._flow_shift = flow_shift
|
||||
|
||||
# Timesteps
|
||||
self.scheduler.set_timesteps(num_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
# Prepare latents
|
||||
num_channels_latents = self.transformer.config.in_channels
|
||||
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
latent_height = height // self.vae_scale_factor_spatial
|
||||
latent_width = width // self.vae_scale_factor_spatial
|
||||
|
||||
# Check if we have an image (I2V mode) or not (T2V mode)
|
||||
if image is not None:
|
||||
# I2V mode: prepare latents with image condition
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
|
||||
video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
|
||||
|
||||
if isinstance(image, PIL.Image.Image):
|
||||
image_tensor = video_processor.preprocess(image, height=height, width=width)
|
||||
else:
|
||||
image_tensor = image
|
||||
image_tensor = image_tensor.to(device=device, dtype=torch.float32)
|
||||
|
||||
latents, latent_condition, first_frame_mask = self.prepare_i2v_latents(
|
||||
image=image_tensor,
|
||||
batch_size=batch_size,
|
||||
num_channels_latents=num_channels_latents,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
generator=generator,
|
||||
latents=req.sampling_params.latents,
|
||||
)
|
||||
else:
|
||||
# T2V mode: prepare random latents
|
||||
latents = self.prepare_latents(
|
||||
batch_size=batch_size,
|
||||
num_channels_latents=num_channels_latents,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
generator=generator,
|
||||
latents=req.sampling_params.latents,
|
||||
)
|
||||
latent_condition = None
|
||||
first_frame_mask = torch.ones(
|
||||
1, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device
|
||||
)
|
||||
|
||||
if attention_kwargs is None:
|
||||
attention_kwargs = {}
|
||||
|
||||
latents = self.diffuse(
|
||||
latents=latents,
|
||||
timesteps=timesteps,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
guidance_scale=guidance_scale,
|
||||
dtype=dtype,
|
||||
attention_kwargs=attention_kwargs,
|
||||
num_latent_frames=num_latent_frames,
|
||||
latent_height=latent_height,
|
||||
latent_width=latent_width,
|
||||
latent_condition=latent_condition,
|
||||
first_frame_mask=first_frame_mask,
|
||||
)
|
||||
|
||||
# Wan2.2 is prone to out of memory errors when predicting large videos
|
||||
# so we empty the cache here to avoid OOM before vae decoding.
|
||||
if current_omni_platform.is_available():
|
||||
current_omni_platform.empty_cache()
|
||||
self._current_timestep = None
|
||||
|
||||
# For I2V mode, blend final latents with condition
|
||||
if latent_condition is not None:
|
||||
latents = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents
|
||||
|
||||
# Decode
|
||||
if output_type == "latent":
|
||||
output = latents
|
||||
else:
|
||||
latents = latents.to(self.vae.dtype)
|
||||
latents_mean = (
|
||||
torch.tensor(self.vae.config.latents_mean)
|
||||
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
||||
.to(latents.device, latents.dtype)
|
||||
)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
||||
latents.device, latents.dtype
|
||||
)
|
||||
latents = latents / latents_std + latents_mean
|
||||
output = self.vae.decode(latents, return_dict=False)[0]
|
||||
|
||||
return DiffusionOutput(output=output)
|
||||
|
||||
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass through transformer to predict noise.
|
||||
|
||||
Args:
|
||||
current_model: The transformer model to use
|
||||
**kwargs: Arguments to pass to the transformer
|
||||
|
||||
Returns:
|
||||
Predicted noise tensor
|
||||
"""
|
||||
if current_model is None:
|
||||
current_model = self.transformer
|
||||
return current_model(**kwargs)[0]
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
negative_prompt: str | list[str] | None = None,
|
||||
do_classifier_free_guidance: bool = True,
|
||||
num_videos_per_prompt: int = 1,
|
||||
max_sequence_length: int = 512,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
"""Encode text prompts using T5 text encoder."""
|
||||
device = device or self.device
|
||||
dtype = dtype or self.text_encoder.dtype
|
||||
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
prompt_clean = [self._prompt_clean(p) for p in prompt]
|
||||
batch_size = len(prompt_clean)
|
||||
|
||||
text_inputs = self.tokenizer(
|
||||
prompt_clean,
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_attention_mask=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
ids, mask = text_inputs.input_ids, text_inputs.attention_mask
|
||||
seq_lens = mask.gt(0).sum(dim=1).long()
|
||||
|
||||
prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state
|
||||
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
|
||||
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
|
||||
prompt_embeds = torch.stack(
|
||||
[torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
|
||||
)
|
||||
|
||||
_, seq_len, _ = prompt_embeds.shape
|
||||
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
||||
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
||||
|
||||
negative_prompt_embeds = None
|
||||
if do_classifier_free_guidance:
|
||||
negative_prompt = negative_prompt or ""
|
||||
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
|
||||
neg_text_inputs = self.tokenizer(
|
||||
[self._prompt_clean(p) for p in negative_prompt],
|
||||
padding="max_length",
|
||||
max_length=max_sequence_length,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_attention_mask=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask
|
||||
seq_lens_neg = mask_neg.gt(0).sum(dim=1).long()
|
||||
negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
|
||||
negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)]
|
||||
negative_prompt_embeds = torch.stack(
|
||||
[
|
||||
torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
|
||||
for u in negative_prompt_embeds
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
||||
negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
||||
|
||||
return prompt_embeds, negative_prompt_embeds
|
||||
|
||||
@staticmethod
|
||||
def _prompt_clean(text: str) -> str:
|
||||
return " ".join(text.strip().split())
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size: int,
|
||||
num_channels_latents: int,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
dtype: torch.dtype | None,
|
||||
device: torch.device | None,
|
||||
generator: torch.Generator | list[torch.Generator] | None,
|
||||
latents: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Prepare random latents for T2V mode."""
|
||||
if latents is not None:
|
||||
return latents.to(device=device, dtype=dtype)
|
||||
|
||||
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
num_latent_frames,
|
||||
int(height) // self.vae_scale_factor_spatial,
|
||||
int(width) // self.vae_scale_factor_spatial,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
return latents
|
||||
|
||||
def prepare_i2v_latents(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
batch_size: int,
|
||||
num_channels_latents: int,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
dtype: torch.dtype | None,
|
||||
device: torch.device | None,
|
||||
generator: torch.Generator | list[torch.Generator] | None,
|
||||
latents: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Prepare latents for I2V mode with image conditioning.
|
||||
|
||||
Returns:
|
||||
latents: Initial noise latents
|
||||
latent_condition: Encoded first frame condition
|
||||
first_frame_mask: Mask (0 for first frame, 1 for rest)
|
||||
"""
|
||||
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
latent_height = height // self.vae_scale_factor_spatial
|
||||
latent_width = width // self.vae_scale_factor_spatial
|
||||
|
||||
shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device=device, dtype=dtype)
|
||||
|
||||
# Prepare first frame condition
|
||||
image = image.unsqueeze(2) # [batch, channels, 1, height, width]
|
||||
image = image.to(device=device, dtype=self.vae.dtype)
|
||||
|
||||
# Encode through VAE
|
||||
latent_condition = retrieve_latents(self.vae.encode(image), sample_mode="argmax")
|
||||
latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
|
||||
|
||||
# Normalize latents
|
||||
latents_mean = (
|
||||
torch.tensor(self.vae.config.latents_mean)
|
||||
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
||||
.to(latent_condition.device, latent_condition.dtype)
|
||||
)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
||||
latent_condition.device, latent_condition.dtype
|
||||
)
|
||||
latent_condition = (latent_condition - latents_mean) * latents_std
|
||||
latent_condition = latent_condition.to(dtype)
|
||||
|
||||
# Create mask: 0 for first frame (condition), 1 for rest (to denoise)
|
||||
first_frame_mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device)
|
||||
first_frame_mask[:, :, 0] = 0
|
||||
|
||||
return latents, latent_condition, first_frame_mask
|
||||
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
image,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
):
|
||||
if height % 16 != 0 or width % 16 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
|
||||
|
||||
if prompt is not None and prompt_embeds is not None:
|
||||
raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Please provide only one.")
|
||||
|
||||
if negative_prompt is not None and negative_prompt_embeds is not None:
|
||||
raise ValueError(
|
||||
"Cannot forward both `negative_prompt` and `negative_prompt_embeds`. Please provide only one."
|
||||
)
|
||||
|
||||
if prompt is None and prompt_embeds is None:
|
||||
raise ValueError("Provide either `prompt` or `prompt_embeds`.")
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""Load weights using AutoWeightsLoader for vLLM integration."""
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights)
|
||||
@@ -21,9 +21,21 @@ from vllm.model_executor.layers.conv import Conv3dLayer
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
from vllm.model_executor.models.utils import (
|
||||
PPMissingLayer,
|
||||
is_pp_missing_parameter,
|
||||
make_empty_intermediate_tensors_factory,
|
||||
make_layers,
|
||||
)
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
get_pipeline_parallel_world_size,
|
||||
is_pipeline_first_stage,
|
||||
is_pipeline_last_stage,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.sp_plan import (
|
||||
SequenceParallelInput,
|
||||
SequenceParallelOutput,
|
||||
@@ -871,14 +883,23 @@ class WanTransformer3DModel(nn.Module):
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
out_channels = out_channels or in_channels
|
||||
|
||||
if get_pipeline_parallel_world_size() == 0 and not is_pipeline_first_stage():
|
||||
raise RuntimeError(
|
||||
"`initialize_model_parallel()` must be called before constructing `WanTransformer3DModel`"
|
||||
)
|
||||
|
||||
# 1. Patch & position embedding
|
||||
self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
|
||||
self.patch_embedding = Conv3dLayer(
|
||||
in_channels=in_channels,
|
||||
out_channels=inner_dim,
|
||||
kernel_size=patch_size,
|
||||
stride=patch_size,
|
||||
)
|
||||
# Patch embedding only on first PP stage; other stages receive hidden_states via P2P
|
||||
if is_pipeline_first_stage():
|
||||
self.patch_embedding = Conv3dLayer(
|
||||
in_channels=in_channels,
|
||||
out_channels=inner_dim,
|
||||
kernel_size=patch_size,
|
||||
stride=patch_size,
|
||||
)
|
||||
else:
|
||||
self.patch_embedding = PPMissingLayer()
|
||||
|
||||
# 2. Condition embeddings
|
||||
self.condition_embedder = WanTimeTextImageEmbedding(
|
||||
@@ -890,30 +911,39 @@ class WanTransformer3DModel(nn.Module):
|
||||
pos_embed_seq_len=pos_embed_seq_len,
|
||||
)
|
||||
|
||||
# 3. Transformer blocks
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
WanTransformerBlock(
|
||||
inner_dim,
|
||||
ffn_dim,
|
||||
num_attention_heads,
|
||||
eps,
|
||||
added_kv_proj_dim,
|
||||
cross_attn_norm,
|
||||
quant_config=quant_config,
|
||||
prefix=f"blocks.{layer_idx}",
|
||||
)
|
||||
for layer_idx in range(num_layers)
|
||||
]
|
||||
# 3. Transformer blocks — partitioned across PP stages via vLLM's `make_layers`.
|
||||
# It computes the [start_layer, end_layer) slice for this rank and fills the remaining slots
|
||||
# with PPMissingLayer so that weight names stay globally consistent.
|
||||
self.start_layer, self.end_layer, self.blocks = make_layers(
|
||||
num_layers,
|
||||
lambda prefix: WanTransformerBlock(
|
||||
inner_dim,
|
||||
ffn_dim,
|
||||
num_attention_heads,
|
||||
eps,
|
||||
added_kv_proj_dim,
|
||||
cross_attn_norm,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
),
|
||||
prefix="blocks",
|
||||
)
|
||||
|
||||
# 4. Output norm & projection
|
||||
self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
|
||||
self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
|
||||
# 4. Output norm & projection — only on the last PP stage
|
||||
if is_pipeline_last_stage():
|
||||
self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
|
||||
self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
|
||||
else:
|
||||
self.norm_out, self.proj_out = PPMissingLayer(), PPMissingLayer()
|
||||
|
||||
# SP helper modules
|
||||
self.timestep_proj_prepare = TimestepProjPrepare()
|
||||
self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
|
||||
if is_pipeline_last_stage():
|
||||
self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
|
||||
else:
|
||||
self.output_scale_shift_prepare = PPMissingLayer()
|
||||
|
||||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], inner_dim)
|
||||
|
||||
# ROPE helper
|
||||
self._cached_rope_emb = None
|
||||
@@ -930,9 +960,10 @@ class WanTransformer3DModel(nn.Module):
|
||||
timestep: torch.LongTensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
encoder_hidden_states_image: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
return_dict: bool = True,
|
||||
attention_kwargs: dict[str, Any] | None = None,
|
||||
) -> torch.Tensor | Transformer2DModelOutput:
|
||||
) -> torch.Tensor | Transformer2DModelOutput | IntermediateTensors:
|
||||
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
||||
p_t, p_h, p_w = self.config.patch_size
|
||||
post_patch_num_frames = num_frames // p_t
|
||||
@@ -949,18 +980,25 @@ class WanTransformer3DModel(nn.Module):
|
||||
self._hidden_states_shape = hidden_states.shape
|
||||
self._cached_rope_emb = rotary_emb
|
||||
|
||||
# Patch embedding and flatten to sequence
|
||||
# (hidden_states is sharded at blocks.0 input by _sp_plan)
|
||||
hidden_states = self.patch_embedding(hidden_states)
|
||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
||||
if is_pipeline_first_stage():
|
||||
# Patch embedding and flatten to sequence
|
||||
# (hidden_states is sharded at blocks.0 input by _sp_plan)
|
||||
hidden_states = self.patch_embedding(hidden_states)
|
||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
||||
else:
|
||||
if intermediate_tensors is None:
|
||||
raise RuntimeError("intermediate_tensors must be provided for non-first PP stages")
|
||||
hidden_states = intermediate_tensors["hidden_states"]
|
||||
|
||||
# Handle timestep shape
|
||||
# Handle timestep shape (2-D for TI2V, 1-D for T2V)
|
||||
if timestep.ndim == 2:
|
||||
ts_seq_len = timestep.shape[1]
|
||||
timestep = timestep.flatten()
|
||||
else:
|
||||
ts_seq_len = None
|
||||
|
||||
# Compute conditioning on all PP stages.
|
||||
# Each stage needs temb/timestep_proj for scale-shift modulation in its local blocks.
|
||||
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
|
||||
timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len
|
||||
)
|
||||
@@ -996,14 +1034,13 @@ class WanTransformer3DModel(nn.Module):
|
||||
hidden_states_mask = None
|
||||
|
||||
# Transformer blocks
|
||||
for block in self.blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
timestep_proj,
|
||||
rotary_emb,
|
||||
hidden_states_mask,
|
||||
)
|
||||
for block in self.blocks[self.start_layer : self.end_layer]:
|
||||
hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask)
|
||||
|
||||
if not is_pipeline_last_stage():
|
||||
# Non-last PP stage: hand the token sequence to the caller via IntermediateTensors.
|
||||
# predict_noise will broadcast it to the next stage before calling that stage's forward.
|
||||
return IntermediateTensors({"hidden_states": hidden_states})
|
||||
|
||||
# Output norm, projection & unpatchify
|
||||
shift, scale = self.output_scale_shift_prepare(temb)
|
||||
@@ -1074,7 +1111,8 @@ class WanTransformer3DModel(nn.Module):
|
||||
if weight_name not in original_name:
|
||||
continue
|
||||
lookup_name = original_name.replace(weight_name, param_name)
|
||||
if lookup_name not in params_dict:
|
||||
# Skip weights that belong to PP stages other than this one
|
||||
if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict:
|
||||
break
|
||||
param = params_dict[lookup_name]
|
||||
weight_loader = param.weight_loader
|
||||
@@ -1098,6 +1136,10 @@ class WanTransformer3DModel(nn.Module):
|
||||
if modulation_alias in params_dict:
|
||||
lookup_name = modulation_alias
|
||||
|
||||
# Skip weights that belong to PP stages other than this one
|
||||
if is_pp_missing_parameter(lookup_name, self):
|
||||
continue
|
||||
|
||||
if lookup_name not in params_dict:
|
||||
logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
|
||||
continue
|
||||
|
||||
@@ -1632,6 +1632,7 @@ class AsyncOmniEngine:
|
||||
data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1
|
||||
tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1
|
||||
cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
|
||||
pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
|
||||
vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
|
||||
enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
|
||||
use_hsdp = normalized_kwargs.get("use_hsdp", False)
|
||||
|
||||
Reference in new issue
Block a user