踩坑实录：微调过程中遇到的显存溢出和内存泄漏

踩坑实录：微调过程中遇到的显存溢出和内存泄漏

最近在尝试对LLM进行LoRA微调时，遇到了严重的显存溢出问题，特此记录。

环境配置

• PyTorch 2.1.0
• transformers 4.33.0
• accelerate 0.21.0
• CUDA 11.8
• 8GB显存RTX 3090

问题复现步骤

使用以下代码进行LoRA微调时出现显存溢出：

from peft import get_peft_model, LoraConfig
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
peft_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.01,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, peft_config)

解决方案

1. 降低LoRA秩r值：将r=8改为r=4，显存使用从7.2GB降至5.8GB
2. 启用梯度检查点：

model.gradient_checkpointing_enable()

3. 设置gradient_accumulation_steps=2
4. 使用混合精度训练：

training_args = TrainingArguments(
    gradient_accumulation_steps=2,
    fp16=True,
    ...)  

内存泄漏排查

发现训练过程中出现内存泄漏，主要原因是未正确释放中间变量。解决方案是：

• 使用torch.cuda.empty_cache()清理缓存
• 避免在循环中累积tensor
• 定期检查显存使用情况

最终通过以上优化，成功完成了微调任务。