xDiT

xDiT（PyPI 包名 xfuser）是一套统一的并行推理框架。本示例在单卡昇腾 NPU 上生成第一张图。

前置条件

硬件

Atlas 900 A2 训练服务器（Ascend 910B），并按需完成物理机或容器内的设备挂载。单卡生成示例需 1 张卡，序列并行示例需 2 张卡。

基础软件

在运行本文档示例之前，你的机器上需要已经装好并可用：

本文档示例在 Python 3.12、CANN 9.1.0 环境下验证通过。

加载 CANN 环境

source /usr/local/Ascend/ascend-toolkit/set_env.sh

安装 PyTorch NPU 栈

torch、torch_npu 与 triton 三者版本严格配套，按 CANN 与 PyTorch 配套表 选择与 CANN 匹配的组合：

pip install torch==2.9.0 torch_npu==2.9.0.post6 triton==3.5.0

安装 xDiT

安装 xfuser（PyPI 包名），并打印安装版本：

pip install xfuser
python -c "from importlib.metadata import version; print('xDiT version:', version('xfuser'))"

输出结果如下：

...
xDiT version: xxx

其中 xxx 是安装的 xDiT（xfuser）版本号。

运行示例：文生图

安装模型下载所需的 ModelScope：

pip install "modelscope==1.37.0"

用 SD3 medium 在单卡上生成一张 256×256 的图。模型约 28 GB。

先把下面这个脚本写入 sd3_npu.py：

from pathlib import Path

script = """
import os
import sys
import torch
import torch_npu
from modelscope import snapshot_download
from transformers import T5EncoderModel
from xfuser import xFuserArgs, xFuserStableDiffusion3Pipeline
from xfuser.config import FlexibleArgumentParser
from xfuser.core.distributed import get_runtime_state, get_world_group

model_path = snapshot_download('stabilityai/stable-diffusion-3-medium-diffusers')

parser = FlexibleArgumentParser(description="xFuser SD3 Arguments")
args = xFuserArgs.add_cli_args(parser).parse_args(['--model', model_path] + sys.argv[1:])
engine_args = xFuserArgs.from_cli_args(args)
engine_config, input_config = engine_args.create_config()
local_rank = get_world_group().rank

text_encoder_3 = T5EncoderModel.from_pretrained(
    model_path, subfolder="text_encoder_3", dtype=torch.float16
)
pipe = xFuserStableDiffusion3Pipeline.from_pretrained(
    pretrained_model_name_or_path=model_path,
    engine_config=engine_config,
    dtype=torch.float16,
    text_encoder_3=text_encoder_3,
).to(f"npu:{local_rank}")
pipe.prepare_run(input_config)

output = pipe(
    height=input_config.height,
    width=input_config.width,
    prompt=input_config.prompt,
    num_inference_steps=input_config.num_inference_steps,
    output_type=input_config.output_type,
    guidance_scale=input_config.guidance_scale,
    generator=torch.Generator(device="npu").manual_seed(input_config.seed),
)
os.makedirs("results", exist_ok=True)
if pipe.is_dp_last_group():
    output.images[0].save("results/sd3_npu.png")
    print("saved: results/sd3_npu.png")
get_runtime_state().destroy_distributed_env()
"""

Path("sd3_npu.py").write_text(script.lstrip("\n"), encoding="utf-8")

脚本已写入 sd3_npu.py。用 torchrun 在单卡上运行：

torchrun --nproc_per_node=1 sd3_npu.py --prompt "a tiny test sketch" --height 256 --width 256 --num_inference_steps 1 --seed 42

输出结果如下：

...
saved: results/sd3_npu.png

多卡运行示例

同一个脚本、同一个模型，加 --ulysses_degree 2 在 2 卡上做序列并行，attention 用 SDPA 后端：

torchrun --nproc_per_node=2 sd3_npu.py --prompt "a tiny test sketch" --height 256 --width 256 --num_inference_steps 1 --seed 42 --ulysses_degree 2 --attention_backend SDPA

输出结果如下：

...
saved: results/sd3_npu.png

更多用法

更多模型与多卡并行（PipeFusion / CFG 并行 / Ring 等）见 xDiT examples。