When does p2p communication uses different NCCL streams?

username: Edenzzzz
cooked: <p>When using point-to-point communications, it might sometimes be desirable to schedule unrelated comm calls on different streams, e.g. in the backward of Ring attention.<br>
Based on this example(<a href="https://github.com/InternLM/InternEvo/blob/f2949df89c15e1c16b6d48412ae9f94122ef463d/internlm/model/ops/ring_flash_attn/zigzag_ring_flash_attn_with_sliding_window.py#L196" class="inline-onebox" rel="noopener nofollow ugc">InternEvo/internlm/model/ops/ring_flash_attn/zigzag_ring_flash_attn_with_sliding_window.py at f2949df89c15e1c16b6d48412ae9f94122ef463d · InternLM/InternEvo · GitHub</a>), I profiled to see 3 NCCL streams for p2p comms matching <code>local_dkv_comm</code>, <code>local_kv_comm</code> and <code>dkv_comm</code><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/2/7282a5466a53f04d4c2fd3a0fd4fb444b502f63e.png" data-download-href="https://discuss.pytorch.org/uploads/default/7282a5466a53f04d4c2fd3a0fd4fb444b502f63e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/7282a5466a53f04d4c2fd3a0fd4fb444b502f63e_2_690x219.png" alt="image" data-base62-sha1="gl0kZo3QrIeGzZnHV0o7ej82v2S" width="690" height="219" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/7282a5466a53f04d4c2fd3a0fd4fb444b502f63e_2_690x219.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/7282a5466a53f04d4c2fd3a0fd4fb444b502f63e_2_1035x328.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/7282a5466a53f04d4c2fd3a0fd4fb444b502f63e_2_1380x438.png 2x" data-dominant-color="CBBBD4"><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">image</span><span class="informations">2344×746 91.4 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>However I couldn’t reproduce them elsewhere with 2 process groups.<br>
Is it because that torch uses different streams only when there are a number of process groups?</p>

username: fduwjj
cooked: <p>IIUC, for each processGroup we assign the job to a NCCL stream. So if you want 3 streams, maybe you want to create a different PG?</p>

username: Edenzzzz
cooked: <p>That’s right, Closing in favour of the answer on github</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/67158#issuecomment-2277092652">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/67158#issuecomment-2277092652" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/67158#issuecomment-2277092652" target="_blank" rel="noopener nofollow ugc">Make streams used for NCCL operations configurable</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-10-24" data-time="17:13:26" data-timezone="UTC">05:13PM - 24 Oct 21 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/wjuni" target="_blank" rel="noopener nofollow ugc">
          <img alt="wjuni" src="https://avatars.githubusercontent.com/u/7837303?v=4" class="onebox-avatar-inline" width="20" height="20">
          wjuni
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: distributed
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: nccl
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🚀 Feature

Make streams used for NCCL operations configurable

## Motivat<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">ion

I've noticed that PyTorch distributed module has introduced P2P send and receive functionality via NCCL (which is listed as "not supported" yet on the document though). However, I found possible cases where NCCL hangs when two nodes exchange tensor with NCCL P2P `send` and `recv`. Here is (possibly) a minimal working example.

```python
import torch
import torch.distributed as dist
from argparse import ArgumentParser

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--rank", type=int)
    args = parser.parse_args()

    local_rank, remote_rank = args.rank, 1 - args.rank
    device = torch.device('cuda', local_rank)
    torch.cuda.set_device(device)

    dist.init_process_group(backend="nccl", init_method='tcp://127.0.0.1:30001', rank=local_rank, world_size=2)
    dist.barrier()
    print("Ready")

    tensor_to_send = torch.empty((1000000,), device=device)
    tensor_to_recv = torch.empty((1000000,), device=device)

    # exchange tensors between two nodes
    dist.send(tensor_to_send, remote_rank)
    dist.recv(tensor_to_recv, remote_rank)
    torch.cuda.synchronize() # never returns
    
    print("Done")
```
I believe running two processes with command arguments `--rank 0` and `--rank 1` should not raise any problems. However, they will likely end up hanging at `torch.cuda.synchronize()` with 100% GPU utilization (if not, please try with larger tensors).

This is because `dist.send` and `dist.recv`, which internally call `ncclSend` and `ncclRecv`, cause deadlock. `ncclSend` blocks until the remote rank calls `ncclRecv`, and vice versa (see https://github.com/NVIDIA/nccl/issues/584). Thus, when `node 0` and `node 1` issue `dist.send` at the same time, and they will not be able to proceed to `dist.recv`. This makes issuing successive calls to `dist.send` and `dist.recv` very fragile to deadlocks.

To handle this, the NCCL documentation recommends grouping `ncclSend` and `ncclRecv` operations with with `ncclGroupStart` and `ncclGroupEnd` (`dist.P2POp` and `dist.batch_isend_irecv` in PyTorch). However, this does not always work for all cases, especially when two operations cannot be grouped. For example, when two nodes randomly send and receive tensors to each other, two nodes may simultaneously call `send` and fall into a deadlock as there is no locking mechanism here. 

In my project, there are multiple nodes exchanging tensors with each other, and any node can send or receive a tensor. The sender transmits a tensor and corresponding metadata to the receiver as needed. To achieve this, metadata containing the size of the tensor is delivered to the receiver out-of-band via RPC. The receiver calls `dist.recv` upon receiving the metadata, and the receiver calls `dist.send` to send the tensor. Here, when two nodes call `dist.send` simultaneously by chance, this causes GPUs hang and it will never return from `synchronize()`.



## Pitch


IMHO, the easiest way to resolve this is using different CUDA streams for the `dist.send` and `dist.recv`.  Unfortunately, it seems that there is no viable way for users to specify a specific stream for the NCCL operations. PyTorch distributed uses a dedicated stream for each NCCL connection (https://github.com/pytorch/pytorch/issues/60511), and `with torch.stream(stream)` does not work as in other CUDA operations.

I would like to have a context manager that can configure a stream used for NCCL operations something like below:
```python
# At the sender side
stream_send = dist.NCCLStream(remote_rank) # Create stream for NCCL operation
with dist.use_stream(stream_send): 
    # this context should not affect other CUDA operations
    dist.send(tensor_to_send, remote_rank)
torch.cuda.synchronize()


# At the receiver side
stream_recv = dist.NCCLStream(remote_rank) # Create stream for NCCL operation
with dist.use_stream(stream_recv): 
    # this context should not affect other CUDA operations
    dist.recv(tensor_to_recv, remote_rank)
torch.cuda.synchronize()
```
so that `dist.send` and `dist.recv` do not block each other.

## Alternatives


Well, we can implement some hacky locking mechanisms to prevent two nodes from calling `dist.send` at the same time. However, this would come with some overhead, and doesn't seem to be a good solution...

## Additional context


I'll be glad to know if other workarounds can be applied to my problem. 

I'd like to work on this and submit a PR if there is no workaround to resolve this problem.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>