CUDA tensors on multiprocessing queue

username: neerajprad
cooked: <p>I am working on a problem where multiple workers send CUDA tensors to a shared queue that is read by the main process. While the code works great with CPU tensors (i.e. the tensors sent by the workers are retrieved correctly by the main process), I am finding that when the workers send CUDA tensors through the shared queue, the tensor values read by the main process are often garbage values.</p>
<p>Related discussion: <a href="https://discuss.pytorch.org/t/invalid-device-pointer-using-multiprocessing-with-cuda/28023">Invalid device pointer using multiprocessing with CUDA</a>.</p>
<p>e.g. the following minimal code reproduces this issue. It works fine for CPU tensors, but repeatedly reads only the last few tensors that were sent for CUDA.</p>
<pre><code class="lang-python">import torch
import torch.multiprocessing as mp


DEVICE = "cuda"
N = 10

done = mp.Event()


def proc(queue):
    t = torch.tensor(0., device=DEVICE)
    n = 0
    while n &lt; N:
        t = t+1
        print("sent:", t)
        queue.put(t)
        n += 1
    done.wait()


if __name__ == "__main__":
    ctx = mp.get_context("spawn")
    queue = ctx.Queue()
    p = ctx.Process(target=proc, args=(queue,))
    p.daemon = True
    p.start()
    for _ in range(N):
        print("recv:", queue.get())
    done.set()
</code></pre>
<p>which prints out:</p>
<pre><code class="lang-auto">  $ python examples/ex2.py 
sent: tensor(1., device='cuda:0')
sent: tensor(2., device='cuda:0')
sent: tensor(3., device='cuda:0')
sent: tensor(4., device='cuda:0')
sent: tensor(5., device='cuda:0')
sent: tensor(6., device='cuda:0')
sent: tensor(7., device='cuda:0')
sent: tensor(8., device='cuda:0')
sent: tensor(9., device='cuda:0')
sent: tensor(10., device='cuda:0')
recv: tensor(9., device='cuda:0')
recv: tensor(10., device='cuda:0')
recv: tensor(9., device='cuda:0')
recv: tensor(10., device='cuda:0')
recv: tensor(9., device='cuda:0')
recv: tensor(10., device='cuda:0')
recv: tensor(9., device='cuda:0')
recv: tensor(10., device='cuda:0')
recv: tensor(9., device='cuda:0')
recv: tensor(10., device='cuda:0')
</code></pre>
<p>Based on the discussion in <a href="https://discuss.pytorch.org/t/invalid-device-pointer-using-multiprocessing-with-cuda/28023">Invalid device pointer using multiprocessing with CUDA</a> and <a class="mention" href="/u/colesbury">@colesbury</a>’s suggestion in <a href="https://discuss.pytorch.org/t/using-torch-tensor-over-multiprocessing-queue-process-fails/2847/2">Using torch.Tensor over multiprocessing.Queue + Process fails</a>, I suspected that the issue is that we need to hold the reference to the CUDA tensor in the worker process until it is read by the main process. To test this, I appended the tensors to a temporary list in the worker and then it worked fine!</p>
<p>This is however quite wasteful because the workers produce a large number of tensors and it is not feasible to hold them in memory. I was wondering what is the best practice to deal with such a use case. I can try to store the tensors in a limited size queue and hope that the last one has had sufficient time to be read by the main process, but that seems too fragile. cc. <a class="mention" href="/u/smth">@smth</a>, <a class="mention" href="/u/colesbury">@colesbury</a></p>

username: colesbury
cooked: <p>Unfortunately, I don’t know of a good way to do this. Trying to manage the lifetimes of CUDA tensors across processes is complicated.  I try to only share CUDA tensors that have the same lifetime as the program (like model weights for example).</p>

username: neerajprad
cooked: <aside class="quote no-group" data-username="colesbury" data-post="2" data-topic="28626">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/colesbury/48/26_2.png" class="avatar"> colesbury:</div>
<blockquote>
<p>Unfortunately, I don’t know of a good way to do this.</p>
</blockquote>
</aside>
<p>Thanks <a class="mention" href="/u/colesbury">@colesbury</a>. This will be nice to have to do parallel MCMC sampling. For the time being, I’m thinking of using events to periodically clear off tensors from the shared memory - lets see how far that gets us. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: JakobHavtorn
cooked: <p>Interested to hear if you made any progress here <a class="mention" href="/u/neerajprad">@neerajprad</a> <a class="mention" href="/u/colesbury">@colesbury</a>?</p>
<p>I’ve just tried running you minimal example with PyTorch 1.11.0 and was not able to reproduce the issue (tensors are received in the correct order). Also tried adding a <code>time.sleep</code> statement in the <code>proc</code> loop and a <code>gc.collect()</code> in both <code>proc</code> and <code>recv</code> loop to trigger garbage collection of shared CUDA tensors, but it did not happen – the order was still correct.</p>
<p>This kind of data prefetching is appealing because it does not suffer from the limitations/challenges of <code>.to(non_blocking=True)</code>. That is, requiring <code>pin_memory=True</code> (which has it’s own set of potential issues) and the inconvenience of not being “allowed” to log values that originate from CUDA tensors to the terminal between forward-backward passes/batches within sacrificing the effect of <code>non_blocking</code>.</p>