低显存环境下训练 LoRA 权重出现梯度爆炸的解决方法探讨

显存压榨到极致时，LoRA 训练最容易在 Epoch 2-3 左右突然出现 Loss 变成 NaN 的情况，这通常不是学习率设高了，而是低显存环境下混合精度计算（FP16）带来的数值溢出。

实测在 24G 显存以下跑 Llama-3 或 SDXL 的 LoRA 时，如果开启了 fp16 且使用了较高的 learning_rate，梯度在反向传播时很容易在某些层瞬间激增。我对比了三种方案的实测表现：

方案一：切换到 BF16（最推荐）
如果显卡是 30 系列或 40 系列，直接把精度从 fp16 换成 bf16。BF16 的动态范围和 FP32 一致，能极大程度避免梯度爆炸。
实测结果： 同样 5e-5 的学习率，FP16 在第 200 步崩了，BF16 稳跑 2000 步且 Loss 下降曲线更平滑。

方案二：引入梯度裁剪（Gradient Clipping）
在训练配置文件或脚本中加入 max_grad_norm。这个参数相当于给梯度设了个“天花板”，超过这个值的梯度会被强制压缩。

# 在训练参数中添加
--max_grad_norm 1.0

实测结果： 开启后虽然 Loss 下降速度稍微变慢，但有效解决了训练中期的突发性崩盘，适合那些必须在老旧显卡（如 2080Ti）上跑 FP16 的场景。

方案三：调整 Optimizer 到 8-bit AdamW
低显存环境下，显存碎片化严重，标准的 AdamW 占用太多内存。换成 bitsandbytes 提供的 8-bit 优化器，不仅省显存，在某些特定权重初始化下，数值稳定性反而有所提升。

# 替换优化器配置
optim="adamw_8bit"

对比总结：

• 稳定性： BF16 > 梯度裁剪 > 8-bit AdamW
• 显存节省： 8-bit AdamW > BF16 ≈ 梯度裁剪
• 收敛速度： BF16 > 梯度裁剪 > 8-bit AdamW

对于绝大多数用户，解决梯度爆炸的最快路径是：优先检查硬件是否支持 BF16 → 开启 max_grad_norm 1.0 → 降低 learning_rate 一个数量级 → 检查数据集是否有异常脏数据（如极长文本导致的 Padding 溢出）。