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

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() 即可。在任何情况下都无需添加此依赖项。我在这里很新，所以如果可以的话，请告诉我如何改进这个答案。