DiffSynth-Studio GPU/NPU 多硬件支持实战：从 NVIDIA/AMD/Apple Silicon 到昇腾 NPU 的推理与训练适配指南

本文是 DiffSynth-Studio 硬件加速指南，系统讲解如何在 NVIDIA、AMD、Apple Silicon 与昇腾（Ascend）NPU 上运行扩散模型的推理与训练。你将掌握各硬件平台的安装要点、"cuda" 到 "npu" 的设备切换方法、NPU 训练脚本中的性能环境变量与专属参数，以及 USP 序列并行等高级特性的启用方式，可直接将项目内任意示例迁移到目标设备上。

硬件支持总览

DiffSynth-Studio 的示例代码默认面向 NVIDIA GPU 编写，设备字符串统一为 "cuda"。项目在 diffsynth/core/device/npu_compatible_device.py 中提供了一层设备抽象，用于自动判断当前机器的计算设备类型：

• get_device_type()：按 cuda → npu → cpu 的优先级返回设备类型；
• get_torch_device()：返回 torch.cuda / torch.npu 等命名空间对象；
• get_device_name()：返回 "cuda:0"、"npu:0" 这类带设备编号的字符串；
• get_nccl_backend() / parse_nccl_backend()：在 CUDA 上返回 nccl，在 NPU 上返回 hccl，用于分布式通信后端选择。

依据 npu_compatible_device.py 中 is_torch_npu_available() 的实现，NPU 的可用性判定为 torch_npu 模块存在且 torch.npu.is_available() 为真，因此不同硬件的适配主要围绕 "设备字符串替换" 展开。整体支持矩阵如下：

设备	适配方式	说明
NVIDIA GPU	无需修改	全部示例代码默认支持
AMD GPU	安装 ROCm 版 PyTorch	极少数依赖 CUDA 指令的模型不兼容
Apple Silicon	"cuda" 替换为 "mps" 或 "cpu"	显存与内存统一，需控制内存占用
昇腾 NPU	"cuda" 替换为 "npu"	需安装 CANN 与 torch_npu，参考 NPU 支持

开始之前，请先按照安装指南完成基础依赖的安装。

NVIDIA GPU：开箱即用

项目提供的所有示例代码均以 NVIDIA GPU 为默认运行环境，无需任何额外修改即可直接执行。安装依赖时按常规方式执行：

pip install -e .

训练场景下如需启用混合精度、梯度检查点等特性，参考各模型的 model_training 目录中的训练脚本即可，本项目的完整训练入门可参见模型训练指南。

AMD GPU：借助 ROCm 运行

AMD 官方提供了基于 ROCm 的 PyTorch 安装包，因此 DiffSynth-Studio 的绝大多数模型可以不做代码改动直接在 AMD GPU 上运行。安装时以 ROCm 6.4（Linux，以文档更新日期 2025 年 12 月 15 日为例）为例：

pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4

需要说明的限制是：极少数模型因为依赖 CUDA 专属指令（如特定 fused kernel），可能与 ROCm 不完全兼容。

Apple Silicon：统一内存下的 mps/cpu 切换

Apple Silicon（M 系列）设备的显存与系统内存是统一的，"cuda" 设备字符串无法使用。适配方法是把代码中所有 "cuda" 替换为 "mps" 或 "cpu"：

• "mps"：使用 Metal Performance Shaders 加速，优先推荐；
• "cpu"：当模型较大或希望严格控制内存占用时使用。

由于统一内存架构下显存即内存，替换后请留意系统内存占用，超大模型建议配合项目的 VRAM 管理机制（offload_device 等配置）运行，详见 VRAM 管理文档。

昇腾 NPU：推理与训练完整适配

环境准备

在昇腾 NPU 上运行需要两步环境准备：

1. 通过华为官方文档安装 CANN 工具包；
2. 从源码安装 DiffSynth-Studio 并携带 NPU 依赖选项：

git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio
# aarch64/ARM 架构
pip install -e .[npu_aarch64]
# x86 架构
pip install -e .[npu] --extra-index-url "https://download.pytorch.org/whl/cpu"

安装完成后，torch_npu 会接管 torch.npu 命名空间。从源码可以看到，当检测到 NPU 可用时，npu_compatible_device.py 会自动执行 torch.npu.config.allow_internal_format = False 以保证精度。

推理适配：cuda → npu

在 NPU 上推理时，核心改动是把代码中的 "cuda" 替换为 "npu"。下面以 Wan2.1-T2V-1.3B 视频生成推理为例，给出完整的适配 diff：

import torch
from diffsynth.utils.data import save_video, VideoData
from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig
from diffsynth.core.device.npu_compatible_device import get_device_name

vram_config = {
    "offload_dtype": "disk",
    "offload_device": "disk",
    "onload_dtype": torch.bfloat16,
    "onload_device": "cpu",
    "preparing_dtype": torch.bfloat16,
-   "preparing_device": "cuda",
+   "preparing_device": "npu",
    "computation_dtype": torch.bfloat16,
-   "computation_device": "cuda",
+   "computation_device": "npu",
}
pipe = WanVideoPipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
-   device="cuda",
+   device="npu",
    model_configs=[
        ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors", **vram_config),
        ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", **vram_config),
        ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth", **vram_config),
    ],
    tokenizer_config=ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="google/umt5-xxl/"),
-   vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 2,
+   vram_limit=torch.npu.mem_get_info(get_device_name())[1] / (1024 ** 3) - 2,
)

video = pipe(
    prompt="Documentary-style photography: a lively puppy running swiftly across lush green grass. The puppy has brownish-yellow fur, upright ears, and an alert, joyful expression. Sunlight bathes its body, making the fur appear exceptionally soft and shiny. The background is an open field with occasional wildflowers, and faint blue sky with scattered white clouds in the distance. Strong perspective captures the motion of the running puppy and the vitality of the surrounding grass. Mid-shot, side-moving viewpoint.",
    negative_prompt="Overly vibrant colors, overexposed, static, blurry details, subtitles, artistic style, painting, still image, overall grayish tone, worst quality, low quality, JPEG artifacts, ugly, distorted, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, fused fingers, motionless scene, cluttered background, three legs, many people in background, walking backward",
    seed=0, tiled=True,
)
save_video(video, "video.mp4", fps=15, quality=5)

这段代码有三个关键适配点，均建议照此模式迁移到其他模型：

1. VRAM 管理配置：vram_config 中 preparing_device 与 computation_device 由 "cuda" 改为 "npu"，其余 offload/onload 策略保持不变（示例中权重在 CPU 与磁盘之间调度，计算在 NPU 上进行）；
2. 管道设备：device="npu" 让 WanVideoPipeline 与 ModelConfig 将各子模块加载到 NPU；
3. 显存上限计算：torch.cuda.mem_get_info("cuda") 换成 torch.npu.mem_get_info(get_device_name())，其中 get_device_name() 由 diffsynth/core/device/npu_compatible_device.py 提供，返回当前 NPU 的设备字符串（如 "npu:0"），避免硬编码设备编号。

从源码实现看，WanVideoPipeline.from_pretrained 内部通过 npu_compatible_device 的设备工具完成张量搬运与同步，因此示例中 import diffsynth.core.device.npu_compatible_device 能确保所有内部调用都感知到 NPU 环境。

USP（Unified Sequence Parallel）序列并行

如果要在 NPU 上启用长序列推理的 USP 特性，需要额外安装两个第三方库：

pip install git+https://github.com/feifeibear/long-context-attention.git
pip install git+https://github.com/xdit-project/xDiT.git

其中 long-context-attention 提供长上下文注意力实现，xDiT 提供分布式推理框架。DiffSynth-Studio 在 diffsynth/utils/xfuser/xdit_context_parallel.py 中封装了 initialize_usp(device_type)、usp_dit_forward、usp_attn_forward 等函数，将 DiT 模型的前向过程改造成 USP 序列并行版本，NPU 上的分布式通信后端将自动选择 hccl（参见 npu_compatible_device.py 的 get_nccl_backend）。

训练：NPU 专用启动脚本

针对 NPU 训练，项目为每类模型都提供了现成的启动脚本，存放于 examples/<模型名>/model_training/special/npu_training/ 目录下，例如：

• examples/wanvideo/model_training/special/npu_training/Wan2.2-T2V-A14B-NPU.sh
• examples/wanvideo/model_training/special/npu_training/Wan2.1-T2V-14B-NPU.sh
• examples/qwen_image/model_training/special/npu_training/Qwen-Image-LoRA-NPU.sh
• examples/z_image/model_training/special/npu_training/Z-Image-Turbo-NPU.sh

这些脚本在原有训练命令之前设置了 NPU 专属的环境变量，并对特定模型开启了相关参数。

环境变量

PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

启用 NPU 显存池的可扩展段（expandable segments）功能，即虚拟显存特性。开启后显存分配器可以动态扩展已分配的内存段，减少频繁分配/释放带来的碎片化与显存不足问题，是 NPU 上训练大模型的首选配置。

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

CPU_AFFINITY_CONF

控制昇腾 NPU 的内核绑定（kernel binding）功能，用于提升算子调度性能：

取值	含义
0 或不设置	不启用绑定功能
1	启用粗粒度内核绑定
2	启用细粒度内核绑定

export CPU_AFFINITY_CONF=1

特定模型参数

模型	参数	说明
Wan 14B 系列	--initialize_model_on_cpu	14B 大模型需要在 CPU 上初始化
Qwen-Image 系列	--initialize_model_on_cpu	模型需要在 CPU 上初始化
Z-Image 系列	--enable_npu_patch	使用 NPU 融合算子替换 Z-Image 模型中的对应算子，提升模型在 NPU 上的性能

这些参数在训练脚本中均有真实对应。例如 examples/wanvideo/model_training/train.py 中定义了 --initialize_model_on_cpu（default=False），并在加载模型时通过 device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device 将模型先放到 CPU 初始化，再迁移到 NPU 训练，从而规避大模型在 NPU 上直接初始化导致的显存压力。examples/z_image/model_training/train.py 则定义了 --enable_npu_patch，并将其传入 ZImagePipeline.from_pretrained(..., enable_npu_patch=enable_npu_patch)。

以 Wan2.2-T2V-A14B 的 NPU 全参训练脚本（Wan2.2-T2V-A14B-NPU.sh）为例，它通过 --max_timestep_boundary 与 --min_timestep_boundary 将训练划分为高噪（timesteps [875, 1000]）与低噪（timesteps [0, 875)）两段，并在两段命令中都追加 --initialize_model_on_cpu：

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export CPU_AFFINITY_CONF=1

modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "wanvideo/Wan2.2-T2V-A14B/*" --local_dir ./data/diffsynth_example_dataset

# 高噪分支：boundary 对应 timesteps [875, 1000]
accelerate launch --config_file examples/wanvideo/model_training/full/accelerate_config_14B.yaml examples/wanvideo/model_training/train.py \
  --dataset_base_path data/diffsynth_example_dataset/wanvideo/Wan2.2-T2V-A14B \
  --dataset_metadata_path data/diffsynth_example_dataset/wanvideo/Wan2.2-T2V-A14B/metadata.csv \
  --height 480 --width 832 --num_frames 49 --dataset_repeat 100 \
  --model_id_with_origin_paths "Wan-AI/Wan2.2-T2V-A14B:high_noise_model/diffusion_pytorch_model*.safetensors,Wan-AI/Wan2.2-T2V-A14B:models_t5_umt5-xxl-enc-bf16.pth,Wan-AI/Wan2.2-T2V-A14B:Wan2.1_VAE.pth" \
  --learning_rate 1e-5 --num_epochs 2 \
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./models/train/Wan2.2-T2V-A14B_high_noise_full" \
  --trainable_models "dit" \
  --max_timestep_boundary 0.417 --min_timestep_boundary 0 \
  --initialize_model_on_cpu

# 低噪分支：boundary 对应 timesteps [0, 875)
accelerate launch --config_file examples/wanvideo/model_training/full/accelerate_config_14B.yaml examples/wanvideo/model_training/train.py \
  --dataset_base_path data/diffsynth_example_dataset/wanvideo/Wan2.2-T2V-A14B \
  --dataset_metadata_path data/diffsynth_example_dataset/wanvideo/Wan2.2-T2V-A14B/metadata.csv \
  --height 480 --width 832 --num_frames 49 --dataset_repeat 100 \
  --model_id_with_origin_paths "Wan-AI/Wan2.2-T2V-A14B:low_noise_model/diffusion_pytorch_model*.safetensors,Wan-AI/Wan2.2-T2V-A14B:models_t5_umt5-xxl-enc-bf16.pth,Wan-AI/Wan2.2-T2V-A14B:Wan2.1_VAE.pth" \
  --learning_rate 1e-5 --num_epochs 2 \
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./models/train/Wan2.2-T2V-A14B_low_noise_full" \
  --trainable_models "dit" \
  --max_timestep_boundary 1 --min_timestep_boundary 0.417 \
  --initialize_model_on_cpu

对于显存受限的场景，NPU 脚本也支持分阶段（split）训练。例如 Qwen-Image-LoRA-NPU.sh 先将 task 设为 "sft:data_process" 完成数据集切分缓存，再以 "sft:train" 阶段仅加载 DiT 权重进行 LoRA 训练，从而把 NPU 显存占用降到最低。

源码级解读：NPU 融合算子（Z-Image 性能优化）

文档参数表中 Z-Image 系列的 --enable_npu_patch 对应的是 diffsynth/core/npu_patch/npu_fused_operator.py 中的三个 NPU 融合算子：

• rms_norm_forward_npu：将 diffsynth/models/general_modules.py 中 RMSNorm 的 forward 替换为 torch_npu.npu_rms_norm 融合实现；
• rms_norm_forward_transformers_npu：替换 transformers 库中 Qwen3RMSNorm 的 forward；
• rotary_emb_Zimage_npu：将 Z-Image Attention.apply_rotary_emb 替换为 torch_npu.npu_rotary_mul（rotary_mode="interleave"）。

补丁的挂载逻辑位于 diffsynth/pipelines/z_image.py 的 apply_npu_patch(enable_npu_patch)：当 IS_NPU_AVAILABLE 且开关开启时，用上述融合算子替换 RMSNorm 与 RoPE 的前向实现，以减少 NPU 上的算子调用次数、提升性能；同时会给出警告提示，可通过 enable_npu_patch=False 关闭该特性。Z-Image 训练脚本（Z-Image-Turbo-NPU.sh）中的 --enable_npu_patch 即用于触发这一优化路径。

小结与排查建议

在 DiffSynth-Studio 中完成多硬件迁移，核心路径可归纳为三步：

1. 选对安装方式：NVIDIA 直接安装；AMD 换装 ROCm 版 PyTorch；Apple Silicon 无需改安装但改设备字符串；NPU 需先装 CANN，再以 [npu] 或 [npu_aarch64] 选项从源码安装；
2. 统一设备字符串：将所有 "cuda" 替换为 "mps"/"cpu"（Apple Silicon）或 "npu"（昇腾），并同步替换 torch.cuda.* 的显存查询 API；
3. 善用官方适配资产：NPU 训练直接复用 examples/<模型名>/model_training/special/npu_training/ 下的脚本，其中已内置 expandable_segments 虚拟显存、CPU 亲和性绑定等性能优化，以及 --initialize_model_on_cpu、--enable_npu_patch 等模型级参数。

如果在 NPU 上遇到算子不支持或性能瓶颈，优先检查是否遗漏了 NPU 专属参数，或参考 VRAM 管理文档 调整 offload 策略，这通常是统一内存设备与昇腾 NPU 上最有效的调优手段。