Synchronization of CUDA operations between `multiprocess` processes

username: cbcase
cooked: <p>A broad question: is there any documentation or info anywhere about the synchronization guarantees when passing CUDA tensors through multiprocessing queues?</p>
<p>As a concrete example, imagine code like the following:</p>
<pre data-code-wrap="python"><code class="lang-python">def run_function(q):
  while(True):
    x = q.get()
    # Use x, break on some condition, etc.

def main():
  ctx = mp.get_context("spawn")
  q = ctx.SimpleQueue()
  proc = ctx.Process(
        target=run_function,
        args=(q, ),
        daemon=True,
    )
  proc.start()
  x = torch.randn(..., device="cuda")
  W = torch.randn(..., device="cuda")
  for _ in range(50):
    x = torch.matmul(W, x)
  q.put(x)
  proc.join()
</code></pre>
<p>What are the guarantees about the state of <code>x</code> when the child process picks it up out of the queue in <code>run_function</code>? Naively, it looks like it could be anything, as the GPU could have run any number of loop iterations by the time the CPU <code>put</code>s x – and the memory allocator could have happily reused pointers / buffers during that time.</p>
<p>Empirically, though, nothing runs in the child process until the parent process finishes its work (on the GPU) – is this guaranteed? How?</p>
<p>And follow-ups:</p>
<ul>
<li>Does anything change with multiple devices?</li>
<li>What if we don’t put <code>x</code> itself in the queue but rather a python class that contains (possibly nested) references to Tensor objects?</li>
</ul>
<p>Thanks!</p>

username: wconstab
cooked: <p>I don’t know for sure, but, I suspect q.put() is serializing tensor data and sending it to the remote process. If so, this is a CPU-side thing, and before it can be done the tensor data must be requested from the GPU.  This would implicitly synchronize the GPU work to the CPU on the source process before allowing the q.put.</p>
<p>(The opposite possibility, is that tensor data remains on GPU and the remote pytorch process just references the existing gpu data, and I don’t think it’s set up that way.)</p>
<p>You can confirm by doing a profile trace using the torch autograd profiler, and observing a call to cuda d2h memcopy happening before the q.put finishes</p>

username: cbcase
cooked: <p>Thanks, Will. It does look like that’s what is going on. If I have the child process work look like this:</p>
<pre data-code-wrap="python"><code class="lang-python">x_recv = recv_q.get()
x = torch.empty_like(x_recv.x, device=my_device)
x.copy_(x_recv.x)
</code></pre>
<p>and then do some matmuls with <code>x</code>, the profile looks like this:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/2/9211b0fe9650784f09daebf43f608a3ae0b3b5e3.png" data-download-href="https://discuss.pytorch.org/uploads/default/9211b0fe9650784f09daebf43f608a3ae0b3b5e3" title="Screenshot 2023-12-13 at 12.41.25 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/2/9211b0fe9650784f09daebf43f608a3ae0b3b5e3_2_690x372.png" alt="Screenshot 2023-12-13 at 12.41.25 PM" data-base62-sha1="kQbC5nttgm1sYKHRMHyx3rBmvN9" width="690" height="372" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/2/9211b0fe9650784f09daebf43f608a3ae0b3b5e3_2_690x372.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/2/9211b0fe9650784f09daebf43f608a3ae0b3b5e3_2_1035x558.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/2/9211b0fe9650784f09daebf43f608a3ae0b3b5e3.png 2x" data-dominant-color="E4E9EC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2023-12-13 at 12.41.25 PM</span><span class="informations">1204×650 31.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>It’s not at all clear to me how the parent process D2H can overlap with the child process H2D, but I agree that explains the apparent synchronization.</p>
<p>This is a little at odds with the description given at <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html#reuse-buffers-passed-through-a-queue" class="inline-onebox" rel="noopener nofollow ugc">Multiprocessing best practices — PyTorch 2.1 documentation</a>, which makes it sound like CUDA tensors are shared directly. I’ll poke at the CUDA API calls going on as well with IPC handles.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="cbcase" data-post="3" data-topic="193585">
<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/cbcase/48/61078_2.png" class="avatar"> cbcase:</div>
<blockquote>
<p>It’s not at all clear to me how the parent process D2H can overlap with the child process H2D, but I agree that explains the apparent synchronization.</p>
</blockquote>
</aside>
<p>The “overlap” could indicate the wait time (thus sync) of the <code>HtoD</code> call, which is launched after the <code>DtoH</code> code. It does not necessarily mean operations are executed, but could still block.</p>

username: cbcase
cooked: <p>So what I said before is partially true. The combined H2D and D2H appears to occur only when you need to move the tensor across two devices. (<code>my_device</code> in the child is different from the device of <code>x_recv</code>).</p>
<p>In the case of same-device, the tensor is transparently shared with the CUDA IPC mechanisms, but there appear to be some cudaStreamSynchronize()s that occur to guarantee some ordering. Still trying to make sense of the profiles.</p>

username: cbcase
cooked: <p>Alright, here’s a complete example for clarity:</p>
<pre data-code-wrap="python"><code class="lang-python"># Testing indirection of tensors passed through a queue
class Box:
    def __init__(self, x):
        self.x = x

# Child process runs this
def run_fn(send_q, recv_q):
    my_device = "cuda:0"
    torch.cuda.set_device(my_device)
    W1 = torch.randn(4096, 4096, device=my_device)
    torch.cuda.synchronize()
    send_q.put(0)
    y_recv, box_recv = recv_q.get()
    ret = torch.matmul(W1, box_recv.x) + y_recv
    send_q.put(ret)

# Parent process
def main():
    my_device = "cuda:0"
    torch.cuda.set_device(my_device)
    ctx = mp.get_context("spawn")
    send_q = ctx.SimpleQueue()
    recv_q = ctx.SimpleQueue()
    proc = ctx.Process(
        target=run_fn,
        args=(recv_q, send_q),
        daemon=True,
    )
    proc.start()
    proc_ready = recv_q.get()
    assert proc_ready == 0

    W = torch.randn(4096, 4096, device=my_device)
    X = torch.randn(4096, 4096, device=my_device)
    y = torch.randn(4096, device=my_device)
    torch.cuda.synchronize()

    for _ in range(10):
        X = torch.matmul(X, W)

    box = Box(X)
    send_q.put((y, box))
    ret = recv_q.get()
    print(ret.sum().item())
    proc.join()
</code></pre>
<p>The resulting profile is kind of a mess, but it <em>looks</em> like the following occurs:</p>
<ul>
<li>At the point where the parent <code>put()</code>s things into queue, it does a cudaEventRecord.</li>
<li>As soon as the child <code>get()</code>s things out of the queue, the first thing <em>it</em> does is to cudaStreamWaitEvent – presumably on the event created by the parent?</li>
<li>The child also does 2x cudaIpcOpenMemHandle to access the shared tensors</li>
</ul>
<p>This all makes sense and explains the apparent synchronization between parent and child – there is an implicit stream synchronization and the enqueue / dequeue boundary.</p>
<p>Note that I cannot for the life of me find where the parent creates the IpcMemHandles, or how those handles get passed between processes. I’m just taking it on faith that it occurs somewhere.</p>
<p>I’m still not sure I have a robust mental model of sync guarantees for multiprocessing, but it does seem like things <em>generally</em> skew toward implicit safety.</p>