diffusers/v0.43.0/zh/optimization/xdit.md

xDiT

xDiT 是一款专为大规模并行部署 Diffusion Transformers（DiTs）而设计的推理引擎。xDiT 提供了一套面向 Diffusion Models 的高效并行方案，以及 GPU 内核加速功能。

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

与并行化正交的优化侧重于提升单 GPU 性能。除了利用知名的 Attention 优化库外，我们还借助 torch.compile 和 onediff 等编译加速技术。

xDiT 的概览如下所示。

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

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 库的 pipeline 对象一起传入 xDiTParallel，即可完成 Diffusers 中特定 pipeline 的并行化。

可以通过在命令行中使用 -h 来查看 xDiT 的运行时参数，有关更多详细信息，请参阅此 用法 示例。

xDiT 需要使用 torchrun 启动，以支持其多节点、多 GPU 并行能力。例如，以下命令可用于 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

支持的模型

xDiT 支持部分 Diffusers 模型，例如 Flux.1、Stable Diffusion 3 等。最新支持的模型列表请参见此处。

基准测试

我们在不同的机器上测试了不同的模型，以下是一些基准测试数据。

Flux.1-schnell

Stable Diffusion 3

HunyuanDiT

更详细的性能指标可在我们的 github 页面上找到。

参考

xDiT-project

USP：面向长上下文生成式 AI 的统一序列并行方法

PipeFusion：用于扩散 Transformer 模型推理的位移补丁流水线并行