CogVideo项目在NPU设备上的优化器选择与多卡训练方案

背景介绍

在深度学习模型训练过程中，优化器的选择对模型性能有着重要影响。CogVideo作为THUDM开发的大规模视频生成模型，其训练过程通常使用FusedEmaAdam优化器以获得更好的训练效果。然而，当我们将训练环境迁移到NPU设备时，会遇到优化器兼容性问题。

NPU设备上的优化器替代方案

FusedEmaAdam优化器依赖于CUDA环境，这在NPU设备上不可用。针对这一问题，技术团队给出了明确的解决方案：

1. AdamW优化器替代方案：在NPU设备上，可以直接使用标准的AdamW优化器替代FusedEmaAdam。AdamW是Adam优化器的改进版本，加入了权重衰减正则化，在大多数情况下能够提供相似的训练效果。

2. 实现差异：虽然FusedEmaAdam通过CUDA实现了算子融合加速，但在NPU环境下，使用AdamW虽然可能损失部分性能优势，但能够保证训练的正常进行。

多NPU卡训练方案

对于拥有多张NPU卡的用户，技术团队还提供了多卡训练的建议方案：

1. 分布式训练配置：需要正确设置分布式训练的环境变量，包括：

   • WORLD_SIZE：全局进程数量
   • RANK：当前进程的全局排名
   • LOCAL_RANK：当前节点上的进程排名
   • LOCAL_WORLD_SIZE：当前节点上的进程数量

2. 启动方式：技术团队表示将很快提供专门的多卡训练脚本，用户可以通过该脚本自动处理多卡训练的配置问题。

实践建议

对于希望在NPU设备上训练CogVideo模型的用户，建议：

1. 首先将优化器替换为AdamW
2. 等待官方发布多卡训练脚本，或根据分布式训练原理自行配置多卡环境
3. 监控训练过程中的性能指标，必要时调整学习率等超参数

通过以上方案，用户可以在NPU设备上顺利完成CogVideo模型的训练任务，虽然可能损失部分CUDA优化带来的性能优势，但整体训练流程可以得到保证。