How to broadcast tensors using NCCL?

vgoklani: <p>I’m unclear on how to broadcast tensors using NCCL from the rank0 process to all other processes. I’ve tried a few approaches, but each attempt freezes the process and puts the GPUs @ 100% utilization (checked via nvidia-smi). Presumably I’m not using <code>broadcast</code> correctly <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>For example, I added a broadcasting snippet to this sample script from huggingface:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/stas00/toolbox/blob/master/pytorch/torch-distributed-gpu-test.py">
  <header class="source">

      <a href="https://github.com/stas00/toolbox/blob/master/pytorch/torch-distributed-gpu-test.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/stas00/toolbox/blob/master/pytorch/torch-distributed-gpu-test.py" target="_blank" rel="noopener nofollow ugc">stas00/toolbox/blob/master/pytorch/torch-distributed-gpu-test.py</a></h4>


      <pre><code class="lang-py">#!/usr/bin/env python

#
# This a `torch.distributed` diagnostics script that checks that all GPUs in the cluster (one or
# many nodes) can talk to each other via nccl and allocate gpu memory.
#
# To run first adjust the number of processes and nodes:
#
# python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# You may need to add --master_addr $MASTER_ADDR --master_port $MASTER_PORT if using a custom addr:port
#
# You can also use the rdzv API: --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT --rdzv_backend c10d
#
# use torch.distributed.launch instead of torch.distributed.run for torch &lt; 1.9
#
# If you get a hanging in `barrier` calls you have some network issues, you may try to debug this with:
#
# NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
</code></pre>



  This file has been truncated. <a href="https://github.com/stas00/toolbox/blob/master/pytorch/torch-distributed-gpu-test.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<p>Please note, the script works perfectly before adding the broadcasting code. That is, the all-reduce results are all correct.</p>
<pre><code class="lang-auto">import fcntl
import os
import socket
import torch
import torch.distributed as dist


def printflock(*msgs):
    """ solves multi-process interleaved print problem """
    with open(__file__, "r") as fh:
        fcntl.flock(fh, fcntl.LOCK_EX)
        try:
            print(*msgs)
        finally:
            fcntl.flock(fh, fcntl.LOCK_UN)


local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
hostname = socket.gethostname()

gpu = f"[{hostname}-{local_rank}]"

try:
    # init
    # dist.init_process_group("nccl")

    os.environ["OMP_NUM_THREADS"] = "1"
    # os.environ["PYTHONHASHSEED"] = str(args.seed)
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'

    # torch.cuda.set_device(args.local_rank)
    dist.init_process_group(
        backend=dist.Backend.NCCL,
        init_method='env://',
        world_size=2,
        rank=local_rank
    )

    assert dist.is_initialized()
    assert dist.is_nccl_available()

    # global rank
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # test all-reduce
    x = torch.ones(1).to(device)
    if local_rank == 0:
        x.add_(1)

    dist.all_reduce(x, op=dist.ReduceOp.SUM)
    dist.barrier()
    printflock(f"{gpu} x: {x}")

    # test all_reduce
    x = torch.tensor(1.0, device='cuda')
    # count = torch.zeros(1).to(device)
    if local_rank == 0:
        x.add_(3)
        # count.add_(1)
        # torch.cuda.comm.broadcast(count, devices=[1])
    #     dist.broadcast(count, src=local_rank)
    dist.all_reduce(x, op=dist.ReduceOp.AVG)
    dist.barrier()
    printflock(f"{gpu} x: {x}")

    ## test broadcast - code freezes in this block!
    # count = torch.zeros(1).to(device)
    # if local_rank == 0:
    #     count.add_(1)
    #     dist.broadcast(tensor=count, src=0)
    # dist.barrier()
    # printflock(f"{gpu} count: {count}")

    # test cuda is available and can allocate memory
    torch.cuda.is_available()
    torch.ones(1).cuda(local_rank)

    printflock(f"{gpu} is OK (global rank: {rank + 1}/{world_size})")

    dist.barrier()
    if rank == 0:
        printflock(f"pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}")
        printflock(f"device compute capabilities={torch.cuda.get_device_capability()}")
        printflock(f"pytorch compute capabilities={torch.cuda.get_arch_list()}")

except Exception:
    printflock(f"{gpu} is broken")
    raise
</code></pre>
<p>I ran it as follows:</p>
<pre><code class="lang-auto">CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 --node_rank=0 torch-distributed-gpu-test.py
</code></pre>
<p>The code freezes in the broadcast portion.</p>
<p>Thanks!</p>
<p>cc: <a class="mention" href="/u/ptrblck">@ptrblck</a> <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>

Li_Shen: <p>I’m having similar issues. <code>broadcast</code> either hangs the program or not broadcasting at all, the other processes still have the old tensor.</p>

wanchaol: <p><a class="mention" href="/u/vgoklani">@vgoklani</a> <a class="mention" href="/u/li_shen">@Li_Shen</a></p>
<p>I answered this in another post <a href="https://discuss.pytorch.org/t/deadlock-when-using-torch-distributed-broadcast/153036" class="inline-onebox">*deadlock* when using torch.distributed.broadcast</a></p>
<p>Basically when you are adding collectives like <code>broadcast</code>, please make sure it’s called on all ranks rather than only rank 0 (like <code>all_reduce</code> in the script), this should resolve this issue. Let me know if it not works.</p>

vgoklani: <p>it works, thank you!!!</p>