Will "dist.all_gather" break the auto gradient graph?

username: bengouawu
cooked: <p>if we use <strong>torch.nn.parallel.gather</strong> to collect data from other GPUs, then we do some operations on the gathered data <strong>A</strong>, the gradient will go back to the original array( since it’s from the ‘nn’ package, that’s what supposed to happen).   We can check grad_fn of  <strong>A</strong>,  it says <strong>CopyBackwards</strong>, that’s what we expect.</p>
<p>but it seems the <strong>all_gather</strong> from <strong>torch.distributed</strong> package is different.  it simply copies the data here and doesn’t have any grad_fn function.  Does this mean if we use this function, we have to manually compute gradients for the gathered data?</p>
<p>related post:</p><aside class="quote" data-post="1" data-topic="35143">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/y/e480ec/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/something-about-dist-all-gather-operation/35143">Something about dist.all_gather operation</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    I have a question. When we apply dist.all_gather() operation, suppose the there are 4 gpus, and each gpu will get the value of others, and when we apply the result of all_gather with ground truth to calculate loss, does loss can backward? or the dist.all_gather operation will break the graph like the operation of detach()?
  </blockquote>
</aside>
<aside class="quote quote-modified" data-post="1" data-topic="45184">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/90ced4/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/is-my-code-the-correct-way-using-distributeddataparallel-in-single-node-multi-gpus/45184">Is my code the correct way using DistributedDataParallel in single node multi GPUs?</a> <a class="badge-category__wrapper " href="/c/distributed/12"><span data-category-id="12" data-drop-close="true" class="badge-category "><span class="badge-category__name">distributed</span></span></a>
  </div>
  <blockquote>
    I use torch.nn.parallel.DistributedDataParallel API in PyTorch1.1 to spawn my multi-card model (2 GPUs) to 8 GPUs.  According to official tutorial <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">GETTING STARTED WITH DISTRIBUTED DATA PARALLEL</a>, DistributedDataParallel is the recommanded way to parallel one’s model. I am not confident about my implementation and I can’t find other valuable tutorials, so come here for help. 
My ideas are simply as follow: 

split my 3D CNN model into 2 GPUs (simply called dev_in and dev_out),
use DistributedDataP…
  </blockquote>
</aside>


username: bengouawu
cooked: <p>Hi, I ended up writing the backward function manually. I used all_gather function from nccl, so Pytorch won’t take care of that of course. I hope they can add support for that in the future.</p>

username: amsword
cooked: <p>i also had similar questions, but finally i find the current design is good. For all_gather, the gradient will not be propagated back to other devices, but the gradient for current device can be calculated correctly. Since each device calculates the loss and propagates the gradient for each device, the final results are correct. For gather, i assume the non-0 device will not calculate the loss and only device 0 does the loss calculation. In that case, the gradient has to be sent back to non-0 devices.</p>

username: maximedb
cooked: <p>Do you have an example of how to do that ?<br>
Thanks!</p>

username: amsword
cooked: <p>Wrote a blog about a way to use all_gather, without the need to calculate the gradient</p>
<aside class="onebox allowlistedgeneric">
  <header class="source">
      <a href="https://amsword.medium.com/gradient-backpropagation-with-torch-distributed-all-gather-9f3941a381f8" target="_blank" rel="noopener nofollow ugc" title="12:16AM - 07 February 2021">Medium – 7 Feb 21</a>
  </header>
  <article class="onebox-body">
    <img src="" class="thumbnail" width="" height="">

<h3><a href="https://amsword.medium.com/gradient-backpropagation-with-torch-distributed-all-gather-9f3941a381f8" target="_blank" rel="noopener nofollow ugc">Gradient backpropagation with torch.distributed.all_gather</a></h3>

<p>This blog presents some tips on how to use torch.distributed.all_gather and make sure the gradient is correctly calculated for deep…</p>

  <p><span class="label1">Reading time: 4 min read</span>
    </p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: kaizhao
cooked: <p>Here is an implementation of <code>all_gather</code> with gradient back-prop: <a href="https://github.com/vlkit/vlkit/blob/master/vlkit/ops/distributed.py#L4-L25" rel="noopener nofollow ugc">vlkit.ops.distributed.all_gather()</a>.</p>

username: DeepakSaini119
cooked: <p>Hi, this link doesn’t work. I wanted to implement something similar. Thanks.</p>

username: kaizhao
cooked: <p>Hi the link is now open. Alternatively you can use <code>vlkit.ops.all_gather</code>.</p>