使用多GPU调用AllenNLP共指消解API时的CUDA内存问题求解

1. 确保多GPU并行正确启用 #

• AllenNLP不会自动启用多GPU，需手动配置：
  • 使用命令行工具时，执行allennlp predict或训练命令时添加--cuda-device all，或指定具体GPU编号如--cuda-device 0,1。
  • 自定义代码中，用torch.nn.DataParallel或DistributedDataParallel包装共指消解模型，确保模型参数分布到所有GPU上。
• 用nvidia-smi实时监控GPU内存占用，确认所有目标GPU都有负载。

2. 长文档内存优化方案 #

2000词的长文档会大幅增加显存占用，可通过以下方式缓解：

• 文档分块处理：将长文档切割为500-1000词的子块，分别处理后合并指代结果。切割时尽量保留句子完整性，避免打断跨句指代关系。
• 降低单GPU批次大小：即使多GPU并行，单GPU的batch size过大仍会爆显存，建议将batch size设为1-2，根据GPU内存容量调整。
• 切换轻量预训练模型：将模型底层的大尺寸预训练模型（如bert-large）替换为bert-base，AllenNLP的共指消解配置支持修改model.text_encoder参数实现这一点。

3. 排查并行化失效原因 #

• 检查设备分配：打印model.device确认模型是否分布在多GPU，若显示仅cuda:0，说明并行化未生效，需检查模型包装逻辑。
• 消除硬编码设备：避免代码中出现torch.cuda.device('cuda:0')这类硬编码，确保数据和模型都能自动分配到可用GPU。
• 优化数据加载：设置DatasetReader的num_workers参数为大于1的值，避免数据加载成为瓶颈导致GPU空闲。

4. 推理阶段并行化技巧 #

• 使用allennlp predict批量处理文档时，指定--cuda-device all和合适的batch-size，让多GPU同时处理不同文档或文档块。
• 自定义推理代码中，将输入数据拆分为多份分别送入对应GPU，利用DataParallel的自动分发功能提升效率。

内容的提问来源于stack exchange，提问作者Yijiang Dong