VLMEvalKit项目中InternVL2-Llama3-76B模型多卡部署问题解析

在开源项目VLMEvalKit中使用InternVL2-Llama3-76B这类大模型时，开发者可能会遇到多GPU设备间的张量分配问题。本文将深入分析该问题的成因并提供解决方案。

问题现象

当尝试在VLMEvalKit中运行InternVL2-Llama3-76B模型时，即使设置了device_map='auto'参数，系统仍会报出"Expected all tensors to be on the same device"错误。这表明模型的不同部分被错误地分配到了不同的GPU设备上（如cuda:0和cuda:1），导致张量运算无法正常进行。

问题根源

1. 模型规模因素：76B参数的大模型通常需要跨多张GPU进行分布式部署，简单的auto分配策略可能无法正确处理模型各组件间的依赖关系。

2. 组件特殊性：InternVL这类视觉-语言混合模型包含视觉编码器和语言模型两部分，它们对计算资源的需求不同，需要特别处理。

3. 设备映射不足：标准的auto设备映射策略可能无法识别模型中的关键组件，导致重要模块被分散到不同设备。

解决方案

方案一：使用原始配置

最简单的解决方案是直接使用项目提供的原始运行命令：

python run.py --model InternVL2-76B --data RealWorldQA

这种方式已经内置了合理的设备分配策略。

方案二：手动设备映射

对于需要自定义部署的情况，可以使用accelerate库进行精细化的设备映射：

from accelerate import infer_auto_device_map

# 初始化模型时不指定device_map
model = AutoModel.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    low_cpu_mem_usage=True).eval()

# 获取模型的不可分割模块
no_split_module = model._no_split_modules

# 自定义设备映射
device_map = infer_auto_device_map(
    model,
    max_memory={0: '40GiB', 1: '40GiB', 2: '40GiB', 3: '40GiB'},
    no_split_module_classes=no_split_module)

# 关键组件强制分配到指定设备
device_map['vision_model'] = 0
device_map['mlp1'] = 0
device_map['language_model.model.tok_embeddings'] = 0
device_map['language_model.model.embed_tokens'] = 0
device_map['language_model.output'] = 0
device_map['language_model.model.norm'] = 0
device_map['language_model.lm_head'] = 0

技术要点

1. 内存分配策略：为每个GPU设置40GiB的内存上限，确保资源合理分配。

2. 关键组件固定：将视觉模型、嵌入层和输出层等关键组件固定到同一设备（如cuda:0），避免跨设备通信。

3. 不可分割模块：通过_no_split_modules属性识别模型中不应被分割的组件，保证模型结构的完整性。

最佳实践建议

1. 对于大多数用户，直接使用项目提供的运行命令是最稳妥的选择。

2. 需要自定义部署时，建议先完整加载模型再构建设备映射，这样可以更准确地评估各组件资源需求。

3. 在实际部署前，建议先在小规模数据上测试设备映射方案的有效性。

4. 根据实际GPU配置调整max_memory参数，避免内存不足或资源浪费。

通过以上方法，开发者可以有效地解决VLMEvalKit中大模型的多GPU部署问题，确保模型能够高效稳定地运行。