使用DeepSpeed微调Llama-7B时预留GPU的方法咨询

以下是两个可行的解决思路，避开--include参数导致的GPU可见性限制问题：

方案1：不限制全局GPU可见性，通过DeepSpeed配置控制训练设备，训练后迁移模型到预留GPU #

• 启动训练时，不要用deepspeed --include过滤GPU，让所有GPU保持可见。比如机器有4块GPU（0-3），想预留3号，直接执行：

  CUDA_VISIBLE_DEVICES=0,1,2,3 deepspeed train.py --deepspeed ds_config.json
  

• 在DeepSpeed配置文件（如ds_config.json）中，指定训练用的GPU数量：

  {
    "train_batch_size": 1,
    "num_gpus": 3,
    "optimizer": {"type": "AdamW", "params": {"lr": 2e-5}},
    "fp16": {"enabled": true},
    "zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu"}}
  }
  

  这里num_gpus设为3，DeepSpeed会自动占用前3块GPU（0、1、2）用于训练，预留的3号GPU不会被训练进程占用。
• 训练完成后，从DeepSpeed的模型包装中提取原始模型，迁移到预留GPU：

  # 训练结束后获取原始模型
  raw_model = trainer.model.module
  # 迁移到预留的GPU
  raw_model.to('cuda:3')
  # 后续非训练前向调用直接用该模型，速度正常
  logits = raw_model(**inputs).logits
  

  这样既保留了DeepSpeed的内存优化（避免OOM），又能让后续前向在GPU上快速运行。

方案2：用torch.distributed.launch指定训练进程数，预留GPU独立使用 #

• 启动训练时，用torch.distributed.launch配合DeepSpeed，指定训练用的GPU进程数。比如预留3号GPU，训练用0、1、2：

  CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=3 --master_port=29500 train.py --deepspeed ds_config.json
  

  --nproc_per_node=3会启动3个训练进程，分别绑定0、1、2号GPU，3号GPU处于空闲状态。
• 训练完成后同样提取原始模型，迁移到cuda:3执行前向调用即可。

重要提醒 #

• 训练过程中不要手动移动DeepSpeed包装后的模型设备，这会破坏其分布式内存优化逻辑，直接导致OOM。必须等训练结束后，提取未被包装的原始模型再做设备迁移。
• 如果需要在训练过程中穿插非训练前向，可以单独启动一个进程，加载训练 checkpoint 到预留GPU执行，避免和训练进程冲突。

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