VLMEvalKit项目中InternVL2-8B模型显存优化方案解析

在开源项目VLMEvalKit中，研究人员发现使用InternVL2-8B模型进行多模态评估时存在显存不足的问题。本文将深入分析该问题的技术背景、解决方案及其实现原理。

问题背景

当在8张24GB显存的显卡上运行InternVL2-8B模型评估MMMU_DEV_VAL数据集时，系统会随机出现显存不足(OOM)错误。这种情况通常发生在模型中间特征(feature)数据申请显存时，表明现有的显存分配机制存在优化空间。

技术挑战

大语言模型如InternVL2-8B具有数十亿参数，即使采用模型并行技术将模型切分到多张显卡上，在处理高分辨率图像输入时，中间特征图仍可能消耗大量显存。特别是在多模态评估场景下，需要同时处理视觉和语言信息，显存需求更为突出。

解决方案

项目团队通过PR #566引入了自动切分功能，用户只需设置环境变量AUTO_SPLIT=1，系统便会自动将模型均匀分配到可用显卡上。该方案的创新点包括：

1. 动态负载均衡：系统根据可用GPU数量和进程数自动计算每个模型的分配策略
2. 显存预分配：优化了中间特征的显存管理机制，减少碎片化
3. 简化部署：用户无需手动指定模型切分方案

使用示例

在实际应用中，用户可以通过以下命令启用自动切分功能：

AUTO_SPLIT=1 python run.py --model InternVL2-8B --data MMMU_DEV_VAL

测试表明，该方案成功在两张RTX 4090显卡(各24GB显存)上稳定运行InternVL2-8B模型完成MMMU_DEV_VAL数据集的评估。

技术展望

这一优化不仅解决了当前模型的显存问题，还为未来更大规模的多模态模型评估提供了可扩展的解决方案。随着模型规模的持续增长，类似的显存优化技术将变得越来越重要。项目团队表示将持续完善相关文档，帮助用户更好地利用这一功能。