- π₯ Meet xDiT
- π’ Updates
- π― Supported DiTs
- π Performance
- π QuickStart
- β¨ the xDiT's secret weapons
- π Develop Guide
- π§ History and Looking for Contributions
- π Cite Us
Diffusion Transformers (DiTs), pivotal in text-to-image and text-to-video models, are driving advancements in high-quality image and video generation. With the escalating input sequence length in DiTs, the computational demand of the Attention mechanism grows quadratically! Consequently, multi-GPU and multi-machine deployments are essential to maintain real-time performance in online services.
To meet real-time demand for DiTs applications, parallel inference is a must. xDiT is an inference engine designed for the parallel deployment of DiTs on large scale. xDiT provides a suite of efficient parallel inference approaches for Diffusion Models.
-
Sequence Parallelism, USP is a unified sequence parallel approach combining DeepSpeed-Ulysses, Ring-Attention.
-
PipeFusion, a patch level pipeline parallelism using displaced patch by taking advantage of the diffusion model characteristics.
-
Data Parallel: Processes multiple prompts or generates multiple images from a single prompt in parallel across images.
-
CFG Parallel, also known as Split Batch: Activates when using classifier-free guidance (CFG) with a constant parallelism of 2.
The four parallel methods in xDiT can be configured in a hybrid manner, optimizing communication patterns to best suit the underlying network hardware.
As shown in the following picture, xDiT offers a set of APIs to adapt DiT models in huggingface/diffusers to hybrid parallel implementation through simple wrappers. If the model you require is not available in the model zoo, developing it yourself is straightforward; please refer to our Dev Guide.
We also have implemented the following parallel stategies for reference:
- Tensor Parallelism
- DistriFusion
The communication and memory costs associated with the aforementioned parallelism, except for the CFG and DP, in DiTs are detailed in the table below. (* denotes that communication can be overlapped with computation.)
As we can see, PipeFusion and Sequence Parallel achieve lowest communication cost on different scales and hardware configurations, making them suitable foundational components for a hybrid approach.
π: Number of pixels; ππ: Model hidden size; π³: Number of model layers; π·: Total model parameters; π΅: Number of parallel devices; π΄: Number of patch splits; πΈπΆ: Query and Output parameter count; π²π½: KV Activation parameter count; π¨ = πΈ = πΆ = π² = π½: Equal parameters for Attention, Query, Output, Key, and Value;
| attn-KV | communication cost | param memory | activations memory | extra buff memory | |
|---|---|---|---|---|---|
| Tensor Parallel | fresh | ||||
| DistriFusion* | stale | ||||
| Ring Sequence Parallel* | fresh | ||||
| Ulysses Sequence Parallel | fresh | ||||
| PipeFusion* | stale- |
- πAugust 9, 2024: Support Latte sequence parallel version. The inference scripts are examples/latte_example.
- πAugust 8, 2024: Support Flux sequence parallel version. The inference scripts are examples/flux_example.
- πAugust 2, 2024: Support Stable Diffusion 3 hybrid parallel version. The inference scripts are examples/sd3_example.
- πJuly 18, 2024: Support PixArt-Sigma and PixArt-Alpha. The inference scripts are examples/pixartsigma_example.py, examples/pixartalpha_example.py.
- πJuly 17, 2024: Rename the project to xDiT. The project has evolved from a collection of parallel methods into a unified inference framework and supported the hybrid parallel for DiTs.
- πJuly 10, 2024: Support HunyuanDiT. The inference script is legacy/scripts/hunyuandit_example.py.
- πJune 26, 2024: Support Stable Diffusion 3. The inference script is legacy/scripts/sd3_example.py.
- πMay 24, 2024: PipeFusion is public released. It supports PixArt-alpha legacy/scripts/pixart_example.py, DiT legacy/scripts/ditxl_example.py and SDXL legacy/scripts/sdxl_example.py.
| Model Name | CFG | SP | PipeFusion |
|---|---|---|---|
| π¬ Latte | β | βοΈ | β |
| π΅ HunyuanDiT-v1.2-Diffusers | βοΈ | βοΈ | βοΈ |
| π Flux | NA | βοΈ | β |
| π΄ PixArt-Sigma | βοΈ | βοΈ | βοΈ |
| π’ PixArt-alpha | βοΈ | βοΈ | βοΈ |
| π Stable Diffusion 3 | βοΈ | βοΈ | βοΈ |
Here are the benchmark results for Pixart-Alpha using the 20-step DPM solver as the scheduler across various image resolutions. To replicate these findings, please refer to the script at ./legacy/scripts/benchmark.sh.
TBD: Updates results on hybrid parallelism.
- The Latency on 4xA100-80GB (PCIe)
- The Latency on 8xL20-48GB (PCIe)
- The Latency on 8xA100-80GB (NVLink)
- The Latency on 4xT4-16GB (PCIe)
1. Install from pip (current version)
pip install xfuser
Install yunchang from feifeibear/long-context-attention.
Please note that it has a dependency on flash attention and specific GPU model requirements. We recommend installing yunchang from the source code rather than using pip install yunchang==0.2.0.
python setup.py install
We provide examples demonstrating how to run models with xDiT in the ./examples/ directory. You can easily modify the model type, model directory, and parallel options in the examples/run.sh within the script to run some already supported DiT models.
bash examples/run.shTo inspect the available options for the PixArt-alpha example, use the following command:
python ./examples/pixartalpha_example.py -h
...
xFuser Arguments
options:
-h, --help show this help message and exit
Model Options:
--model MODEL Name or path of the huggingface model to use.
--download-dir DOWNLOAD_DIR
Directory to download and load the weights, default to the default cache dir of huggingface.
--trust-remote-code Trust remote code from huggingface.
Runtime Options:
--warmup_steps WARMUP_STEPS
Warmup steps in generation.
--use_parallel_vae
--seed SEED Random seed for operations.
--output_type OUTPUT_TYPE
Output type of the pipeline.
Parallel Processing Options:
--do_classifier_free_guidance
--use_split_batch Use split batch in classifier_free_guidance. cfg_degree will be 2 if set
--data_parallel_degree DATA_PARALLEL_DEGREE
Data parallel degree.
--ulysses_degree ULYSSES_DEGREE
Ulysses sequence parallel degree. Used in attention layer.
--ring_degree RING_DEGREE
Ring sequence parallel degree. Used in attention layer.
--pipefusion_parallel_degree PIPEFUSION_PARALLEL_DEGREE
Pipefusion parallel degree. Indicates the number of pipeline stages.
--num_pipeline_patch NUM_PIPELINE_PATCH
Number of patches the feature map should be segmented in pipefusion parallel.
--attn_layer_num_for_pp [ATTN_LAYER_NUM_FOR_PP ...]
List representing the number of layers per stage of the pipeline in pipefusion parallel
--tensor_parallel_degree TENSOR_PARALLEL_DEGREE
Tensor parallel degree.
--split_scheme SPLIT_SCHEME
Split scheme for tensor parallel.
Input Options:
--height HEIGHT The height of image
--width WIDTH The width of image
--prompt [PROMPT ...]
Prompt for the model.
--no_use_resolution_binning
--negative_prompt [NEGATIVE_PROMPT ...]
Negative prompt for the model.
--num_inference_steps NUM_INFERENCE_STEPS
Number of inference steps.Hybriding multiple parallelism techniques togather is essential for efficiently scaling. It's important that the product of all parallel degrees matches the number of devices. For instance, you can combine CFG, PipeFusion, and sequence parallelism with the command below to generate an image of a cute dog through hybrid parallelism. Here ulysses_degree * pipefusion_parallel_degree * cfg_degree(use_split_batch) == number of devices == 8.
torchrun --nproc_per_node=8 \
examples/pixartalpha_example.py \
--model models/PixArt-XL-2-1024-MS \
--pipefusion_parallel_degree 2 \
--ulysses_degree 2 \
--num_inference_steps 20 \
--warmup_steps 0 \
--prompt "A small dog" \
--use_cfg_parallelwarmup_steps, also required in DistriFusion, typically set to a small number compared with num_inference_steps.
The warmup step impacts the efficiency of PipeFusion as it cannot be executed in parallel, thus degrading to a serial execution.
We observed that a warmup of 0 had no effect on the PixArt model.
Users can tune this value according to their specific tasks.
The exceptional capabilities of xDiT stem from our innovative technologies.
PipeFusion: Displaced Patch Pipeline Parallelism for Diffusion Models
USP: A Unified Sequence Parallelism Approach for Long Context Generative AI
The implement and design of xdit framework
We conducted a major upgrade of this project in August 2024.
The latest APIs is located in the xfuser/ directory, supports hybrid parallelism. It offers clearer and more structured code but currently supports fewer models.
The legacy APIs is in the legacy/ directory, limited to single parallelism. It supports a richer of parallel methods, including PipeFusion, Sequence Parallel, DistriFusion, and Tensor Parallel. CFG Parallel can be hybrid with PipeFusion but not with other parallel methods.
For models not yet supported by the latest APIs, you can run the examples in the legacy/scripts/ directory. If you wish to develop new features on a model or require hybrid parallelism, stay tuned for further project updates.
We also welcome developers to join and contribute more features and models to the project. Tell us which model you need in xDiT in discussions.
@article{wang2024pipefusion,
title={PipeFusion: Displaced Patch Pipeline Parallelism for Inference of Diffusion Transformer Models},
author={Jiannan Wang and Jiarui Fang and Jinzhe Pan and Aoyu Li and PengCheng Yang},
year={2024},
eprint={2405.07719},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
@article{fang2024unified,
title={USP: a Unified Sequence Parallelism Approach for Long Context Generative AI},
author={Fang, Jiarui and Zhao, Shangchun},
journal={arXiv preprint arXiv:2405.07719},
year={2024}
}