Wenet项目中FP16训练出现Loss为NaN问题的分析与解决

在语音识别领域的Wenet项目实践中，使用混合精度训练时可能会遇到一个典型问题：当训练迭代到一定步数（如1万步左右）后，损失函数（loss）开始频繁出现NaN（非数值）的情况。这种现象在使用DeepSpeed Stage2优化器和FP16（半精度浮点数）模式时尤为常见。

问题本质分析

FP16训练出现NaN的根本原因在于半精度浮点数的数值表示范围有限（约±65,504），当模型在训练过程中产生过大的梯度值时，很容易超出这个范围，导致数值溢出（overflow）。这种数值不稳定现象通常表现为：

1. 训练初期可能正常，随着模型参数更新逐渐出现异常
2. 损失函数突然变为NaN或无限大
3. 模型性能急剧下降

解决方案

首选方案：使用BF16格式

对于支持BF16（Brain Floating Point）的硬件（如较新的NVIDIA GPU），BF16是更好的选择。BF16具有与FP32相同的指数位（8位），但尾数位较少（7位），这使得它能够：

• 保持与FP32相似的数值范围
• 减少内存占用
• 避免FP16常见的数值溢出问题

次优方案：FP32全精度训练

在不支持BF16的硬件上（如NVIDIA V100），可以采用的解决方案是：

1. 全程使用FP32训练：虽然会牺牲一些训练速度和内存效率，但能保证数值稳定性
2. 混合精度策略：先使用FP32训练模型至相对稳定状态，再切换到FP16进行微调
   • 初期FP32阶段：帮助模型找到相对稳定的参数空间
   • 后期FP16阶段：加速训练并减少内存占用

实践建议

1. 梯度裁剪：即使使用FP32，适当设置梯度裁剪阈值（如1.0或5.0）有助于防止梯度爆炸
2. 学习率调整：出现NaN时，尝试降低学习率或使用学习率warmup策略
3. 损失监控：实现NaN检测机制，当检测到NaN时自动回滚到上一个有效检查点
4. 混合精度配置：如果使用混合精度，确保正确设置了loss scaling参数

总结

在Wenet等语音识别模型的训练中，数值稳定性是成功训练的关键因素之一。根据硬件条件选择合适的浮点精度格式，配合适当的训练策略，可以有效解决FP16训练中出现的NaN问题，确保模型训练的稳定性和最终性能。