Pytorch cudagraph with nccl operation failed

youkaichao1: <p>Here is a simple example:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.distributed as dist

def main():
    import os
    # Initialize the NCCL backend
    dist.init_process_group(backend='nccl')
    world_size = int(os.environ['WORLD_SIZE'])
    rank = int(os.environ['RANK'])

    # Create a tensor on the GPU
    tensor = torch.rand(10).cuda(rank)

    # Start CUDA graph capture
    stream = torch.cuda.Stream(device=f"cuda:{rank}")
    graph = torch.cuda.CUDAGraph()
    stream.synchronize()
    with torch.cuda.graph(graph, stream=stream):
        # Perform all-reduce operation
        dist.all_reduce(tensor)
    stream.synchronize()

    # Execute the graph
    graph.replay()

    print("All-reduce completed:", tensor)

if __name__ == "__main__":
    main()

</code></pre>
<p>Run with <code>torchrun --nproc-per-node 8 test.py</code>, got the following error:</p>
<pre><code class="lang-auto">torch/distributed/distributed_c10d.py", line 2050, in all_reduce
    work = group.allreduce([tensor], opts)
torch.distributed.DistBackendError: NCCL error in: ../torch/csrc/distributed/c10d/NCCLUtils.hpp:219, unhandled cuda error (run with NCCL_DEBUG=INFO for details), NCCL version 2.18.1
ncclUnhandledCudaError: Call to CUDA function failed.
Last error:
Cuda failure 'operation not permitted when stream is capturing'
</code></pre>
<p>Per the <a href="https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/cudagraph.html" rel="noopener nofollow ugc">nccl documentation</a>:</p>
<blockquote>
<p>Starting with NCCL 2.9, NCCL operations can be captured by CUDA Graphs.</p>
</blockquote>
<p>So why does pytorch cudagraph fail to capture allreduce operation?</p>

ptrblck: <p>From the <a href="https://pytorch.org/docs/stable/notes/cuda.html#usage-with-distributeddataparallel">docs</a>:</p>
<blockquote>
<p>Disable DDP’s internal async error handling<br>
Before full-backward capture, DDP must be constructed in a side-stream context<br>
Your warmup must run at least 11 DDP-enabled eager iterations before capture.</p>
</blockquote>

youkaichao1: <p>If I have a large network, do I need to warm up the whole network for at least 11 runs, or I can just launch some small DDP/pytorch iterations to fulfill the requirement?</p>

ptrblck: <p>You need full iterations of the DDP workload you want to capture and replay.</p>

youkaichao1: <p>What if I’m just using inference without grad, i.e. just using <code>torch.distributed.all_reduce</code>? Do I still have to run the full workload for more than 11 runs?</p>

ptrblck: <p>I would claim it wouldn’t hurt, but if you don’t need to reduce any gradients and initialize the grad buckets, you might be able to reduce the warmup iterations to 3.</p>

youkaichao1: <p>These numbers are sooooo magic <img src="https://discuss.pytorch.org/images/emoji/apple/rofl.png?v=12" title=":rofl:" class="emoji" alt=":rofl:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>There are of course technical reasons to make sure the heuristics were able to select the fastest kernels, kernel fusions and other optimizations are done, buffers were allocated… But the true reason is of course:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/4/c/4c5a198368c34f8f367a3a1cdb92e26e13eba15a.png" alt="1000076234" data-base62-sha1="aTrkaNa9I8U2w8g9Ektx1IZABwm" width="487" height="327"></p>

youkaichao1: <p>lol is this an existing meme or you just created one for for this thread?</p>

ptrblck: <p>Yes, I made it, but I already had it ready from another occasion <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

wenqibiao: <p>so this problem be solved?</p>