How to share a CUDA tensor between processes

Yann: <p>Hello,</p>
<p>I have an application where transferring large torch Tensors from CPU to CUDA is a bottleneck. I am trying to solve this by preloading Tensors on CUDA using background worker processes with <code>torch.multiprocessing</code>.</p>
<p>Ideally, my plan would be to store Tensors in shared CPU memory so that my different workers can sample Tensors from this shared memory, put them on CUDA, and finally transfer the CUDA tensors to the main process.</p>
<p>However, my initial tests have been unsuccessful so far. More specifically, I seem to be able to work on shared Tensors in CPU memory, but not in CUDA memory.</p>
<p>I am using this toy code snippet taken from StackOverflow for testing. It generates a list of 2 Tensors, puts them in shared memory, asks other processes to modify these tensors, and then reads the Tensors values in the main process:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch.multiprocessing as mp
import torch

DEVICE = "cpu"

def foo(worker,tl):
    # print(f"WORKER {worker}: {tl}")
    tl[worker] += (worker+1) * 1000

if __name__ == '__main__':
    print(f"USING DEVICE {DEVICE}")
    mp.set_start_method('spawn', force=True)
    tl = [torch.randn(2, device=DEVICE), torch.randn(3, device=DEVICE)]

    for t in tl:
        t.share_memory_()

    print("before mp: tl=")
    print(tl)

    p0 = mp.Process(target=foo, args=(0, tl))
    p1 = mp.Process(target=foo, args=(1, tl))
    p0.start()
    p1.start()
    p0.join()
    p1.join()

    print("after mp: tl=")
    print(tl)
</code></pre>
<p>Depending on <code>DEVICE</code>, the output is:</p>
<ul>
<li>CPU: Seems to work properly:</li>
</ul>
<pre><code class="lang-auto">USING DEVICE cpu
before mp: tl=
[tensor([-0.7902,  0.1714]), tensor([-0.0894, -0.3304, -0.4405])]
after mp: tl=
[tensor([ 999.2098, 1000.1714]), tensor([1999.9106, 1999.6696, 1999.5596])]
</code></pre>
<ul>
<li>CUDA: Looks like the CUDA tensors are zero-filled instead of random-filled when accessed by the background processes, and yet they still get modified by the background processes:</li>
</ul>
<pre><code class="lang-auto">USING DEVICE cuda:0
before mp: tl=
[tensor([-0.4746, -0.2319], device='cuda:0'), tensor([0.8965, 1.0782, 0.8496], device='cuda:0')]
after mp: tl=
[tensor([1000., 1000.], device='cuda:0'), tensor([2000., 2000., 2000.], device='cuda:0')]
</code></pre>
<p>It looks like CUDA doesn’t behave well with this.</p>
<p>Can someone explain this behavior please? Also, is there a better way of preloading from CPU to GPU, since this is what I am trying to achieve eventually?</p>
<p>Thanks!</p>
<p>PS: This is on Windows, running Python with conda. In Jupyter, the behavior is even weirder.</p>

Yann: <p>Hello again,</p>
<p>This seems to be a bug in PyTorch ? Here is a simpler script illustrating a similar issue. The tensor received by the main process is a CUDA tensor of the right shape, but filled with zeros:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.multiprocessing as mp


def cuda_loader_process(queue_tensors: mp.Queue, queue_msgs: mp.Queue):
    # Create (load...) a CPU tensor:
    t_cpu = torch.ones((1,2), device = "cpu")
    # Move to CUDA:
    t_cuda = t_cpu.to("cuda:0")
    # Send the CUDA tensor to the main process:
    print(f"CHILD PROCESS - Sending: {t_cuda} to the main process")
    queue_tensors.put(t_cuda)
    # Wait for the main process to allow the child process to exit:
    msg = queue_msgs.get()
    print(f"CHILD PROCESS - Received msg: {msg}")


if __name__ == '__main__':
    # set multiprocessing to use the "spawn" method for child processes:
    mp.set_start_method("spawn", force=True)

    # inter-process messaging Queues:
    queue_tensors = mp.Queue()
    queue_msgs = mp.Queue()

    # instantiate child CUDA loader process:
    p_cuda_loader = mp.Process(target=cuda_loader_process, args=(queue_tensors, queue_msgs))
    p_cuda_loader.start()

    # wait for CUDA tensor from the child process:
    t_cuda = queue_tensors.get().clone()  # clone the tensor to own reference (avoids handling CUDA IPC)
    print(f"MAIN PROCESS - Received tensor: {t_cuda}")

    # ask the child process to exit:
    queue_msgs.put("done")

    # join the child process:
    p_cuda_loader.join()
</code></pre>
<p>OUTPUT:</p>
<pre data-code-wrap="terminal"><code class="lang-terminal">CHILD PROCESS - Sending: tensor([[1., 1.]], device='cuda:0') to the main process
MAIN PROCESS - Received tensor: tensor([[0., 0.]], device='cuda:0')
CHILD PROCESS - Received msg: done
</code></pre>