python - 如何在 Pytorch 分布式中使用自定义 Rank 到 GPU 映射

Question

我正在运行分布在单个 4-gpu 节点上的 PyTorch，我不是唯一一个运行项目的节点。因此，有时某些 GPU 太忙，而其他 GPU 则处于空闲状态。

假设 GPU 0 和 GPU 2 很忙，但 GPU 1 和 GPU 3 没有，所以我只想使用这些 GPU。似乎默认情况下，在 pytorch 分布式中，Rank 0 始终是 GPU 0。

如何向 PyTorch 指定 Rank 0 应该对应于 GPU 1，而 Rank 1 应该对应于 GPU 3？

标签： pythonpytorch

Answer 1

您可以通过设置变量来告诉任何命令哪些 GPU 应该对它们可见CUDA_VISIBLE_DEVICES。例如

CUDA_VISIBLE_DEVICES=2,3 python run.py <args>

这只会使 GPU 2 和 3 对该进程可见，因此 DDP 将仅在您指定的那些 GPU 上运行。（请注意，这是零索引的，即0第一个 GPU。）