Diffusers

加入 Hugging Face 社区

并获得增强的文档体验

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

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

切换文档主题

开始使用

xDiT

xDiT 是一个专为大规模并行部署扩散变换器（Diffusion Transformers, DiTs）而设计的推理引擎。xDiT 提供了一套针对扩散模型的高效并行方法，以及 GPU 内核加速功能。

xDiT 支持四种并行方法，包括统一序列并行（Unified Sequence Parallelism）、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 库中的管道对象一起传递给 xDiTParallel，即可完成 Diffusers 中特定管道的并行化。

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