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

Question

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

如何让每个进程检索同一台机器上运行ning 本地进程的数量？我可以用

检测到 world size

torch.distributed.get_world_size()

和 global rank 与

torch.distributed.get_rank()

但是，考虑到我不想对参数进行硬编码，有没有办法恢复每个节点上运行ning 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

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

# -*- 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'                                                                                                                                     ```

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

Get local world size in torch distributed training

pytorch

gpu

distributed-computing