解决sd-scripts项目中全量微调时保存检查点的内存溢出问题

问题背景

在使用kohya-ss/sd-scripts项目进行Stable Diffusion模型的全量微调(full finetuning)时，许多用户在RTX 3090等高端显卡上训练过程可以正常进行，但在保存检查点(checkpoint)时却会遇到CPU内存不足(OOM)的问题。这主要是因为保存完整模型权重时需要将GPU上的参数全部加载到CPU内存中，对于大型模型来说，32GB的系统内存可能不足以容纳整个模型。

技术分析

当使用全量微调而非LoRA微调时，模型的所有参数都会被更新，这使得保存检查点时的内存需求大大增加。特别是在使用bf16精度训练时，虽然减少了GPU显存的占用，但在保存模型时仍需要将整个模型的状态加载到CPU内存中进行序列化。

解决方案

项目维护者提供了一个有效的解决方案：使用--mem_eff_save参数。这个参数的作用是：

1. 使用自定义的模型保存实现替代默认的safetensors库
2. 优化内存使用方式，减少保存过程中的峰值内存消耗
3. 保持与标准保存方法相同的功能性和模型完整性

实施建议

对于使用全量微调的用户，建议在训练命令中添加以下参数组合：

--save_model_as safetensors \
--mem_eff_save \
--mixed_precision bf16 \
--save_precision bf16

这种配置可以：

• 保持bf16精度的训练和保存
• 使用更高效的内存管理方式保存模型
• 最终输出仍为safetensors格式，保持兼容性

注意事项

1. 即使使用了内存优化选项，在保存大型模型时仍建议确保系统有足够的内存余量
2. 对于极端情况，可以考虑减少--save_every_n_steps的频率来降低保存频率
3. 监控系统内存使用情况，特别是在长时间训练过程中

通过采用这些优化措施，用户可以在资源有限的硬件配置上顺利完成Stable Diffusion模型的全量微调任务。