作为强化学习训练系统的一部分,我使用四个 GPU 并行训练四个策略。对于每个模型,有两个进程——actor 和 learner,它们只使用它们特定的 GPU(例如,对应于模型 #2 的actor 和 learner 只使用 GPU #2 作为它们的所有张量)。 Actor 和学习器通过 torch 的 share_memory_()
共享模型层。
由于四个训练“子系统”是完全对称的,我希望它们在四个 GPU 中的每一个上使用完全相同数量的 GPU 内存。然而,在实践中,我看到在第一个 GPU ( cuda:0
) 上分配了更多的 GPU 内存。
感觉所有的内存共享都是通过 GPU #0 以某种方式完成的。有没有办法解决这个问题?
到目前为止,我尝试通过显式更改进程 CUDA_VISIBLE_DEVICES
函数中的 os.environ
在子进程中设置 start
。这似乎没有任何影响,可能是因为子进程是从主进程派生出来的,其中 PyTorch CUDA 已经初始化,而此时 envvars 似乎被忽略了。
最佳答案
好的,到目前为止我想出了一个解决方法。我的假设是正确的,如果 PyTorch CUDA 子系统在子进程 fork 之前已经初始化,那么将 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 进程并读取标准输出。
我不会将此答案标记为已接受,因为我想了解正确的解决方案(如果存在)。
关于deep-learning - PyTorch 在第一个可用的 GPU 上分配更多内存(cuda :0),我们在Stack Overflow上找到一个类似的问题:https://stackoverflow.com/questions/59873577/