DiffSynth-Studio MUSA Patch 使用文档

1. 项目简介

DiffSynth-Studio MUSA Patch 是为 DiffSynth-Studio 扩散模型引擎开发的摩尔线程（Moore Threads）GPU 适配补丁。它使 DiffSynth-Studio 能够在摩尔线程 MUSA 架构的 GPU 上高效运行视频生成和训练任务，同时保持与 NVIDIA CUDA 环境的完全兼容。

该补丁采用双后端自适应设计：导入时自动检测当前 GPU 类型（MUSA 或 CUDA），并透明切换底层设备调用，用户代码无需任何修改即可在两种硬件上运行。

2. 功能特性

功能	说明
CUDA-MUDA 透明转换	将 `torch.cuda.` 调用自动映射为 `torch.musa.`，无需修改业务代码
设备自动检测	运行时自动识别 MUSA 或 CUDA 环境，选择对应后端
Disk Offload 兼容	修复 safetensors 在非 CUDA 设备上的加载问题
FSDP2 分布式训练	支持 FSDP2 全分片数据并行，含 text_encoder 和 image_encoder 分片
FusedAdamW 优化器	在 MUSA 上使用融合 AdamW 优化器提升训练速度
梯度检查点控制	可配置 DiT Block 的梯度检查点层数，精确控制显存占用
批量训练	支持多样本批量训练（batch_size > 1），提供自定义 collate 函数
Profiling 工具	内置 PyTorch Profiler 封装，支持 MUSA/CUDA 双后端性能分析
训练路径修复	修复本地路径 + glob 模式的模型加载问题

3. 系统要求

MUSA 环境

依赖	版本	说明
MUSA DDK (Driver)	4.3.5-server	摩尔线程 GPU 驱动
MUSA Toolkit	4.3.3	含 mcc、mublas、mufft、murand、musolver 等库
mccl	compatible with torch_musa	摩尔线程集合通信库
Python	3.10	/
PyTorch	2.7.1	commit: `e2d141dbde55c2a4370fac5165b0561b6af4798b`
torch_musa	2.7.1	commit: `9f1bb312567dde50f75fb2c86c74d40ad794a08e`
DiffSynth-Studio	2.0.9	commit: `5c89a15b9aa1e19d311108511970dca1718f4d69`
uv	latest	Python 包管理器

— 已验证的 MUSA 驱动及工具链版本：

组件	版本
MUSA Driver (DDK)	4.3.5-server
musa_toolkits	4.3.3
mcc (MUSA C Compiler)	4.3.3
mublas / mublasLt (BLAS)	1.10.3
mufft (FFT)	1.10.3
murand (Random)	1.1.3
musolver (Solver)	1.4.3
musparse (Sparse)	1.5.3
muPP (Image Processing)	1.11.3
MUPTI (Profiling)	1.0.3
mtjpeg (JPEG)	1.0.2
musa_runtime	4.3.3

查看版本信息

在宿主机上查看 DDK 驱动版本：

dpkg -s musa

在容器内查看 MUSA 工具链版本：

cat /usr/local/musa/version.json

CUDA 环境

依赖	版本要求	说明
Python	>= 3.10	/
PyTorch	>= 2.5	CUDA 版本
DiffSynth-Studio	latest (main branch)	基础框架

4. 安装指南

基础 Docker 镜像

docker run -it --privileged --name <container_name> --env MTHREADS_VISIBLE_DEVICES=all --shm-size=256g -v <your_data/model_path> sh-harbor.mthreads.com/mt-ai/mtwan:4.3.3-pt2.7-v0.2 /bin/bash

— 推荐使用以上基础镜像，已包含 MUSA DDK 4.3.5、torch_musa 2.7.1 及相关依赖，无需从零构建 torch_musa。

4.1 完整安装（MUSA 环境）

# ========== Step 1: 构建 torch_musa（如已安装可跳过）==========
git clone git@sh-code.mthreads.com:ai/torch_musa.git
cd torch_musa
bash build.sh    # 使用 main 分支
cd ..

# ========== Step 2: 克隆 DiffSynth-Studio ==========
git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio
# 请按照官方原版 DiffSynth-Studio 的安装指南安装。
# 若是 MUSA 环境，请绕过 torch 和 nvidia 相关依赖的包。

# ========== Step 3: 安装 uv ==========
curl -LsSf https://astral.sh/uv/install.sh | sh
grep -qF 'export PATH="$HOME/.local/bin:$PATH"' ~/.bashrc || echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

# ========== Step 4: 克隆 MUSA Patch ==========
git clone https://sh-code.mthreads.com/ai/DiffSynth-Studio-musa-patch.git

# ========== Step 5: 创建虚拟环境 ==========
cd DiffSynth-Studio-musa-patch
uv venv --python 3.10 --system-site-packages
source .venv/bin/activate

# ========== Step 6: 安装依赖 ==========
uv pip install --no-deps -r requirements.txt --index-strategy unsafe-best-match
cd ..
# 安装 diffsynth
uv pip install -e . --no-deps

4.2 CUDA 环境安装

在 CUDA 环境下，patch 包同样适用（自动检测后端）。安装流程相同，跳过 torch_musa 构建步骤即可。

# 克隆 DiffSynth-Studio
git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio

# 安装 uv
curl -LsSf https://astral.sh/uv/install.sh | sh
grep -qF 'export PATH="$HOME/.local/bin:$PATH"' ~/.bashrc || echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

# 克隆 MUSA Patch
git clone https://github.com/your-org/DiffSynth-Studio-musa-patch.git

# 创建虚拟环境并安装依赖
cd DiffSynth-Studio-musa-patch
uv venv --python 3.10 --system-site-packages
source .venv/bin/activate
uv pip install --no-deps -r requirements.txt
cd ..

# 安装 diffsynth
uv pip install -e . --no-deps

4.3 验证安装

import musa_patches

# 检测当前后端
if musa_patches.IS_MUSA:
        print("运行在摩尔线程 MUSA 后端")
else:
    print("运行在 NVIDIA CUDA 后端")

5. 核心架构

补丁通过 Python 模块导入自动生效。只需在代码顶部添加两行即可完成所有适配：

import sys
sys.path.insert(0, "DiffSynth-Studio-musa-patch")
import musa_patches   # 导入即生效

musa_patches/__init__.py 执行以下操作：

py_patch() — 为 Python < 3.9 添加 math.lcm 兼容
patch_before_import_diffsynth() — 在 diffsynth 导入前修补其模块
各子模块的 apply_*() 函数自动在导入时运行

6. 推理使用

以下示例展示如何使用 Wan2.1-T2V-1.3B 模型生成视频。代码自动检测 GPU 后端，在 MUSA 和 CUDA 上均可运行。

import torch
import sys
sys.path.insert(0, "DiffSynth-Studio-musa-patch")

import musa_patches

from diffsynth.utils.data import save_video
from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig

# 自动检测设备类型 / Auto-detect device type
if musa_patches.IS_MUSA:
    target_device = "musa:0"
    total_vram = torch.musa.mem_get_info("musa")[1] / (1024 ** 3)
else:
    target_device = "cuda:0"
    total_vram = torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3)

# VRAM 配置（磁盘卸载）/ VRAM config (disk offload)
vram_config = {
    "offload_dtype": "disk",
    "offload_device": "disk",
    "onload_dtype": torch.bfloat16,
    "onload_device": "cpu",
    "preparing_dtype": torch.bfloat16,
    "preparing_device": target_device,
    "computation_dtype": torch.bfloat16,
    "computation_device": target_device,
}

model_base_path = "/workspace/models/Wan2.1-T2V-1.3B"

pipe = WanVideoPipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device=target_device,
    model_configs=[
        ModelConfig(model_id=model_base_path,
                    origin_file_pattern="diffusion_pytorch_model*.safetensors",
                    skip_download=True, **vram_config),
        ModelConfig(model_id=model_base_path,
                    origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth",
                    skip_download=True, **vram_config),
        ModelConfig(model_id=model_base_path,
                    origin_file_pattern="Wan2.1_VAE.pth",
                    skip_download=True, **vram_config),
    ],
    tokenizer_config=ModelConfig(model_id=model_base_path,
                                  origin_file_pattern="google/umt5-xxl/",
                                  skip_download=True),
    redirect_common_files=False,
    vram_limit=total_vram - 2,
)

# 生成视频 / Generate video
video = pipe(
    prompt="A cute puppy running on a green lawn, documentary style.",
    negative_prompt="static, blurry, low quality",
    seed=0,
    tiled=True,
)
save_video(video, "output.mp4", fps=15, quality=5)

7. 训练使用

7.1 环境准备

训练前使用 setup_env.sh 自动配置 GPU 环境：

# 自动检测并配置 8 GPU 环境 / Auto-detect and configure 8 GPU environment
source examples/wanvideo/model_training/setup_env.sh 8

# 脚本会自动设置 / The script auto-configures:
# - MUSA: MUSA_VISIBLE_DEVICES, MCCL_*, TORCH_MUSA_FSDP2_*
# - CUDA: CUDA_VISIBLE_DEVICES, NCCL_*, PYTORCH_CUDA_ALLOC_CONF

— setup_env.sh 和 accelerate_config_fsdp2.yaml 位于 model_training/ 目录下，由 full/ 和 lora/ 下的训练脚本共享引用：

examples/wanvideo/model_training/
├── setup_env.sh                    # 共享环境配置脚本 / Shared env setup
├── accelerate_config_fsdp2.yaml    # 共享 FSDP2 配置 / Shared FSDP2 config
├── train.py
├── full/                           # 全参数训练脚本 / Full-parameter training
│   ├── Wan2.1-T2V-1.3B.sh
│   └── ...
└── lora/                           # LoRA 训练脚本 / LoRA training
    ├── Wan2.1-T2V-1.3B.sh
    └── ...

7.2 准备数据集

训练数据集使用 CSV 或 JSON 格式的元数据文件：

video,prompt
videos/clip1.mp4,"A sunset over a small town"
videos/clip2.mp4,"Waves crashing on the shore"

对于 I2V 训练，添加 input_image 列 / For I2V training, add an input_image column:

video,prompt,input_image
videos/clip1.mp4,"A puppy running","images/frame1.jpg"

7.3 配置训练数据路径

训练脚本通过两个关键变量指定数据集路径，所有训练脚本（examples/wanvideo/model_training/full/*.sh 和 lora/*.sh）头部均有以下配置段：

# ============================================================
# Dataset Configuration
# ============================================================
DATASET_NAME="/workspace/datasets/demo_video_dataset/videos"
DATASET_META_NAME="/workspace/datasets/demo_video_dataset/metadata.csv"

DATASET_NAME：对应启动参数 --dataset_base_path，即数据集的根目录。metadata.csv 中所有相对路径（如 videos/clip1.mp4）都基于此目录解析。
DATASET_META_NAME：对应启动参数 --dataset_metadata_path，即元数据 CSV/JSON 文件的绝对路径。

推荐的目录结构 / Recommended directory structure：

/workspace/datasets/my_dataset/
├── metadata.csv              # 元数据文件
├── videos/                   # 视频/图像文件目录
│   ├── clip1.mp4
│   ├── clip2.mp4
│   └── ...
└── images/                   # I2V 训练用的参考图像目录（可选）
    ├── ref1.jpg
    └── ref2.jpg

修改数据集路径：将脚本中的两个变量改为您自己的路径即可，例如：

DATASET_NAME="/data/my_training_data"
DATASET_META_NAME="/data/my_training_data/metadata.csv"

配置模型路径：训练脚本中的 --model_id_with_origin_paths 参数指定预训练模型文件，格式为 模型目录:文件匹配模式，多个文件用逗号分隔。修改为您的模型存放路径：

# T2V 模型路径示例（以 Wan2.1-T2V-1.3B 为例）
--model_id_with_origin_paths "/workspace/models/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors,/workspace/models/Wan2.1-T2V-1.3B:models_t5_umt5-xxl-enc-bf16.pth,/workspace/models/Wan2.1-T2V-1.3B:Wan2.1_VAE.pth"

# I2V 模型需要额外加载 CLIP 图像编码器
--model_id_with_origin_paths "/workspace/models/Wan2.1-I2V-14B-480P:diffusion_pytorch_model*.safetensors,/workspace/models/Wan2.1-I2V-14B-480P:models_t5_umt5-xxl-enc-bf16.pth,/workspace/models/Wan2.1-I2V-14B-480P:Wan2.1_VAE.pth,/workspace/models/Wan2.1-I2V-14B-480P:models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"

# Wan2.2 模型分 high_noise / low_noise 两个子目录
--model_id_with_origin_paths "/workspace/models/Wan2.2-T2V-A14B:high_noise_model/diffusion_pytorch_model*.safetensors,..."

# ============================================================
# Dataset Configuration
# ============================================================
DATASET_NAME="/workspace/datasets/demo_video_dataset/videos"
DATASET_META_NAME="/workspace/datasets/demo_video_dataset/metadata.csv"

DATASET_NAME → --dataset_base_path: The dataset root directory. All relative paths in metadata.csv (e.g. videos/clip1.mp4) are resolved relative to this directory.
DATASET_META_NAME → --dataset_metadata_path: The absolute path to the CSV/JSON metadata file.

To use your own dataset, simply update these two variables in the script:

DATASET_NAME="/data/my_training_data"
DATASET_META_NAME="/data/my_training_data/metadata.csv"

Configure model paths: The --model_id_with_origin_paths parameter specifies pretrained model files in the format model_dir:file_glob_pattern, with multiple entries separated by commas. Update to your model directory:

# T2V model example (Wan2.1-T2V-1.3B)
--model_id_with_origin_paths "/workspace/models/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors,/workspace/models/Wan2.1-T2V-1.3B:models_t5_umt5-xxl-enc-bf16.pth,/workspace/models/Wan2.1-T2V-1.3B:Wan2.1_VAE.pth"

# I2V models require an additional CLIP image encoder
--model_id_with_origin_paths "/workspace/models/Wan2.1-I2V-14B-480P:diffusion_pytorch_model*.safetensors,...,/workspace/models/Wan2.1-I2V-14B-480P:models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"

# Wan2.2 models have high_noise/low_noise sub-directories
--model_id_with_origin_paths "/workspace/models/Wan2.2-T2V-A14B:high_noise_model/diffusion_pytorch_model*.safetensors,..."

7.4 启动训练

单机 8 卡训练（T2V）

# Wan2.1-T2V-1.3B 全参数训练 / Full-parameter training
bash examples/wanvideo/model_training/full/Wan2.1-T2V-1.3B.sh

训练脚本的核心参数说明 / Key training script parameters:

参数	默认值	说明 / Description
`--dataset_base_path`	—	数据集根目录 / Dataset root directory
`--dataset_metadata_path`	—	元数据文件路径 / Metadata file path
`--height`	480	视频高度 / Video height
`--width`	832	视频宽度 / Video width
`--num_frames`	81	帧数 / Number of frames
`--model_id_with_origin_paths`	—	模型路径:文件模式 / Model path:file pattern
`--learning_rate`	1e-5	学习率 / Learning rate
`--num_epochs`	30	训练轮数 / Number of epochs
`--trainable_models`	dit	可训练模型 / Trainable models
`--batch_size`	1	每 GPU 批量大小 / Per-GPU batch size
`--recompute_num_layers`	None	梯度检查点层数 / Gradient checkpoint layers
`--gradient_accumulation_steps`	4	梯度累积步数 / Gradient accumulation steps
`--initialize_model_on_cpu`	False	在 CPU 上初始化 / Initialize on CPU
`--max_timestep_boundary`	1.0	最大时间步边界 / Max timestep boundary
`--min_timestep_boundary`	0.0	最小时间步边界 / Min timestep boundary

I2V 训练

# I2V 训练需要添加 extra_inputs 参数 / I2V training requires extra_inputs parameter
bash examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh

Wan2.2 双阶段训练

Wan2.2-A14B 模型需要分两个阶段训练：先训练高噪声模型（timestep 0.417~~1.0），再训练低噪声模型（timestep 0~~0.417）：

Wan2.2-A14B requires two-stage training: first train the high-noise model (timestep 0.417~~1.0), then the low-noise model (timestep 0~~0.417):

# 脚本内包含两个连续的 accelerate launch 命令
# Script contains two consecutive accelerate launch commands
bash examples/wanvideo/model_training/full/Wan2.2-T2V-A14B.sh

7.5 自定义训练脚本

source examples/wanvideo/model_training/setup_env.sh 8
source DiffSynth-Studio-musa-patch/.venv/bin/activate

accelerate launch \
  --config_file "$ACCELERATE_CONFIG" \
  --num_processes 8 \
  --mixed_precision bf16 \
  DiffSynth-Studio-musa-patch/examples/wanvideo/model_training/train.py \
  --dataset_base_path /path/to/dataset \
  --dataset_metadata_path /path/to/metadata.csv \
  --height 480 --width 832 --num_frames 81 \
  --model_id_with_origin_paths "/path/to/model:diffusion_pytorch_model*.safetensors,/path/to/model:models_t5_umt5-xxl-enc-bf16.pth,/path/to/model:Wan2.1_VAE.pth" \
  --learning_rate 1e-5 --num_epochs 30 \
  --trainable_models dit --save_steps 100 \
  --initialize_model_on_cpu \
  --recompute_num_layers 40

许可证

本项目遵循 DiffSynth-Studio 的原始许可证。

目录​

1. 项目简介​

2. 功能特性​

3. 系统要求​

MUSA 环境​

查看版本信息​

CUDA 环境​

4. 安装指南​

基础 Docker 镜像​

4.1 完整安装（MUSA 环境）​

4.2 CUDA 环境安装​

4.3 验证安装​

5. 核心架构​

6. 推理使用​

7. 训练使用​

7.1 环境准备​

7.2 准备数据集​

7.3 配置训练数据路径​

7.4 启动训练​

单机 8 卡训练（T2V）​

I2V 训练​

Wan2.2 双阶段训练​

7.5 自定义训练脚本​

许可证​

这篇文章对您有帮助吗？

目录