VLMEvalKit中CogVLM2多卡推理优化方案解析

问题背景

在VLMEvalKit项目中使用CogVLM2-Llama3-Chat-19B模型进行推理时，当模型大小超过单卡显存容量时，开发者尝试通过设置device_map="auto"参数来实现自动多卡分配，但遇到了"Expected all tensors to be on the same device"的错误提示，表明模型部分组件被错误地分配到了不同设备上。

问题分析

该问题源于CogVLM2模型的特殊架构，特别是其视觉组件与语言模型的结合方式。当使用标准的device_map="auto"参数时，Hugging Face的自动设备映射机制无法正确处理模型中某些特定模块的设备分配，导致部分张量被分散到不同GPU上，引发运行时错误。

解决方案

经过深入分析，我们开发了一套定制化的设备映射方案，主要包含以下几个关键点：

1. 显存精确计算：通过get_memory()函数获取每张GPU的实际可用显存，考虑保留一定比例(alpha参数)作为缓冲。

2. 分布式环境适配：考虑多机多卡场景下的设备映射，通过get_rank_and_world_size()获取当前进程的rank和world_size。

3. 模块化设备分配：

   • 为视觉组件(model.vision)指定固定设备
   • 为EVA2CLIPModel等关键模块设置不分割标记
   • 根据计算出的显存容量分配模型各部分

4. 空权重初始化：使用init_empty_weights()上下文管理器先构建模型结构，再进行设备分配，减少内存占用。

实现细节

核心实现分为三个部分：

1. 显存检测与分配：

def get_memory():
    total_memory = torch.cuda.get_device_properties(0).total_memory
    return total_memory / 1024 / 1024 / 1024  # 转换为GB

2. 设备映射构建：

def build_device_map(model, default_map=None, no_split=None, alpha=0.97, beta=0.9):
    # 计算每GPU可用显存
    per_gpu_mem = get_memory() * alpha
    memory_map = {rank: f'{beta * per_gpu_mem:.2f}GiB'}
    
    # 构建完整设备映射
    device_map = infer_auto_device_map(
        model,
        max_memory=memory_map,
        no_split_module_classes=no_split_module
    )
    
    # 应用设备映射
    return dispatch_model(model, device_map=device_map).eval()

3. 模型加载与分配：

with init_empty_weights():
    model = AutoModelForCausalLM.from_pretrained(...)
model, _ = build_device_map(model, ['model.vision'], ["EVA2CLIPModel"])

技术要点

1. alpha参数：控制显存使用率(默认0.97)，避免OOM
2. beta参数：主卡显存保留比例(默认0.9)，确保系统稳定性
3. no_split_module：防止关键模块被分割到不同设备
4. default_map：强制特定模块分配到指定设备

应用效果

该方案成功解决了CogVLM2大模型在多GPU环境下的部署问题，具有以下优势：

1. 精确控制各GPU显存使用量
2. 确保模型关键组件位于合理设备
3. 支持复杂模型架构的特殊分配需求
4. 兼容分布式训练环境

总结

对于VLMEvalKit中的大模型推理任务，特别是像CogVLM2这样的多模态模型，标准的设备自动分配机制往往不能满足需求。通过定制化的设备映射方案，开发者可以更精细地控制模型各部分的设备分配，实现高效稳定的多卡推理。这一方案不仅适用于CogVLM2，也可为其他复杂架构的大模型部署提供参考。