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

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 rank

0，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

启动器会在开始时设置环境变量，本地世界大小可以从操作系统环境变量中获取（默认为节点中的 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'                                                                                                                                     ```

Answer 3

如果您使用torchrun，您可以使用torchrun设置的环境变量获取局部世界大小。

如果是这样的话，您问题的直接答案是

local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])