CogVideo项目中的视频日志内存溢出问题分析与解决方案

问题背景

在CogVideo项目的SFT(监督式微调)训练过程中，当尝试记录训练视频到WandB平台时，出现了显存不足(OOM)的问题。这个问题特别发生在使用H100(80GB)显卡进行5B模型的微调过程中。

问题现象

用户报告在将sft.yaml配置文件中的only_log_video_latents参数设置为False并启用WandB日志记录后，系统抛出显存不足错误。错误信息显示，尽管GPU总容量为79.11GB，当前空闲6.68GB，但PyTorch尝试分配8.08GB时失败。

技术分析

1. 显存需求分析：5B参数的CogVideo模型在训练时本身就需要大量显存，特别是在解码视频帧时会产生额外的显存开销。

2. 配置问题：原配置是为16块H100显卡设计的分布式训练环境，当在单卡环境下运行时，显存明显不足。

3. 视频解码开销：当only_log_video_latents设为False时，系统会解码完整的视频帧用于日志记录，这一过程需要大量显存。

解决方案

1. 使用Diffusers版本：官方建议使用Diffusers版本的实现，通过cogvideox-factory进行训练，这能更好地管理显存使用。

2. 保持仅记录潜在空间：将only_log_video_latents参数保持为True，只记录视频的潜在表示而非完整解码帧。

3. 分布式训练：如果必须使用完整视频日志，建议在多GPU环境下运行，以分摊显存压力。

4. 显存优化配置：可以尝试设置PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True环境变量来减少显存碎片。

最佳实践建议

对于CogVideo这类大规模视频生成模型的训练，建议：

1. 始终在分布式环境下进行训练，特别是对于5B及以上规模的模型。

2. 谨慎使用视频日志功能，评估其对训练稳定性的影响。

3. 考虑使用梯度检查点等技术来优化显存使用。

4. 在单卡环境下，优先考虑使用模型并行或参数高效的微调方法。

通过以上分析和建议，研究人员可以更有效地在资源受限的环境下进行CogVideo模型的训练和调试工作。