使用DeepSpeed推理dolly-7B时未充分利用全部GPU的问题

核心问题定位 #

CUDA OOM和单GPU未充分利用的问题，主要是**world_size变量未正确初始化**，导致DeepSpeed无法识别多GPU配置；同时torch.float全精度显存占用过高（7B模型单卡全精度约占28GB，远超单卡24Gi显存），这也是OOM的直接诱因。

具体修复步骤 #

1. 正确获取world_size
   DeepSpeed启动多GPU时会自动分配环境变量，需通过官方API获取正确的多GPU数量，而非手动指定。修改代码如下：

import deepspeed
from transformers import pipeline
import torch

generator = pipeline(model="databricks/dolly-v2-7b", torch_dtype=torch.float16, trust_remote_code=True)
# 从DeepSpeed获取实际启用的GPU数量
world_size = deepspeed.comm.get_world_size()
generator.model = deepspeed.init_inference(
    generator.model,
    mp_size=world_size,
    dtype=torch.float16,  # 改用半精度，大幅降低显存占用
    replace_with_kernel_inject=True
)

2. 降低模型精度
   将torch.float替换为torch.float16（半精度），7B模型半精度显存占用约14GB，4卡并行后单卡仅需约3.5GB，完全适配24Gi显存。若GPU支持（如Ampere及以上架构），也可尝试torch.bfloat16。

3. 修正启动命令拼写
   确保启动命令的脚本名拼写正确：deepspeed --num_gpus 4 script.py（原命令中scipt.py是拼写错误，可能导致启动异常）。

4. 验证多GPU生效
   可在代码中加入打印语句，确认多GPU分配情况：

print(f"当前进程GPU编号: {torch.cuda.current_device()}, 总GPU数量: {world_size}")

运行后应看到4个进程分别对应GPU 0-3。

额外注意事项 #

• 若使用旧版DeepSpeed，可能存在replace_with_kernel_inject兼容性问题，建议升级至最新版本：pip install --upgrade deepspeed
• 确保transformers版本与DeepSpeed兼容，优先使用最新稳定版

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