xDiT

xDiT adalah mesin inferensi yang dirancang untuk deployment paralel berskala besar dari Diffusion Transformers (DiTs). xDiT menyediakan serangkaian pendekatan paralel yang efisien untuk Diffusion Models, serta akselerasi kernel GPU.

Ada empat metode paralel yang didukung dalam xDiT, termasuk Unified Sequence Parallelism, PipeFusion, paralelisme CFG, dan paralelisme data. Keempat metode paralel dalam xDiT dapat dikonfigurasi secara hibrida, mengoptimalkan pola komunikasi agar paling sesuai dengan perangkat keras jaringan yang mendasarinya.

Optimasi yang ortogonal terhadap paralelisasi berfokus pada percepatan kinerja GPU tunggal. Selain memanfaatkan pustaka optimasi Attention yang sudah dikenal luas, kami memanfaatkan teknologi percepatan kompilasi seperti torch.compile dan onediff.

Ikhtisar xDiT ditunjukkan sebagai berikut.

pip install xfuser

Berikut adalah contoh penggunaan xDiT untuk mempercepat inferensi model Diffusers.

 import torch
 from diffusers import StableDiffusion3Pipeline

 from xfuser import xFuserArgs, xDiTParallel
 from xfuser.config import FlexibleArgumentParser
 from xfuser.core.distributed import get_world_group

 def main():
+    parser = FlexibleArgumentParser(description="xFuser Arguments")
+    args = xFuserArgs.add_cli_args(parser).parse_args()
+    engine_args = xFuserArgs.from_cli_args(args)
+    engine_config, input_config = engine_args.create_config()

     local_rank = get_world_group().local_rank
     pipe = StableDiffusion3Pipeline.from_pretrained(
         pretrained_model_name_or_path=engine_config.model_config.model,
         dtype=torch.float16,
     ).to(f"cuda:{local_rank}")
    
# do anything you want with pipeline here

+    pipe = xDiTParallel(pipe, engine_config, input_config)

     pipe(
         height=input_config.height,
         width=input_config.height,
         prompt=input_config.prompt,
         num_inference_steps=input_config.num_inference_steps,
         output_type=input_config.output_type,
         generator=torch.Generator(device="cuda").manual_seed(input_config.seed),  # or "mps", "xpu", "cpu"
     )

+    if input_config.output_type == "pil":
+        pipe.save("results", "stable_diffusion_3")

if __name__ == "__main__":
    main()

Seperti yang dapat Anda lihat, kita hanya perlu menggunakan xFuserArgs dari xDiT untuk mendapatkan parameter konfigurasi, lalu meneruskan parameter tersebut bersama objek pipeline dari pustaka Diffusers ke xDiTParallel untuk menyelesaikan paralelisasi pipeline tertentu di Diffusers.

Parameter runtime xDiT dapat dilihat di command line menggunakan -h, dan Anda dapat merujuk pada contoh penggunaan ini untuk detail lebih lanjut.

xDiT harus dijalankan menggunakan torchrun untuk mendukung kemampuan paralel multi-node dan multi-GPU-nya. Sebagai contoh, perintah berikut dapat digunakan untuk inferensi paralel 8-GPU:

torchrun --nproc_per_node=8 ./inference.py --model models/FLUX.1-dev --data_parallel_degree 2 --ulysses_degree 2 --ring_degree 2 --prompt "A snowy mountain" "A small dog" --num_inference_steps 50

Model yang didukung

Sebagian model Diffusers didukung di xDiT, seperti Flux.1, Stable Diffusion 3, dan lain-lain. Model-model terbaru yang didukung dapat ditemukan di sini.

Benchmark

Kami menguji berbagai model di berbagai mesin, dan berikut adalah sebagian data benchmark-nya.

Flux.1-schnell

Stable Diffusion 3

HunyuanDiT

Metrik kinerja yang lebih detail dapat ditemukan di halaman github kami.

Referensi

xDiT-project

USP: Pendekatan Parallelism Urutan Terpadu untuk AI Generatif Konteks Panjang

PipeFusion: Paralelisme Pipeline Patch yang Dipindahkan untuk Inferensi Model Diffusion Transformer