xDiT

xDiT là một engine suy luận được thiết kế cho việc triển khai song song quy mô lớn của Diffusion Transformers (DiTs). xDiT cung cấp một bộ các phương pháp song song hiệu quả cho Diffusion Models, cũng như các tăng tốc kernel GPU.

Có bốn phương pháp song song được hỗ trợ trong xDiT, bao gồm Unified Sequence Parallelism, PipeFusion, song song CFG và song song dữ liệu. Bốn phương pháp song song trong xDiT có thể được cấu hình theo cách lai, tối ưu hóa các mẫu giao tiếp để phù hợp nhất với phần cứng mạng nền tảng.

Tối ưu hóa độc lập với song song hóa tập trung vào việc tăng tốc hiệu năng của một GPU đơn. Ngoài việc sử dụng các thư viện tối ưu hóa Attention phổ biến, chúng tôi tận dụng các công nghệ tăng tốc biên dịch như torch.compile và onediff.

Tổng quan về xDiT được trình bày như sau.

pip install xfuser

Dưới đây là một ví dụ về việc sử dụng xDiT để tăng tốc suy luận của một mô hình 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()

Như bạn có thể thấy, chúng ta chỉ cần sử dụng xFuserArgs từ xDiT để lấy các tham số cấu hình, sau đó truyền các tham số này cùng với đối tượng pipeline từ thư viện Diffusers vào xDiTParallel để hoàn tất việc song song hóa một pipeline cụ thể trong Diffusers.

Các tham số runtime của xDiT có thể được xem trên dòng lệnh bằng cách sử dụng -h, và bạn có thể tham khảo ví dụ sử dụng này để biết thêm chi tiết.

xDiT cần được khởi chạy bằng torchrun để hỗ trợ khả năng song song đa nút, đa GPU. Ví dụ, lệnh sau có thể được sử dụng cho suy luận song song 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

Các mô hình được hỗ trợ

Một tập con các mô hình Diffusers được hỗ trợ trong xDiT, chẳng hạn như Flux.1, Stable Diffusion 3, v.v. Các mô hình được hỗ trợ mới nhất có thể được tìm thấy tại đây.

Đánh giá hiệu năng

Chúng tôi đã kiểm tra các mô hình khác nhau trên nhiều máy khác nhau, và đây là một số dữ liệu benchmark.

Flux.1-schnell

Stable Diffusion 3

HunyuanDiT

Các chỉ số hiệu suất chi tiết hơn có thể được tìm thấy trên trang github của chúng tôi.

Tham khảo

xDiT-project

USP: Một phương pháp song song hóa chuỗi thống nhất cho AI tạo sinh ngữ cảnh dài

PipeFusion: Song song hóa pipeline theo từng khối bị dịch chuyển cho suy luận các mô hình Diffusion Transformer