Unexpected Behavior with torch.distributed.isend and irecv in Asynchronous Communication

zhanghb55: <p>I am attempting to utilize asynchronous send (<code>isend</code>) and receive (<code>irecv</code>) for non-blocking communication in a PyTorch distributed setup. My understanding of asynchronous operations implies that <code>torch.distributed.isend</code> should return immediately, allowing for non-blocking execution. However, the behavior observed during execution does not align with this expectation.</p>
<h4><a name="steps-to-reproduce-1" class="anchor" href="#steps-to-reproduce-1"></a>Steps to Reproduce:</h4>
<ol>
<li>Initialize the distributed environment with NCCL backend and CUDA for GPU support.</li>
<li>Use <code>torch.distributed.isend</code> and <code>irecv</code> for sending and receiving tensors between processes.</li>
<li>Introduce a <code>time.sleep(20)</code> delay in the receiving process to simulate asynchronous behavior.</li>
<li>Observe the timestamps logged after the send and receive operations.</li>
</ol>
<h4><a name="expected-behavior-2" class="anchor" href="#expected-behavior-2"></a>Expected Behavior:</h4>
<p>The expectation is that the sending process (process 0) would log its timestamp immediately after initiating the send operation, potentially showing a 20-second earlier timestamp than the receiving process (process 1), which is delayed due to the <code>time.sleep(20)</code>.</p>
<h4><a name="actual-behavior-3" class="anchor" href="#actual-behavior-3"></a>Actual Behavior:</h4>
<p>Both processes log their timestamps almost simultaneously, indicating that the send and receive operations block until both are completed, contrary to the expected asynchronous behavior. This observation suggests that the <code>isend</code> operation might not be returning immediately or that the asynchronous operations are somehow synchronized before logging the timestamps.</p>
<pre><code class="lang-auto">import torch
import os
import time

os.environ['MASTER_ADDR'] = '127.0.0.1'
os.environ['MASTER_PORT'] = '29500'

torch.distributed.init_process_group(backend="nccl")
torch.cuda.set_device(torch.distributed.get_rank() % torch.distributed.get_world_size())

if torch.distributed.get_rank() == 0:
    t = torch.randn(16, 256, 128).cuda()
else:
    t = torch.zeros(16, 256, 128).cuda()

if torch.distributed.get_rank() == 0:
    req = torch.distributed.isend(tensor=t, dst=1)
    timestamp = time.time()
else:
    time.sleep(20)
    req = torch.distributed.irecv(tensor=t, src=0)
    timestamp = time.time()

print(f"{torch.distributed.get_rank()} finishes ", timestamp)
</code></pre>
<p>output is:</p>
<pre><code class="lang-auto">torchrun --nproc_per_node=2 asyncSendRecv.py
[2024-03-25 14:43:13,341] torch.distributed.run: [WARNING] 
[2024-03-25 14:43:13,341] torch.distributed.run: [WARNING] *****************************************
[2024-03-25 14:43:13,341] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2024-03-25 14:43:13,341] torch.distributed.run: [WARNING] *****************************************
[W ProcessGroupNCCL.cpp:1856] Warning: 0NCCL_AVOID_RECORD_STREAMS=1 has no effect for point-to-point collectives. (function operator())
[W ProcessGroupNCCL.cpp:1856] Warning: 0NCCL_AVOID_RECORD_STREAMS=1 has no effect for point-to-point collectives. (function operator())
1 finishes  1711349017.1358058
0 finishes  1711349017.1358182
</code></pre>