PyTorch在第一个可用的GPU上分配更多内存（CUDA：0）

Question

作为强化学习训练系统的一部分，我正在使用四个GPU并行训练四个策略。对于每个模型，都有两个过程-角色和学习者，仅使用其特定的GPU（例如，与模型＃2对应的演员和学习者仅将GPU＃2用于其所有张量）。演员和学习者通过火炬的share_memory_()共享模型层。 由于四个训练“子系统”是完全对称的，因此我希望它们在四个GPU的每一个上使用完全相同的GPU内存量。但是实际上，我发现在第一个GPU（cuda:0）上分配了更多的GPU内存。

感觉所有内存共享都是通过GPU＃0完成的。有办法解决这个问题吗？

到目前为止，我尝试通过在流程CUDA_VISIBLE_DEVICES函数中显式更改os.environ在子流程中设置start。这似乎没有任何效果，可能是因为子进程是从主进程派生的，该主进程已经初始化了PyTorch CUDA，此时似乎似乎忽略了envvars。

Answer 1

好，到目前为止，我想出了一种解决方法。我的假设是正确的，如果在派生子进程之前已经对PyTorch CUDA子系统进行了初始化，则将CUDA_VISIBLE_DEVICES设置为子进程的其他值不会执行任何操作。

更糟糕的是，调用torch.cuda.device_count()足以初始化CUDA，因此我们甚至无法从PyTorch查询GPU的数量。解决方案是对其进行硬编码，作为参数传递或在单独的过程中查询PyTorch API。我对后者的实现：

import sys


def get_available_gpus_without_triggering_pytorch_cuda_initialization(envvars):
    import subprocess
    out = subprocess.run([sys.executable, '-m', 'utils.get_available_gpus'], capture_output=True, env=envvars)
    text_output = out.stdout.decode()
    from utils.utils import log
    log.debug('Queried available GPUs: %s', text_output)
    return text_output


def main():
    import torch
    device_count = torch.cuda.device_count()
    available_gpus = ','.join(str(g) for g in range(device_count))
    print(available_gpus)
    return 0


if __name__ == '__main__':
    sys.exit(main())

基本上，此函数作为单独的python进程调用其自己的脚本，并读取stdout。

我不会将此答案标记为已接受，因为如果存在的话，我想学习适当的解决方案。