Distributed broadcast operation slows down training speed

overfitted: <p>Hi everyone! I use a piece of PyTorch code that runs on a single machine distributed setting. The code contains <code>all_gather</code> and <code>all_reduce</code> operations to gather predictions from each gpu and calculate metrics, respectively, without any noticeable slowing down in training speed. I recently added an extra bit in the code for a new thing that I’m trying where I need to broadcast a probability sampled from a uniform distribution (single number) in all gpus, so all of them have the same probability value. I added it within the training loop and looks like this:</p>
<pre><code class="lang-auto">if dist.get_rank() == 0:
    prob = torch.rand(1).cuda()
else:
    prob = torch.zeros(1).cuda()
dist.broadcast(prob, 0)
</code></pre>
<p>After adding this, the code becomes significantly slower! Any ideas why? Thank you!</p>

huahuanZ: <p>I think in your example the creation of <code>prob</code> when <code>rank != 0</code> should be <code>prob=torch.zeros(1, device=rank)</code>.<br>
Your codes:</p>
<pre><code class="lang-auto">def worker(gpu, *args):
    ...
    if gpu == 0:
        prob = torch.rand(1).cuda() #.to(device=gpu)
        print(gpu, prob)
    else:
        prob = torch.zeros(1).cuda() # .to(device=gpu)
        print(gpu, prob)

# output
0 tensor([0.1045], device='cuda:0')
1 tensor([0.], device='cuda:0') &lt;- same on cuda:0
</code></pre>
<p>Specify device</p>
<pre><code class="lang-auto">def worker(gpu, *args):
    ...
    if gpu == 0:
        prob = torch.rand(1, device=gpu)
        print(gpu, prob)
    else:
        prob = torch.zeros(1, device=gpu)
        print(gpu, prob)

# output
0 tensor([0.8787], device='cuda:0')
1 tensor([0.], device='cuda:1')
</code></pre>

overfitted: <p>Hi David! Thank you for your comment. I also thought of that but the PyTorch documentation got me confused. In particular, in <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.cuda.html#torch.Tensor.cuda" class="inline-onebox" rel="noopener nofollow ugc">torch.Tensor.cuda — PyTorch 1.9.0 documentation</a> for the <code>device</code> argument, they say:<br>
<strong>device</strong> (<a href="https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.device" rel="noopener nofollow ugc"><code>torch.device</code></a>) – The destination GPU device. Defaults to the current CUDA device.</p>
<p>So I thought that if you don’t specify it the default would be the current CUDA device rather than always  <code>'cuda:0'</code>??</p>

huahuanZ: <p>That’s up to whether you have set the current device.</p>
<pre><code class="lang-python">def worker(gpu, *args):
    ...
    torch.cuda.set_device(gpu) # &lt;- This line makes the difference
    if gpu == 0:
        prob = torch.rand(1).cuda()
        print(gpu, prob)
    else:
        prob = torch.zeros(1).cuda()
        print(gpu, prob)

# output
0 tensor([0.6916], device='cuda:0')
1 tensor([0.], device='cuda:1') # &lt;- on correct device
</code></pre>

overfitted: <p>This is super helpful! Thank you so much! I’ll keep you updated for the effect the fix has on the training speed. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>