PyTorch allocates more memory on the first available GPU (cuda:0)
As a part of the reinforcement learning training system, I am training four policies in parallel using four GPUs. For each model, there are two processes - the actor and the learner, which only use their specific GPU (e.g. actor and learner corresponding to model #2 only use GPU #2 for all their tensors). Actor and learner share the model layers via torch's share_memory_().
Since the four training "subsystems" are completely symmetric, I would expect them to use the exact same amount of GPU memory on each of the four GPUs. In practice, however, I see a lot more GPU memory allocated on the first GPU (cuda:0).

It feels like all the memory sharing is somehow done via GPU #0. Is there a way to fix this?

So far, I tried setting CUDA_VISIBLE_DEVICES in child processes by explicitly altering os.environ in the process start function. This does not seem to have any effect, probably because child processes are forked from the main process, where PyTorch CUDA is already initialized, and envvars just seem to be ignored at this point.



Ok, so far I came up with a workaround. My hypothesis was right, if PyTorch CUDA subsystem is already initialized before the child process is forked, setting CUDA_VISIBLE_DEVICES to a different value for a subprocess does not do anything.

Even worse, calling torch.cuda.device_count() is enough to initialize CUDA, so we can't even query the number of GPUs from PyTorch. Solution is either hardcode it, pass as a parameter, or query PyTorch API in a separate process. My implementation for the latter:

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())


Basically this function calls its own script as separate python process and reads stdout.

I won't mark this answer as accepted because I would like to learn a proper solution if it exists.
