xDiT

xDiTは、Diffusion Transformers (DiTs) の大規模並列デプロイメント用に設計された推論エンジンです。xDiTは、Diffusion Models向けの効率的な並列アプローチのスイートと、GPUカーネルのアクセラレーションを提供します。

xDiTでは、Unified Sequence Parallelism、PipeFusion、CFG並列化、データ並列化の4つの並列化手法がサポートされています。xDiTの4つの並列化手法はハイブリッドに設定でき、基盤となるネットワークハードウェアに最も適した通信パターンを最適化できます。

並列化とは直交する最適化は、単一 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を指定して表示でき、詳細についてはこの使用例を参照してください。

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