PyTorch使用DistributedDataParallel时如何指定特定GPU设备ID

一、指定仅使用4、5、6、7号GPU的方法 #

两种常用实现方式，按需选择即可：

• 方式一：环境变量限制可见GPU（最易用，无代码侵入）
  启动训练脚本前通过CUDA_VISIBLE_DEVICES环境变量过滤可见设备，命令如下：

  CUDA_VISIBLE_DEVICES=4,5,6,7 python your_training_script.py
  

  此时PyTorch会自动将4、5、6、7号物理GPU映射为本地可见的0、1、2、3号逻辑GPU，无需额外修改DDP的设备绑定逻辑。
• 方式二：代码内显式绑定设备
  如果不想修改启动命令，可以在代码中手动指定设备：
  首先定义目标设备列表：

  target_gpus = [4,5,6,7]
  

  进程初始化时，将模型、数据加载到对应设备，DDP初始化时传入device_ids参数：

  import torch
  from torch.nn.parallel import DistributedDataParallel as DDP
  
  # local_rank为当前进程的本地序号，取值0~3
  current_device = target_gpus[local_rank]
  model = model.to(current_device)
  ddp_model = DDP(model, device_ids=[current_device])
  

二、world_size与rank参数设置 #

你判断的world_size=4是正确的，该参数代表全局训练进程总数，单卡对应一个进程的场景下，4张卡的world size固定为4。
rank的分配逻辑分两种场景：

• 场景1：使用官方启动工具（torchrun/torch.distributed.launch，推荐）
  不需要手动分配rank，启动工具会自动为每个进程分配全局rank与本地rank。
  启动命令示例（配合CUDA_VISIBLE_DEVICES方案）：

  CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node=4 your_training_script.py
  

  代码内直接通过分布式接口获取即可：

  import os
  import torch.distributed as dist
  
  # 初始化进程组
  dist.init_process_group(backend="nccl", world_size=4)
  global_rank = dist.get_rank() # 全局rank，取值0、1、2、3，分别对应4、5、6、7号GPU
  local_rank = int(os.environ["LOCAL_RANK"]) # 节点内本地rank，取值0~3
  

• 场景2：手动实现多进程启动
  如果是自行编写多进程逻辑，需要为4个进程分别传入0、1、2、3的全局rank值，和4张GPU一一对应即可。

补充注意：单机多卡场景下使用默认的init_method="env://"即可完成进程间通信组网，不需要额外配置。

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