MLX-Examples项目中Phi-3模型微调遇到的Loss NaN问题解析

问题背景

在使用MLX-Examples项目对Phi-3-mini-128k-instruct-8bit模型进行微调(LoRA)时，开发者遇到了训练过程中损失值(Loss)变为NaN(非数字)的问题。具体表现为训练开始时损失值就显示为NaN，同时验证损失也显示为NaN，但模型生成功能正常。

问题表现

训练日志显示：

• 可训练参数占比极低(0.042%)
• 第一次迭代验证损失即为NaN
• 第10次迭代训练损失仍为NaN
• 训练吞吐量正常(约530 tokens/秒)

问题根源

经过MLX技术团队分析，这个问题源于MLX框架中的一个量化(quantization)问题。具体来说，当量化过程中遇到全零张量时，会导致NaN值的产生。这个问题在Phi-3模型上尤为明显。

解决方案

MLX技术团队已经修复了这个问题，修复内容包含在MLX框架的PR #1028中。解决方案包括：

1. 修复量化全零张量时产生NaN的问题
2. 需要重新使用最新版MLX对Phi-3原始权重进行量化

对于用户来说，有两种选择：

• 从源码构建最新版MLX
• 等待官方发布包含此修复的新版本

验证结果

根据用户反馈，在应用此修复后，Phi-3模型的LoRA微调已能正常进行，不再出现Loss NaN的问题。

技术启示

1. 量化过程中的边界条件处理非常重要，特别是对于全零张量的处理
2. 模型微调时出现NaN通常表明数值计算存在问题，可能是框架层面的bug
3. 不同模型架构可能暴露框架中的不同问题，Phi-3的特殊性揭示了MLX量化模块的这一问题

最佳实践建议

对于使用MLX-Examples进行模型微调的开发者：

1. 确保使用最新版本的MLX框架
2. 对于Phi-3等新模型，建议从原始权重重新量化
3. 训练初期出现NaN时，首先考虑框架层面的问题而非超参数设置
4. 关注项目更新日志，及时获取问题修复信息