Ludwig项目中Mistral-7B模型训练出现NaN损失值问题的分析与解决方案

问题背景

在使用Ludwig框架对Mistral-7B模型进行LoRA微调时，当训练批次大小(batch_size)设置为大于1或"auto"时，会出现损失值变为NaN的情况，导致训练在第一轮结束时失败。而当batch_size设置为1时，虽然可以避免这个问题，但会导致GPU利用率低下，特别是对于像Tesla V100 32GB这样的高性能GPU来说尤其明显。

问题现象

在训练过程中，系统会检测到模型权重中出现NaN或无限大值，具体报错信息显示在模型的自注意力层(self_attn)的查询投影层(q_proj)的LoRA适配器权重中发现了这些异常值。训练会因此被强制终止，且无法保存有效的检查点。

根本原因分析

经过深入调查，这个问题实际上是Mistral模型的一个已知问题，并非Ludwig框架特有的。当使用4位量化(bits=4)对Mistral模型进行训练时，较大的批次大小会导致在反向传播过程中出现数值溢出，特别是在使用float16计算精度时。

解决方案

推荐解决方案

对于支持bfloat16数据类型的GPU(如Ampere架构及更新的NVIDIA GPU，包括A100、A5000、A6000等)，可以通过在量化配置中添加以下参数来解决：

quantization:
  bits: 4
  bnb_4bit_compute_dtype: bfloat16

bfloat16具有与float32相同的指数位数(8位)，但尾数位数较少(7位)，这使得它能够更好地处理大数值范围，有效防止训练过程中的数值溢出问题。

替代方案

对于不支持bfloat16的GPU(如Tesla V100等Volta架构GPU)，可以考虑以下替代方案：

1. 使用Llama-2-7B模型替代：Llama-2系列模型在相同条件下不会出现这个问题，可以作为Mistral的替代选择。

2. 保持batch_size=1：虽然效率较低，但可以确保训练稳定进行。

3. 使用云服务：考虑使用支持Ampere架构GPU的云服务提供商，如AWS的A5000实例或Predibase的A10G实例。

技术建议

1. 硬件选择：对于大规模语言模型微调，推荐使用Ampere架构或更新的NVIDIA GPU，以获得更好的数值稳定性和训练效率。

2. 监控机制：在训练过程中实现数值稳定性监控，可以在问题出现早期进行干预。

3. 梯度裁剪：虽然不能完全解决这个问题，但适当的梯度裁剪可以帮助提高训练稳定性。

4. 学习率调整：尝试降低学习率，可能有助于缓解数值不稳定的问题。

总结

Mistral-7B模型在4位量化训练时对数值精度较为敏感，特别是在不支持bfloat16的GPU上。开发者需要根据自身硬件条件选择合适的解决方案，权衡训练效率和稳定性。随着硬件的发展和模型优化技术的进步，这类数值稳定性问题有望得到更好的解决。