xDiT

xDiT — это движок инференса, разработанный для крупномасштабного параллельного развёртывания Diffusion Transformers (DiTs). xDiT предлагает набор эффективных параллельных подходов для Diffusion Models, а также ускорения на уровне GPU-ядер.

В xDiT поддерживаются четыре параллельных метода, включая Unified Sequence Parallelism, PipeFusion, параллелизм CFG и параллелизм данных. Четыре параллельных метода в xDiT можно настроить в гибридном режиме, оптимизируя паттерны коммуникации для наилучшего соответствия базовому сетевому оборудованию.

Оптимизация, ортогональная параллелизации, направлена на ускорение производительности одного GPU. Помимо использования известных библиотек оптимизации Attention, мы применяем технологии компиляционного ускорения, такие как torch.compile и onediff.

Обзор xDiT представлен ниже.

You can install xDiT using the following command:

pip install xfuser

Вот пример использования xDiT для ускорения инференса модели 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()

Как видно, для получения параметров конфигурации достаточно использовать xFuserArgs из xDiT, а затем передать эти параметры вместе с объектом пайплайна из библиотеки Diffusers в xDiTParallel, чтобы завершить параллелизацию конкретного пайплайна в Diffusers.

Параметры времени выполнения xDiT можно просмотреть в командной строке с помощью -h, а для получения подробной информации обратитесь к этому примеру использования.

Для поддержки многоузловой и много-GPU параллельной работы xDiT необходимо запускать с помощью torchrun. Например, для параллельного вывода на 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