TurboDiffusion

环境要求

• Python >= 3.9（README 示例使用 3.12）
• PyTorch >= 2.7.0，推荐 torch==2.8.0（更高版本可能导致 OOM）
• 建议使用 NVIDA 显卡（加速数据基于单张 RTX 5090）
• 下载模型权重/数据集前需安装 git lfs
• 训练 rCM/SLA 需额外依赖 megatron-core、hydra-core、wandb、webdataset、transformer_engine

安装与启动步骤

1. 1创建并激活环境

   用 conda 新建一个 Python 3.12 的独立环境，避免与系统环境冲突，随后激活该环境。

   conda create -n turbodiffusion python=3.12
   conda activate turbodiffusion

   复制

2. 2pip 安装 TurboDiffusio

   官方推荐的安装方式，务必带上 --no-build-isolation 参数，否则编译可能失败。

   pip install turbodiffusion --no-build-isolation

   复制

3. 3或从源码编译安装

   想改代码或参与训练时用源码方式：拉取仓库、同步子模块后用 -e 可编辑安装。

   git clone https://github.com/thu-ml/TurboDiffusion.git
   cd TurboDiffusion
   git submodule update --init --recursive
   pip install -e . --no-build-isolation

   复制

4. 4安装 SageSLA 加速算子

   要启用基于 SageAttention 的快速 SLA 前向，必须先安装 SpargeAttn，同样需要 --no-build-isolation。

   pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation

   复制

5. 5安装训练额外依赖

   仅在做 rCM/SLA 训练时需要，普通推理可跳过；transformer_engine 需要单独指定 --no-build-isolation。

   pip install megatron-core hydra-core wandb webdataset
   pip install --no-build-isolation transformer_engine[pytorch]

   复制

6. 6下载模型权重

   通过 HuggingFace 仓库克隆到 assets/checkpoints 目录，执行前请确认已安装 git lfs，否则大文件会缺失。

   git clone https://huggingface.co/worstcoder/Wan assets/checkpoints

   复制

7. 7转换检查点格式

   FSDP2 依赖 Distributed Checkpoint，训练前需先把 .pth 教师权重转成 .dcp；训练后可用 scripts/dcp_to_pth.py 转回。

   python -m torch.distributed.checkpoint.format_utils torch_to_dcp assets/checkpoints/Wan2.1-T2V-1.3B.pth assets/checkpoints/Wan2.1-T2V-1.3B.dcp

   复制

8. 8下载训练数据集

   只有需要训练时才下载，同样是 git lfs 大文件仓库，克隆到 assets/datasets 目录。

   git clone https://huggingface.co/datasets/worstcoder/Wan_datasets assets/datasets

   复制

如何确认成功

README 未给出推理启动命令；执行 pip show turbodiffusion 能显示包信息、且 assets/checkpoints 下有权重文件即表示安装与下载成功。

常见问题

注意事项

• README 提示 checkpoints 和论文尚未定稿，后续会更新以提升质量。
• 当前模型仅针对长英文提示词训练，其他语言或短提示词效果可能不佳。
• 核心加速组件为 SageAttention、SLA（Sparse-Linear Attention）与 rCM 时间步蒸馏。
• README 未提供推理/生成视频的运行命令与参数，实际使用请以仓库后续更新为准。