Diffusers

加入 Hugging Face 社区

并获得增强的文档体验

在模型、数据集和 Spaces 上协作

通过加速推理获得更快的示例

切换文档主题

开始使用

xDiT

xDiT 是一个推理引擎，专为大规模并行部署 Diffusion Transformers (DiTs) 而设计。xDiT 为扩散模型提供了一套高效的并行方法，以及 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,
         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 库中的 pipeline 对象传递到 xDiTParallel 中，即可完成 Diffusers 中特定 pipeline 的并行化。

xDiT 运行时参数可以使用命令行 -h 查看，您可以参考此用法示例了解更多详情。

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