🧩 Deepspeed-diffusers

Read this in English.

deepspeed-diffusers 是一个原生结合Deepspeed、Diffusers、Peft库来训练扩散模型（Diffusion Models）的项目。

动机

训练过Diffusers的开发者都知道目前已经有了Huggingface另一个产品Accelerate来并行训练扩散模型。

为什么我还开源了这个项目呢？

• 一个非常滑稽的事实是当初我没有仔细看到diffusers的训练示例文件是部分支持Zero3而不是不支持。本项目最后实现下来，结果Zero3实现思路与其不谋而和。
• 可以安慰自己的是，虽然现在集成框架很多，使用起来也很便捷，但是集成的东西太多了，就不容易二次开发了！
• Slurm调度系统的脚本起码我也提供了，使用和修复了Deepspeed本来就支持却又不太行的Slurm启动器。

为了充分单独发挥Deepspeed的能力，本项目就这么诞生了。

此外，本项目在原生结合的时候也借鉴了OvJat 、afiaka87、liucongg的代码，在此非常感谢这些项目的付出！

在这些项目中，本项目的缺点是什么？

1. 我在实现的时候，ZeRO-1,2也只支持UNet2DConditionModel。目前我还没有发现我差别人具体在哪，还得充分测试。
2. 无法保存Lora训练时候的检测点，这和Peft库太过集成绑定有些关系，除非自己写个Lora微调逻辑。

除此之外，Accelerate有的功能本项目也有，本项目主打一个简洁实用🥹

最近更新 🔥

• [2024/05/30] 类似Accelerate示例，不把stable diffusion类完全加入DeepSpeed，仅加载Unet，这样减少了大概3GB+的显存
• [2024/05/25] 支持数据集自定义，支持EMA
• [2024/04/18] 修复训练文本不匹配的问题，增加推理验证
• [2024/04/10] 支持slurm多节点和单节点训练。
• [2024/04/01] 支持wandb。
• [2024/03/31] 支持lora微调，修复了全量微调后的生成图片总是黑色图片的问题，增加了slurm单节点训练脚本
• [2024/03/29] deepspeed-diffusers 发布了，支持Unet全量微调。

演示

安装依赖

在运行脚本之前，请确定安装了所有的依赖：

• 请保证源码是最新的：git clone https://github.com/dyedd/deepspeed-diffusers
• 然后cd到文件夹并运行：pip -r install requirements.txt

宝可梦数据集示例

强烈推荐直接把数据集下载到本地，而不是通过脚本自动下载的是缓存文件。否则哪天断网也要断Huggingface的数据反馈以及对于国内环境不是非常友好。

根据下载的目录，修改cfg.json的dataset_dir内容。

下载权重

接下来的示例结果都是在stable-diffusion-1.5的权重下实验的。

注意：如果您使用 stable-diffusion-2 768x768 模型，请将cfg/default.json的resolution更改为 768。

在此，仍然强烈自己下载git clone权重，不要通过huggingface自动下载~然后修改cfg/default.json的pretrained_model_name_or_path。

配置文件解读

项目所需的参数都写在了cfg/default.json： 以下是每个字段的解释：

• "num_epochs": 训练的总轮数。
• "validation_epochs": 每多少轮进行一次验证。
• "max_train_steps": 限制最大训练步数，0表示不限制。
• "lr_warmup_steps": 学习率预热的步数。
• "save_interval": 模型保存的间隔步数。
• "seed": 随机种子，用于确保实验的可重复性。
• "validation_prompts": 验证时使用的提示。
• "pretrained_model_name_or_path": 预训练模型的名称或路径。
• "dataset_dir": 数据集的目录。
• "imagefolder": 是否使用包含图像的文件夹作为数据源。
• "checkpoint_dir": 检查点（模型状态）保存的目录。
• "output_dir": 输出（如训练日志、生成的图像等）的目录。
• "use_lora": 是否使用LoRA（Long Range Attention）技术，以及相关的参数。
  • "action": 是否启用LoRA。
  • "rank": LoRA的秩。
  • "alpha": LoRA的α参数。
• "dataloader_num_workers": 数据加载器的工作线程数。
• "resume_from_checkpoint": 是否从检查点恢复训练。
• "use_ema": 是否使用指数移动平均（Exponential Moving Average）。
• "offline": 是否在wandb离线模式下运行。
• "resolution": 图像的分辨率。
• "center_crop": 是否对图像进行中心裁剪。
• "random_flip": 是否对图像进行随机翻转。
• 之后都是DeepSpeed常用的配置，如有不解，请查看官方文档。如果有添加，不要忘记同时修改utils.py的deepspeed_config_from_args函数。

训练

本项目支持2种训练模式。

• 全量微调unet，在混合精度FP16下，如果不开启Zero，每张卡的bacth_size为4, 显存大致在11.61到23.32GB。
• Lora+unet，在混合精度FP16下，如果不开启Zero，每张卡的bacth_size为4, 显存大致在2到13GB。

只要开启梯度累积和Zero，那么全量微调也能实现实现16GB以下的显卡训练了！

如果你在本地，可以直接通过bash scripts/train_text_to_image.sh运行脚本；

如果你在slurm系统，在修改部分信息后，可以通过sbatch scripts/train_text_to_image.slurm下提交。

推理/采样

同样支持本地和Slurm调度，命令分别为

bash scripts/test_text_to_image.sh
sbatch scripts/test_text_to_image.slurm

可能出现的问题

Note

1. 生成的图像都是黑色图片或者报错RuntimeWarning: invalid value encountered in cast images = (images * 255).round().astype("uint8")

请注意，本项目配置文件中关于优化器，学习率的参数仅对于宝可梦这个数据集而言。这个问题是因为优化器，学习率不适合训练集，而造成训练的损失一直是none。 2. 提交的slurm脚本运行时间为0秒就退出 这是因为slurm脚本里写了日志保存在log文件夹，而该文件夹目前不存在，就无法运行。

引用

如果您在论文和项目中使用了Deepspeed-diffusers，请使用以下BibTeX引用它。

@Misc{Deepspeed-diffusers,
  title =        {Deepspeed-diffusers: training diffusers with deepspeed.},
  author =       {Ximiao Dong},
  howpublished = {\url{https://github.com/dyedd/deepspeed-diffusers}},
  year =         {2024}
}