解决lmms-eval项目中视频数据集后处理时的OOM问题

背景介绍

在评估大型多模态模型Oryx-34B时，研究人员遇到了一个典型的内存问题。当使用8块H100 GPU在VideoMME数据集上进行评估时，模型推理阶段("Model Responding")运行正常，但在后处理阶段(Postprocessing)完成后出现了内存不足(OOM)错误，具体表现为NCCL警告"Cuda failure 2 'out of memory'"。

问题分析

这种内存问题通常发生在以下几个环节：

1. 指标计算阶段：在多GPU环境下，当各个GPU的计算结果需要汇总时，可能会产生大量的临时内存占用
2. 数据收集阶段：从多个GPU收集计算结果时，如果没有合理的内存管理策略，容易导致内存峰值
3. 批处理设置不当：虽然模型推理阶段正常，但后处理阶段可能没有考虑到多GPU环境下的内存分配

特别值得注意的是，这个问题在单GPU环境下不会出现，说明它与多GPU通信和内存管理密切相关。

解决方案

针对这个问题，项目团队提出了有效的解决方案：

1. 优化多GPU通信：改进了NCCL通信模式，减少了内存占用
2. 分批处理指标计算：将大型计算结果分批处理，避免一次性占用过多内存
3. 内存管理改进：实现了更智能的内存分配和释放策略

技术实现细节

该解决方案的核心在于重新设计了后处理阶段的内存管理策略：

• 采用异步通信模式，减少内存峰值
• 实现计算结果的分批聚合，避免一次性加载所有数据
• 优化了GPU间的数据传输机制

实际效果

经过这些优化后，系统能够：

• 在8块H100 GPU上顺利完成整个评估流程
• 保持原有的计算准确性
• 显著降低内存使用峰值

总结

这个案例展示了在多GPU环境下评估大型多模态模型时可能遇到的内存管理挑战。通过优化后处理阶段的内存使用策略，特别是改进多GPU通信和计算结果聚合方式，可以有效解决这类OOM问题。这对于其他类似的大规模多模态模型评估工作具有参考价值。