在火炬分布式训练中获取本地世界大小

Question

假设我有 2 台机器，每台机器有 4 个 GPU。假设训练算法的每个实例需要 2 个 GPU。我想运行 4 个进程，每台机器 2 个，每个进程使用 2 个 GPU。

如何让每个进程检索同一台计算机上运行的本地进程的数量？我可以检测world size到

torch.distributed.get_world_size()

和global rank与

torch.distributed.get_rank()

但是，鉴于我不想对参数进行硬编码，有没有办法恢复每个节点上运行 2 个进程？这对于我将 GPU 平均分配给每个进程很有用。

示例：假设我知道一台机器有 4 个 GPU，并且上面有 2 个进程，我将分配 GPU[0, 1]来处理local rank0 级的进程，分配 GPU[2, 3]来处理本地等级 1 的进程。我知道进程总数，但我无法理解它们是否是在同一台机器上，所以我无法决定他们可以使用多少个 GPU。

我需要一个可以调用的函数torch.distributed.get_local_world_size()

Answer 1

torch.cuda.device_count()本质上是本地世界的大小，可用于确定每个设备上有多少可用的 GPU。如果由于某种原因无法做到这一点，使用普通 MPI 可能会有所帮助

from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank() # device rank - [0,1]

torch.cuda.device(i)
ngpus = torch.cuda.device_count()
print(ngpus, " gpus on machine", i) # here's local world size for each process

torch.cuda.device_count()但我认为在任何情况下只要调用而不添加这种依赖关系就可以了。我是新来的，所以如果可以的话，请告诉我如何改进这个答案。

Answer 2

启动器会在开始时设置 envs，本地世界大小可以从操作系统环境变量中获取（默认为节点中 GPU 的数量）：

# -*- coding: utf-8 -*-                                                                                                                                                                                                                                                                                                                                 import os                                                                                                                                                                   import torch.distributed as dist                                                                                                                                            import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=0)
args = parser.parse_args()

dist.init_process_group('nccl')
local_rank = args.local_rank
local_world_size = os.environ["LOCAL_WORLD_SIZE"]                                                                                                                           
print(f'{local_rank = }; { local_world_size = }')

运行： python3 -m torch.distributed.launch --nproc_per_node=4 test.py 输出：

local_rank = 0;  local_world_size = '4'                                                                                                                                     
local_rank = 3;  local_world_size = '4'                                                                                                                                     
local_rank = 1;  local_world_size = '4'                                                                                                                                     
local_rank = 2;  local_world_size = '4'                                                                                                                                     ```