PEFT微调Gemma模型时训练损失NaN问题的分析与解决

问题背景

在使用Hugging Face的PEFT(Parameter-Efficient Fine-Tuning)库对Gemma 7B模型进行4位量化微调时，开发者遇到了训练过程中损失值变为NaN的问题。具体表现为：训练初期损失值正常下降，但经过几个训练步骤后突然变为NaN，同时验证损失也显示为NaN。

技术细节分析

该问题出现在以下技术配置环境中：

1. 模型量化配置：使用4位NF(正态浮点)量化类型，计算数据类型为bfloat16
2. LoRA适配器配置：针对q_proj、v_proj、embed_tokens和lm_head模块进行适配
3. 词汇表扩展：通过扩展SentencePiece模型而非标准add_tokens方法增加新词元
4. 训练参数：使用paged_adamw_8bit优化器，学习率2e-4，bf16精度训练

可能原因分析

1. 数值精度问题：在4位量化基础上使用bfloat16可能导致数值精度不足，梯度计算时出现数值下溢或上溢
2. 新词元初始化：将所有新词元的嵌入权重初始化为零可能造成训练不稳定
3. 混合精度训练：bf16与4位量化的组合可能引入数值不稳定性
4. 梯度检查点：启用的梯度检查点功能可能与量化训练存在兼容性问题

解决方案

PEFT核心开发者Benjamin Bossan提出的解决方案是：

model = get_peft_model(...)
# 将所有可训练参数转换为float32精度
for param in model.parameters():
    if param.requires_grad:
        param.data = param.data.float()

这一方案的关键点在于：

1. 精度提升：将LoRA适配器的可训练参数从量化/低精度状态转换为float32，确保训练稳定性
2. 选择性转换：仅转换需要梯度的参数，保持其他参数的量化状态，平衡精度与效率
3. 兼容性保持：不影响原始量化模型的推理效率，仅在训练阶段使用更高精度

最佳实践建议

1. 渐进式词汇表扩展：考虑使用更温和的新词元初始化策略，如小随机数初始化而非全零
2. 精度监控：在训练初期添加梯度范数监控，提前发现数值不稳定迹象
3. 学习率调整：对于量化模型，可能需要更保守的学习率策略
4. 混合精度配置：可以尝试禁用bf16或使用更稳定的混合精度组合

结论

在PEFT框架下微调大型量化语言模型时，数值精度管理是关键挑战。通过将适配器参数临时提升到float32精度，可以在保持模型整体量化效率的同时，确保训练过程的数值稳定性。这一解决方案既保留了PEFT的参数效率优势，又解决了量化训练中的数值不稳定问题，为类似场景提供了可靠的技术路径。