Do gradients propagate through all_reduce & all_gather?

username: derJaeger
cooked: <p>Hi,<br>
this is probably a very stupid question but I would like to make sure I am not making a mistake.<br>
Some part of my loss needs to be computed based on ALL samples in the batch, not just the ones allocated to each GPU. I can aggregate the values I need with all_gather or all_reduce and then compute my final loss. Will the gradients of that loss then properly ‘travel back’ to each individual GPU through the all_gather/all_reduce operation?<br>
Thanks!<br>
Best,<br>
Fabian</p>

username: wanchaol
cooked: <p><a class="mention" href="/u/derjaeger">@derJaeger</a> when you refer to “travel back”, do you mean the gradient flow back to each individual GPU? If so, the answer is not, it will not automatically flow the gradients back to each individual GPU samples if you use the c10d collective. Because currently c10d collective is not autograd enabled yet.</p>
<p>We are working on making the c10d collective autograd enabled (there’s a version of implementation that you can try to use and refer to in <a href="https://github.com/pytorch/pytorch/blob/04d8da88a6a1abf0da2b11096c85244bf38d3b2a/torch/distributed/nn/functional.py#L207" rel="noopener nofollow ugc">here</a>, but it’s not publicly documented and it’s not been publicly released yet, not maintaining well either, so when using it please take your own risk (we might delete this in the future release and make the c10d collective directly autograd enabled). If you want to use it, I recommend you refer to this implementation and write your own version.</p>

username: derJaeger
cooked: <p>Dear <a class="mention" href="/u/wanchaol">@wanchaol</a> , thanks so much for your response! Great to hear you are on it <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"> I have conducted some more research and believe I found what I am looking for:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/Lightning-AI/lightning/blob/ab59f308b18622421edc67048d3b9fbfde96a9f4/src/pytorch_lightning/utilities/distributed.py#L143">
  <header class="source">

      <a href="https://github.com/Lightning-AI/lightning/blob/ab59f308b18622421edc67048d3b9fbfde96a9f4/src/pytorch_lightning/utilities/distributed.py#L143" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/Lightning-AI/lightning/blob/ab59f308b18622421edc67048d3b9fbfde96a9f4/src/pytorch_lightning/utilities/distributed.py#L143" target="_blank" rel="noopener nofollow ugc">Lightning-AI/lightning/blob/ab59f308b18622421edc67048d3b9fbfde96a9f4/src/pytorch_lightning/utilities/distributed.py#L143</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="133" style="counter-reset: li-counter 132 ;">
          <li>    # sync all processes before reduction</li>
          <li>    torch.distributed.barrier(group=group)</li>
          <li>    torch.distributed.all_reduce(result, op=op, group=group, async_op=False)</li>
          <li>
          </li>
<li>    if divide_by_world_size:</li>
          <li>        result = result / torch.distributed.get_world_size(group)</li>
          <li>
          </li>
<li>    return result</li>
          <li>
          </li>
<li>
          </li>
<li class="selected">class AllGatherGrad(torch.autograd.Function):</li>
          <li>    @staticmethod</li>
          <li>    def forward(</li>
          <li>        ctx: Any,</li>
          <li>        tensor: Tensor,</li>
          <li>        group: Optional["torch.distributed.ProcessGroup"] = group.WORLD,</li>
          <li>    ) -&gt; Tensor:</li>
          <li>        ctx.group = group</li>
          <li>
          </li>
<li>        gathered_tensor = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())]</li>
          <li>
      </li>
</ol>
    </code></pre>



  </article>

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

  <div style="clear: both"></div>
</aside>
<p>
I must confess I am not super knowledgeable about how gradients have to be handled in that situation but this implementation appears to do what I need it to. Would you agree?<br>
Best,<br>
Fabian</p>

username: vikigenius
cooked: <p>Hi, I am wondering is this implemented now I see that AllGather does have a backward implementation now in the main branch.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/5a114f72bf0aa93268ba51707a20f99b79ed048d/torch/distributed/nn/functional.py#L316">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/5a114f72bf0aa93268ba51707a20f99b79ed048d/torch/distributed/nn/functional.py#L316" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/5a114f72bf0aa93268ba51707a20f99b79ed048d/torch/distributed/nn/functional.py#L316" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/5a114f72bf0aa93268ba51707a20f99b79ed048d/torch/distributed/nn/functional.py#L316</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="306" style="counter-reset: li-counter 305 ;">
          <li>        tensor = tensor.contiguous()</li>
          <li>        input_tensor_list = tuple(t.contiguous() for t in input_tensor_list)</li>
          <li>        dist.reduce_scatter(tensor, list(input_tensor_list), op=op, group=group)</li>
          <li>        return tensor</li>
          <li></li>
          <li>    @staticmethod</li>
          <li>    def backward(ctx, grad_output):</li>
          <li>        return (None, None, None) + _AllGather.apply(ctx.group, grad_output)</li>
          <li></li>
          <li></li>
          <li class="selected">class _AllGather(Function):</li>
          <li>    @staticmethod</li>
          <li>    def forward(ctx, group, tensor):</li>
          <li>        # Need contiguous tensors for collectives.</li>
          <li>        tensor = tensor.contiguous()</li>
          <li></li>
          <li>        ctx.group = group</li>
          <li>        out_tensor_list = [</li>
          <li>            torch.empty_like(tensor) for _ in range(dist.get_world_size(group=group))</li>
          <li>        ]</li>
          <li></li>
      </ol>
    </code></pre>



  </article>

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

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

<p>Does this mean that the gradients flow back to each GPU ?</p>

username: fduwjj
cooked: <p>yes, the backward of all-gather will be reduce-scatter.</p>