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()

확인할 수 있듯이, xDiT의 xFuserArgs만 사용하여 구성 매개변수를 가져오고, 이 매개변수와 Diffusers 라이브러리의 파이프라인 객체를 xDiTParallel에 전달하면 Diffusers의 특정 파이프라인 병렬화를 완료할 수 있습니다.

xDiT 런타임 매개변수는 -h를 사용하여 명령줄에서 확인할 수 있으며, 자세한 내용은 이 사용법 예시를 참고하세요.

xDiT는 멀티 노드, 멀티 GPU 병렬 기능을 지원하기 위해 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