VLMEvalKit项目中InternVL-76B模型推理显存优化实践

背景介绍

在大型视觉语言模型(VLM)的推理过程中，显存管理是一个关键挑战。VLMEvalKit作为一个开源评估工具包，支持对InternVL等大型视觉语言模型进行评测。本文针对InternVL-76B模型在MMMU数据集评测过程中遇到的显存问题进行分析，并提供解决方案。

问题现象

当使用8张40GB显存的GPU对InternVL-76B微调后的模型进行评测时，系统在推理几个问题后出现CUDA out of memory错误。即使升级到8张80GB显存的GPU，仍然会遇到"Expected all tensors to be on the same device"的错误提示。

原因分析

1. 模型规模与显存需求：InternVL-76B作为760亿参数规模的模型，其显存占用非常庞大。即使在多GPU环境下，也需要精细的显存管理。

2. 推理过程中的动态显存增长：MMMU数据集推理时生成的回复较长，导致显存需求随推理过程动态增加。

3. 模型切分问题：当未正确配置模型切分时，部分模型组件可能被错误地放置在CPU上，导致设备不匹配错误。

解决方案

显存优化策略

1. 启用自动模型切分： 在运行命令前设置环境变量AUTO_SPLIT=1，确保模型被正确切分到多个GPU上。这是VLMEvalKit中split_model函数的关键配置。

2. 设备映射验证： 检查internvl init中的切分步骤，打印输出device_map和visible_devices信息，确保与本地GPU配置一致。

3. 显存监控： 使用nvidia-smi工具实时监控显存使用情况，及时发现潜在的显存瓶颈。

最佳实践建议

1. 硬件配置：

   • 对于InternVL-76B模型，推荐使用8张80GB显存的GPU进行推理
   • 确保GPU间有高速互连(NVLink)以减少通信开销

2. 参数调整：

   • 适当限制生成文本的最大长度
   • 调整批处理大小(batch size)以平衡显存使用和推理效率

3. 软件配置：

   • 使用最新版本的深度学习框架(如PyTorch)
   • 启用混合精度推理(AMP)以减少显存占用

实施效果

通过正确配置模型切分和显存管理策略，InternVL-76B模型可以在多GPU环境下稳定运行，完成MMMU等复杂数据集的评测任务。监控数据显示各GPU显存使用均衡，避免了单卡过载的情况。

总结

大型视觉语言模型的推理需要综合考虑模型规模、硬件配置和软件优化。VLMEvalKit提供的模型切分功能是解决显存问题的有效手段，但需要正确配置才能发挥作用。未来随着模型规模的持续增大，显存优化技术将变得更加重要。