在 Amazon P5 上进行 SDXL Controlnet 微调训练

背景及业务场景

Controlnet 原理

Controlnet 是 Stable Diffusion 微调神经网络的一种，它使得用户可以通过施加额外条件，细粒度地控制扩散模型的生成过程。这一技术最初由 Adding Conditional Control to Text-to-Image Diffusion Models 这篇论文提出，并很快地风靡了扩散模型的开源社区。主要的功能是将大模型（比如预训练 Stable Diffusion）在自己的小数据集上 finetune，使新的模型能够适应自己的数据集（比如自定义的风格/任务）且对原模型的影响较小（保留原模型绝大部分能力）。

Controlnet finetune 的主要原理如下图所示：

Controlnet 将我们需要修改的网络层复制一份，在前后分别增加参数全为 0 的卷积层，作为右侧的 trainable 分支。trainable 分支的输出和原来的网络输出加在一起作为最终的输出。由于复制出来的网络前后的增加的卷积层参数均为 0，输出也一定是 0，因此整个网络的输出还和原来的网络完全一样。在自己的数据集训练时将原来的网络固定，只训练 trainable 的部分即可。

通过 Controlnet 可以使用不同的方式控制生成的内容，目前有 8 个开源的不同 controlnet 模型，包括姿态估计、深度图、边缘图、素描图等，在众多业务场景上有着丰富的应用方式，比如：可以通过 Controlnet 的 pose，自定义人体姿态或者从图片生成姿态，绘制新的置顶姿态的图；或者使用 canny；固定人物主体轮廓, 改变外观和周边的环境，从而生成新的图片。

Controlnet finetune 场景

在电商及广告数字人领域，通常会用 Stable Diffusion 和 controlnet 进行模特穿戴，广告素材等图像的生成，但针对效果和精准程度较高的复杂场景，如面部多角度的人物姿态图片生成，业界没有合适的开源的 controlnet 模型，需要使用自己的图像，使用 Contolnet 网络进行二次微调训练，以满足个性化和持续优化的需求。

本文下述章节介绍了基于 Controlnet 技术，使用 AWS P5（nvidia H100）高性能 GPU 服务器资源，进行微调的过程及训练完成后模型部署，文中涉及到众多实施落地的技术实现和代码/脚本示例，可帮助感兴趣的小伙伴基于 Stable Diffusion XL 版本进行 Controlnet 微调及部署参考。

训练准备

与 LoRA 微调类似，Controlnet 训练数据集可以采用 HuggingFace dataset 数据集格式，以便使用 HugginFace Dataset API 进行训练数据的加载。

HuggingFace Dataset 数据集中存放图像的目录格式如下：

如上所示，通过 train/validata 子目录即可使用 HuggingFace dataset 的 split 方法方便的进行训练集类别的加载，再通过 lable 子目录获取更详细的数据集标签字段信息，当然也可以自定义图像路径，在 dataset 的 jsonl 配置文件中写明即可。

我们可以参考 HugginFace 开源的 Controlnet 数据集（https://huggingface.co/datasets/fusing/fill50k）来看下数据集字段内容及格式：

如上图所示

ground truth 图片（image）：这里指的就是真实人脸图片
条件图片（image_seg）：这里指的就是画出来的关键点
说明文字（image_caption）：描述图片的文字

按照以上字段和格式准备 Controlnet 微调的训练数据集即可，这里给出一条整理后的训练数据示例：

准备好的训练数据集可以托管在各种存储设施上，比如 Local 本地磁盘 / AWS S3 / AWS EFS…etc，当然也可以 upload 到 Huggingface Datasets。

制作完成后建议存两份，一份 Public 可以快速验证数据集格式和类型、另外一份 Private 可以进行训练。

关于制作 HuggingFace Dataset 训练集细节这里不再赘述，感兴趣的小伙伴可以参考 HF 官方文档进行制作：https://huggingface.co/docs/datasets/create_dataset。

训练过程

训练环境及数据

这里使用 AWS 最新的 P5 机器学习实例进行 controlnet 的模型训练。

Amazon P5 实例是由最新的 NVIDIA H100 GPU 提供支持，带有 8 张 NVIDIA H100 GPU 卡，每张卡 80Gb 的显存，为深度学习（DL）和高性能计算（HPC）应用程序提供高性能的训练机器。与上一代 P4 GPU 的实例机器相比，同等任务的训练时间缩短 6 倍，并且总体成本降低 40%，从而帮助客户能更快训练和部署 SDXL 及 LLM 模型，为 GenAI 应用程序提供动力，快速地迭代解决方案，更快地进入市场。

在 EC2 上配置和使用 P5 实例如下步骤所示：

P5 实例创建及 EFS 配置

新建 EC2 实例，选择 ubuntu deeplearning AMI 镜像

选择 P5.48xlarge（H100 GPU，8 卡 80G 显存）实例类型

配置 EFS 存储，用于高速读取训练数据

EFS 网络配置中绑定安全组，并设置 NFS Inbound

sudo apt-get update
sudo apt-get -y install git binutils
git clone https://github.com/aws/efs-utils

cd efs-utils/
./build-deb.sh
sudo apt-get -y install ./build/amazon-efs-utils*deb
cd ..
mkdir efs
sudo mount -t efs -o tls fs-09984b9c17da1cbad.efs.us-east-1.amazonaws.com /home/ubuntu/efs/

确认 P5 实例类型 CPU/GPU 资源

至此 P5 实例成功创建。

训练数据下载

训练数据下载采用 Git-LFS，Git-LFS 采用了不同的传输协议，加上 AWS 高速网络的加持，可以达到 300 MiB/s 左右的速度，而传统方式下载 HuggingFace 的 Dataset 速度在 30 MiB/s 左右，如下图所示：

训练脚本

这里我们采用 Diffuser 官方 controlnet fine tune 脚本，在 diffuser 源代码的 example/controlnet 目录下，我们使用 accelerate 拉起单机多卡分布式训练的程序，代码示例如下所示：

!accelerate launch train_controlnet_sdxl.py \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" \
 --output_dir="model_out" \
 --dataset_name=multimodalart/facesyntheticsspigacaptioned \
 --conditioning_image_column=spiga_seg \
 --image_column=image \
 --caption_column=image_caption \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" \
 --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" \
 --train_batch_size=4 \
 --num_train_epochs=3 \
 --tracker_project_name="controlnet" \
 --enable_xformers_memory_efficient_attention \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --push_to_hub

其中主要的参数说明如下：

pretrained_model_name_or_path：基础的 Stable Diffusion 模型，这里我们使用 v2-1 版本，因为这一版生成人脸效果更好
output_dir：保存模型的目录文件夹
dataset_name：用于训练的数据集，这里我们使用 Face Synthetics SPIGA with captions
conditioning_image_column：数据集中包含条件图片的这一栏的名称，这里我们用 spiga_seg
image_column：数据集中包含 ground truth 图片的这一栏的名称，这里我们用 image
caption_column：数据集中包含文字说明的这一栏的名称，这里我们用 image_caption
resolution：ground truth 图片和条件图片的分辨率，这里我们用 512x512
learning_rate：学习率
validation_image：这里是让你在训练过程中偷窥一下效果的。每隔 validation_steps 步训练，这些验证图片都会跑一下，让你看看当前的训练效果。请在这里插入一个指向一系列条件图片的本地路径
validation_prompt：这里是一句文本提示，用于和你的验证图片一起验证当前模型。你可以根据你的需要设置
train_batch_size：这是训练时使用的 batch size。因为我们用的是 V100，所以我们还有能力把它设成 4。但如果你的 GPU 显存比较小，我们推荐直接设成 1
num_train_epochs：训练模型使用的轮数
checkpointing_steps：每隔这么多步，我们都会保存一下模型的中间结果检查点。这里我们设置成 5000，也就是每训练 5000 步就保存一下检查点
validation_steps：每隔这么多步，validation_image 和 validation_prompt 就会跑一下，来验证训练过程
report_to：向哪里报告训练情况。这里我们使用 Weights and Biases 这个平台
push_to_hub：将最终结果推到 Hugging Face Hub

主要的调参优化如下：

对于大量数据集，Controlnet SDXL 的模型训练，当 batch size/gradient step 过大时，容易显存 OOM，因此需要根据数据规模，调整其训练参数，其调整的参数如下所示：
```
train_batch_size=1 ###根据需要调小训练批次
gradient_accumulation_steps=4  ####根据需要调小梯度更新批次
```
num_train_epochs 这个参数决定多少轮训练，每一轮训练模型都会跑完整个数据集。需要根据你的 loss 情况决定 epochs 数量，我们实验用的是 3 轮，但最好的结果是出现在一轮多一点的地方。当训练了 3 轮时，我们的模型过拟合了
8bit 的优化器量化技术，安装 bitsandbytes 库，并传参—use_8bit_adam enable 该功能
resolution – 取决于训练集图片分辨率，直接影响显存占用和训练速度，测试下来 resolution 为 512 大约 48GB 显存占用，resolution 为 1024 则会高达 78GB 显存占用，倍数增长
learning_rate – 我们发现设成 1e-5 效果很好，也可以试试介于 1e-4 和 2e-6 之间的其它值
checkpointing_steps – 对于 AWS P5 实例，可以采用 Spot 竞价实例，而采用竞价实例时，机器资源可能会被更高竞价者回收，因此该参数比较重要，一定要设置到 EFS 上

训练资源

其它平台提供商（如 Lamdba Labs）也提供了 GPU 机器的资源，在目前各个平台的对比上，我们选择了 AWS，主要考虑到其稳定性和高性能 GPU 卡资源（如 P5）的 Capacity。

除了刚才提到的 AWS P5 实例外，还有其它训练机器和 GPU 卡可以供选择，如 NVIDIA A10, A100，这需要从性能、时间、费用整体取舍，经过测试对比，P5 在该场景的训练层面，从性价比比 A10，A100 的显卡更加有优势，其测试数据如下：

	A	B	C	D	E	F
1	训练机型	显卡类型	训练时长	训练数据量	训练 epcho	费用
2	ml.g5.48xlarge	A10	10032m	19w	20	$3344
3	ml.p4d.24xlarge	A100	3306m	19w	20	$2052
4	ml.p5.48xlarge	H100	720 m	19w	20	$840

另外我们需要考虑训练方式上多机多卡/多机单卡 / 单机多卡 /单机单卡…etc，目前我们采用单机多卡的方式，主要考虑其任务调度、训练数据分配、进度协调等的复杂度。

训练任务

我们采用 EC2 AMI 自带的默认环境，以便在此基础上通过 MiniConda 环境自主安装 Cuda Cudnn Torch 版本，从而确保在兼容的依赖库版本上拉起 SDXL controlnet 的训练任务，其脚本如下所示：

# Install MiniConda, Create Env and Activate It
wget https://repo.anaconda.com/miniconda/Miniconda3-py311_23.5.2-0-Linux-x86_64.sh
bash Miniconda3-py311_23.5.2-0-Linux-x86_64.sh
conda create -n controlnet_sdxl python=3.11
conda activate controlnet_sdxl

# Install Torch Torchvision First
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

# Clone Train Code
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install -e .

# Login Huggingface & Wandb
huggingface-cli login
pip install wandb
wandb login

# Install Git-LFS
curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get install git-lfs

# Clone Train Dataset
git clone https://huggingface.co/datasets/p5_test/controlnet_fs_dataset_df ./dataset
## Input Twice Username & Pwd

# Train in Tmux
tmux
conda activate condtrolnet_sdxl
bash train_controlnet_sdxl-h100.sh

也可以使用 Amazon SageMaker 拉起 P5 训练任务，其启动方式与 EC2 上基本一致，细微区别在于训练镜像的打包和上传。

关于在 SageMaker 上拉起 P5/P4d 的训练实例方式在这里不再赘述，其示例脚本可以参考：https://github.com/qingyuan18/sm_sample/blob/main/controlnet/stablediffusion_controlnet_finetuning.zh.ipynb。

训练故障解决及调优

根据 Diffuser 官方代码，Controlnet 训练步骤如下：

其中 Map 阶段对训练图像数据做 comput_embedding 操作，这里不是 SD 的 Clip 文本嵌入的 embedding，还有 Controlnet 附加的嵌入，需要 SD XL UNet 模型的参与；然后的 train 阶段才是 Controlnet 核心的训练步骤。

训练过程中遇到各种 Cuda 异常、Torch 异常、OOM 异常，分别出现在 Map 和 Train 阶段，其解决思路和方案参考如下：

思考问题：

各种 Cuda 花式报错很容易将注意力引到 Cuda 版本上，但是相同的 AMI 在 SD 1.5 单机上训练良好，所以推测可能还有别的原因引发这些问题，例如：分布式训练问题、 Torch 问题、Diffuser 训练脚本问题、SDXL 问题，期间还多次出现在 Mapping 阶段卡死的问题。
引发问题的是 Accelrate，是 Diffusers，还是算子？还是 Cuda 版本问题，Torch 兼容性问题，训练脚本问题？

问题排查：

使用同样的训练在 A100 上训练成功，排除训练脚本、SDXL 的问题；
不使用 Accelerate 运行训练脚本，问题复现，所以不是分布式 Accelerate 的问题；
升降级 Torch 版本，高版本可以过 Mapping 过程，低版本会导致 Mapping 过程卡死；
降级 Cuda 版本到8，问题依旧；
通过报错信息，定位到还是 PyTorch 底层算子问题，翻看 Torch 文档，发现有最新 Preview 版本可用，但直接使用 preview 版本，会卡死在 Mapping 阶段；
采用Torch 2.0.1 低版本过 Mapping 阶段算子，再使用 Preview 版本执行后阶段算子，问题解决 !

修改后的训练脚本：

# Install MiniConda, Create Env and Activate It
wget https://repo.anaconda.com/miniconda/Miniconda3-py311_23.5.2-0-Linux-x86_64.sh
bash Miniconda3-py311_23.5.2-0-Linux-x86_64.sh
conda create -n controlnet_sdxl python=3.11
conda activate controlnet_sdxl

# Install Torch Torchvision First
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

# Clone Train Code
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install -e .

# Login Huggingface & Wandb
huggingface-cli login
pip install wandb
wandb login

# Install Git-LFS
curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get install git-lfs

# Clone Train Dataset
git clone https://huggingface.co/datasets/p5_test/controlnet_fs_dataset_df ./dataset
## Input Twice Username & Pwd

# Train in Tmux
tmux
conda activate condtrolnet_sdxl
bash train_controlnet_sdxl-h100.sh

# Run this after 1st raise error
pip uninstall torch torchvision
pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu118
bash train_controlnet_sdxl-h100.sh

训练监控

这里采用 wandb 默认的监控指标，其训练期间 matric 如下：

模型部署

整体部署架构如下所示：

业务部署架构图如下：

部署规则需要匹配业务场景：

部分模型服务访问量不大
部分模型服务访问量弹性过大
部分模型服务需要较快的冷启动速度
部分模型服务推理需要较好的稳定性
整体架构希望可以满足各项应急需要
整体架构希望可以用较低成本来运行

解决方案：

一个服务对应一个计算组，一个计算组由服务对应 SQS，EC2 OD + EC2 Spot + AMI 组成。
EC2 Spot 负责主力推理服务，当 EC2 Spot 被回收时，启动对应 EC2 OD 临时接替推理服务，并启动另外一台 EC2 Spot，等待另一台 EC2 Spot 启动后，对 EC2 OD 进行关机。
国内可以使用 Warm Pool 对 OD 启动进行加速。
Script 控制 OD 开关 + ASG 不支持 EC2 停止状态，独立配置的 Spot Fleet 支持停止状态。
处理任务可以使用从 SQS 拉取任务或者使用 Distribute Server 进行分发任务。
长时间推理任务同样可以分段进行任务分发和回收。
适用于多云计算、混合计算、峰值计算和冷启动计算。

总结

本文介绍了在 AWS P5 机器学习实例上对 SD XL 模型进行 Controlnet 微调训练的方式方法，以及具体实施落地实现的示例。文中脚本代码及笔记本训练示例，可做为用户基于 SD 1.5/2.0 版本到 SD XL 的 AIGC ML 平台的 Controlnet 微调训练工程化落地参考。

亚马逊AWS官方博客