Kohya_SS项目中SDXL模型训练的显存优化方案

在Kohya_SS项目的实际应用中，许多用户反馈在使用16GB显存的GPU进行SDXL 1.0模型的DreamBooth训练时会遇到CUDA内存不足的问题。经过技术验证和参数调优，我们发现通过特定的配置组合可以突破这一硬件限制。

显存瓶颈分析

SDXL模型因其庞大的参数量（特别是包含双文本编码器）对显存需求较高。典型训练场景下：

• 完整训练（UNet+文本编码器）需要超过19GB显存
• 仅UNet训练时显存需求可降至16GB范围内
• 默认参数下batch size即使设为1也难以满足

关键技术优化方案

核心参数配置

1. 训练模式选择

   • 必须采用UNet-only训练模式
   • 将文本编码器学习率(TE1/TE2)设置为0

2. 显存优化技术

   • 启用梯度检查点(Gradient Checkpointing)
   • 使用全FP16或BF16精度训练
   • 激活文本编码器输出缓存(Cache Text Encoder Outputs)

3. 训练参数调整

   • 训练批量大小(Train batch size)必须保持为1
   • 推荐使用Adafactor或Lion8bit优化器

特殊参数说明

对于GUI界面未显示"Cache text encoder outputs"选项的情况，需要在附加参数栏手动添加： --cache_text_encoder_outputs

实施建议

1. 建议从基础配置开始逐步增加参数
2. 监控显存使用情况可使用nvidia-smi工具
3. 训练过程中注意观察loss曲线稳定性
4. 可尝试混合精度训练进一步优化

注意事项

该方案虽然能在16GB显存下运行，但仍属于临界状态。建议：

• 关闭所有非必要后台进程
• 降低图像分辨率至模型允许的最小值
• 准备应对可能的训练中断情况

通过以上优化措施，开发者可以在有限的硬件资源下完成SDXL模型的微调训练，为资源受限环境下的模型优化提供了可行方案。