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
• Step-by-Step Implementation
• Testing
• Troubleshooting
• 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.

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:

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:

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:

# 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:

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.

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:

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:

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: