CogVideo模型LoRA微调中的NaN损失问题分析与解决

问题背景

在使用CogVideo项目进行文本到视频(T2V)模型的LoRA微调过程中，部分开发者遇到了训练过程中出现NaN(非数值)损失值的问题。具体表现为训练初期损失值突然变为NaN，同时伴随梯度范数(grad_norm)也变为NaN的情况。这类问题在深度学习模型训练中并不罕见，但需要针对具体模型和场景进行分析。

问题现象

在CogVideo 5B模型的LoRA微调过程中，开发者观察到以下典型现象：

• 训练开始后不久(约23步后)，损失值突然变为NaN
• 梯度范数同样显示为NaN
• 系统日志中出现"NaN or Inf found in input tensor"的警告信息
• 该问题在使用默认训练脚本(train_ddp_t2v.sh)和参数配置时出现

可能原因分析

1. 浮点精度问题

CogVideo这类大型生成模型对数值精度较为敏感。虽然开发者使用了BF16混合精度训练(这是官方推荐的配置)，但仍需注意：

• 某些操作在低精度下可能产生数值不稳定
• 梯度累积过程中可能放大数值误差

2. 数据质量问题

经过深入排查，发现问题根源在于训练数据：

• 输入数据包含异常值(如NaN或Inf)
• 数据预处理不充分导致数值范围异常
• 视频帧数据格式或范围不符合模型预期

3. 学习率设置不当

对于5B参数量的模型：

• 初始学习率过高可能导致梯度爆炸
• LoRA适配器的学习率需要特别调整

4. 损失函数数值稳定性

视频生成任务的损失计算可能涉及：

• 多尺度特征匹配
• 时序一致性约束
• 复杂的正则化项

解决方案

数据质量检查与修复

1. 数据清洗：

   • 检查训练数据集中是否包含损坏的视频文件
   • 验证所有视频帧的像素值是否在合理范围内(如[0,255]或[0,1])
   • 移除或修复包含NaN/Inf值的数据样本

2. 数据预处理规范化：

   • 确保数据预处理流程与原始训练时一致
   • 对输入数据进行标准化处理(如归一化到[-1,1]或[0,1])
   • 检查数据增强操作是否引入异常值

训练配置优化

1. 学习率调整：

   • 尝试降低初始学习率
   • 使用学习率预热策略
   • 对LoRA参数和主干网络采用不同的学习率

2. 梯度裁剪：

   • 添加梯度裁剪防止梯度爆炸
   • 设置合理的最大梯度范数阈值

3. 数值稳定性增强：

   • 检查模型中是否存在不稳定的数学操作(如除法、指数运算)
   • 在关键计算处添加数值保护(如添加微小epsilon防止除零)

监控与调试

1. 训练过程监控：

   • 在训练初期增加验证频率
   • 监控梯度统计量(均值、方差、最大值等)
   • 记录各层激活值的分布情况

2. 逐步排查法：

   • 先在小型数据集上验证训练稳定性
   • 逐步增加数据量和模型复杂度
   • 使用梯度检查点技术减少内存压力

最佳实践建议

1. 数据准备阶段：

   • 实现数据质量检查脚本
   • 建立标准化的数据预处理流水线
   • 保留数据预处理中间结果以便调试

2. 模型训练阶段：

   • 从小规模实验开始验证配置
   • 使用稳定的优化器(如AdamW)
   • 实现自动异常检测和恢复机制

3. 调试工具：

   • 利用PyTorch的autograd异常检测功能
   • 实现自定义的数值检查回调函数
   • 在关键计算节点添加断言检查

总结

CogVideo等大型生成模型的微调过程中出现NaN损失通常是多方面因素共同作用的结果。通过系统性的数据检查、训练配置优化和严格的数值稳定性控制，可以有效解决这类问题。特别需要注意的是，对于视频生成任务，数据质量往往是最关键的因素，建立完善的数据质量保障机制可以显著提高训练成功率。