用 peft 0.13 加 transformers 4.45 微调 Llama-3.1-8B,LoRA rank 16,batch 4 加梯度累积 4,24GB 卡在第一个 step 就 OOM,把 batch 降到 1 能跑。

真正麻烦的地方在于,几个原因的现象几乎一样。只算了权重显存,没算优化器状态与激活值;序列长度 4096 时激活值占用远超权重;梯度检查点没开,中间激活全部保留;LoRA 层用了 float32 而基座是 bf16。

量出模型权重加载后的显存基线

python -c "import torch;print(torch.cuda.mem_get_info()[0]/2**30)"

再往下打开梯度检查点并缩短序列长度重试

python train.py --gradient-checkpointing --max-len 2048 --batch 2 --accum 8

接着按参数量粗算各部分的显存占用

python -c "p=8e9;print('weight',p*2/2**30,'adam',p*8/2**30)"

等前面都确认过,再用 4bit 量化基座加 LoRA 再跑一次

python train.py --load-in-4bit --lora-r 16 --batch 4 --max-len 2048

同事提醒我看一下最近一次发布,时间点确实对得上。

命令输出我都留了记录,免得改完记不清原来的值。

改之前先确认了最近有没有人动过配置,避免白查一遍。

改完先别撤,确认这几项。在 24GB 卡上 batch 4 能稳定跑完一个 epoch,显存峰值低于 22GB,loss 曲线正常下降。

LoRA 只省权重和优化器,激活值仍按序列长度平方增长,估算显存必须把激活单独算一遍。拖得越久,排查成本越高,因为现场状态会被后续操作覆盖。

作者 9823189