xDiT

xDiT 是一款为大规模并行部署扩散 Transformer (DiT) 而设计的推理引擎。xDiT 提供了一套针对扩散模型的高效并行方法，以及 GPU 内核加速。

xDiT 支持四种并行方法，包括统一序列并行、PipeFusion、CFG 并行和数据并行。xDiT 中的四种并行方法可以以混合方式配置，优化通信模式以最适合底层网络硬件。

与并行化正交的优化侧重于加速单个 GPU 的性能。除了利用众所周知的注意力优化库之外，我们还利用编译加速技术，例如 torch.compile 和 onediff。

xDiT 的概述如下所示。

你可以使用以下命令安装 xDiT：

bash

pip install xfuser

以下是使用 xDiT 加速 Diffusers 模型推理的示例。

diff

 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,
         torch_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),
     )

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

if __name__ == "__main__":
    main()

如你所见，我们只需要使用 xDiT 中的 xFuserArgs 来获取配置参数，并将这些参数与 Diffusers 库中的管道对象一起传递给 xDiTParallel，即可完成 Diffusers 中特定管道的并行化。

你可以在命令行中使用 -h 查看 xDiT 运行时参数，并参考此使用示例以获取更多详细信息。

xDiT 需要使用 torchrun 启动才能支持其多节点、多 GPU 并行功能。例如，以下命令可用于 8 GPU 并行推理：

bash

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