优化kohya-ss/sd-scripts项目中LoRA提取的内存使用问题

在kohya-ss/sd-scripts项目中，用户报告了一个关于LoRA提取过程中内存占用过高的问题。特别是在使用SDXL DreamBooth模型进行128x128 rank常规LoRA提取时，即使在拥有29GB RAM和15GB VRAM的Kaggle环境中也会失败。

问题分析

LoRA提取过程需要同时加载原始模型和经过微调的模型，这导致了较高的内存需求。当处理较大模型如SDXL时，内存需求会显著增加，使得在资源有限的环境中难以完成提取任务。

解决方案

项目维护者kohya-ss在dev分支中实现了以下优化方案：

1. 设备分配优化：通过--load_original_model_to cuda参数，可以将原始模型加载到GPU显存中，而微调模型保留在系统内存中。这种分离加载策略有效降低了单设备的内存压力。

2. 精度控制优化：新增了--load_precision参数，允许用户指定模型加载时的精度。例如：

   • 对于fp16模型，使用--load_precision fp16
   • 对于bf16模型，使用--load_precision bf16

这种精度控制不仅减少了内存使用量，还能保持模型的精度不受损失。

技术细节

1. 内存优化原理：通过将模型分散到不同设备(GPU和CPU)上，避免了单设备内存的峰值压力。同时，降低加载精度可以直接减少模型在内存中的占用空间。

2. 精度选择建议：

   • 当原始模型是fp16格式时，建议使用fp16精度加载
   • 当原始模型是bf16格式时，建议使用bf16精度加载
   • 这样可以确保在减少内存使用的同时不损失模型精度

实际应用

这些优化特别适用于资源受限的环境，如：

• 云服务提供商提供的有限资源实例
• 个人开发者的工作站
• Kaggle等在线平台的计算环境

通过合理配置这些参数，用户可以在不升级硬件的情况下完成更大模型的LoRA提取任务。

结论

kohya-ss/sd-scripts项目通过引入设备分配和精度控制参数，有效解决了LoRA提取过程中的高内存占用问题。这些改进使得在资源有限的环境中处理大型模型成为可能，为更广泛的用户群体提供了便利。