How to synchronize a variable across entire batch in ddp

username: hij
cooked: <p>I am using ddp to train my model now. In my model I want to calculate the standard deviation across the batch. However since I am wondering if i can calculate the standard deviation across the entire batch instead of within each device. The standard deviation will be part of my computation graph.</p>
<p>I feel that this is similar to synchronize batchnorm and should be doable. How would I go about doing this? Here is an example of what I want to do</p>
<pre><code class="lang-auto">def forward(self, input):
        feats = conv(input)
        batch, channel, height, width = feats.shape
        stddev = feats.view(batch, -1, 1, channel, height, width)
        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
        stddev = stddev.repeat(batch, 1, height, width)
        feats = torch.cat([feats, stddev], 1)
        output = conv_last(feats)
        return output
</code></pre>
<p>Basically when I compute stddev, I want it to do it over entire batch.</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/hij">@hij</a></p>
<p>Buffers are broadcast from rank 0 to other processes in the beginning of every forward pass. See the code below:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/1f8267931176cc9ecbf00493e5a359a57baca3df/torch/nn/parallel/distributed.py#L509-L513" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/1f8267931176cc9ecbf00493e5a359a57baca3df/torch/nn/parallel/distributed.py#L509-L513" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/1f8267931176cc9ecbf00493e5a359a57baca3df/torch/nn/parallel/distributed.py#L509-L513</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="509" style="counter-reset: li-counter 508 ;">
<li># Synchronize buffers across processes.</li>
<li># The process with rank 0 is considered the authoritative copy.</li>
<li>self._distributed_broadcast_coalesced(</li>
<li>    self.modules_buffers[0],</li>
<li>    self.broadcast_bucket_size)</li>
</ol></code></pre>


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

<p>If this is what you need, you can register the <code>stddev</code> as a buffer. If you need sth different (e.g., square <code>stddev</code> and sum it across all processes over the wire), you can call <a href="https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_reduce" rel="nofollow noopener"><code>allreduce</code></a> or <a href="https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_gather" rel="nofollow noopener"><code>allgather</code></a> in the forward function to do that.</p>

username: hij
cooked: <p>If i understand you correctly, registering it as a buffer only allows the stddev of rank 0 to be distributed to other processes. What I want to do is allow all tensors in all processes to contribute to calculating stddev.</p>

username: mrshenli
cooked: <aside class="quote no-group" data-username="hij" data-post="3" data-topic="79081">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hij/48/14113_2.png" class="avatar"> hij:</div>
<blockquote>
<p>If i understand you correctly, registering it as a buffer only allows the stddev of rank 0 to be distributed to other processes.</p>
</blockquote>
</aside>
<p>Yes.</p>
<blockquote>
<p>What I want to do is allow all tensors in all processes to contribute to calculating stddev.</p>
</blockquote>
<p>You can do this by using the collective communication APIs (allreduce/allgather) in the forward pass. One caveat here is that, the collective communication API requires all processes in the same group to invoke the same API in the same order, otherwise, it might hang or the result might be wrong. If you are not sure whether the allreduce/allgather for <code>stddev</code> would interleave with other collectives, you can use the <a href="https://pytorch.org/docs/stable/distributed.html#torch.distributed.new_group" rel="noopener nofollow ugc"><code>new_group</code></a> API to create a dedicated process group for collecting/summing <code>stddev.</code></p>

username: hij
cooked: <p>I have read that all_gather do not retain the gradient information. For my application, I want stddev to be part of the computation graph. How would I go about doing this? And could you also point me to an example/tutorial for these usage. I have not done any distributed training so I am not sure how to use these functions.</p>

username: mrshenli
cooked: <p>I am trying to understand the use case here:</p>
<blockquote>
<p>I have read that all_gather do not retain the gradient information. For my application, I want stddev to be part of the computation graph. How would I go about doing this?</p>
</blockquote>
<p>IIUC, <code>stddev</code> in this case is an intermediate output in <code>forward</code> and it’s <strong>not</strong> a model parameter? So you need its gradient during the backward pass to compute parameter gradients, but you don’t need to retain its gradient for optimizer <code>step()</code>? If above is true, why existing parameter gradient synchronization in DDP not sufficient?</p>
<blockquote>
<p>And could you also point me to an example/tutorial for these usage.</p>
</blockquote>
<p>Sure, below is the tutorial, and please search for “All-Reduce example.”</p>
<p><a href="https://pytorch.org/tutorials/intermediate/dist_tuto.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/dist_tuto.html</a></p>

username: hij
cooked: <p>I might have been confused. So this code snippet is part of my model which I intend to train. However, since I have limited gpu memory, I can only train with batch size of 1. And there isn’t a point to calculate stddev over a batchsize of 1 even if I do DDP. So in this case, would the allgather allreduce work?</p>
<aside class="quote no-group" data-username="hij" data-post="1" data-topic="79081">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hij/48/14113_2.png" class="avatar"> hij:</div>
<blockquote>
<pre><code class="lang-auto">def forward(self, input):
        feats = conv(input)
        batch, channel, height, width = feats.shape
        stddev = feats.view(batch, -1, 1, channel, height, width)
        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
        stddev = stddev.repeat(batch, 1, height, width)
        feats = torch.cat([feats, stddev], 1)
        output = conv_last(feats)
        return output
</code></pre>
</blockquote>
</aside>

username: mrshenli
cooked: <aside class="quote no-group" data-username="hij" data-post="7" data-topic="79081">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hij/48/14113_2.png" class="avatar"> hij:</div>
<blockquote>
<p>However, since I have limited gpu memory, I can only train with batch size of 1. And there isn’t a point to calculate stddev over a batchsize of 1 even if I do DDP. So in this case, would the allgather allreduce work?</p>
</blockquote>
</aside>
<p>I see. I am not sure if the result would still be correct in this case even if <code>allgather</code> and <code>allreduce</code> can retain gradients. IIUC, if this is trained without DDP (assume there are large enough GPU memory), then both <code>feats</code> and <code>stddev</code> are calculated based on all inputs. When trained with DDP, <code>feats</code> are now only derived from local inputs, and you would like to have <code>stddev</code> to be based on global inputs. So, when you cat <code>feats</code> and <code>stddev</code>, the output of the forward now represents a different thing. I am not sure if the loss function can handle that. Even if it can, what does averaging gradient mean in this case?</p>
<p>If above (local <code>feats</code> + global <code>stddev</code>) is the expected behavior, there might be a few ways to implement this.</p>
<ol>
<li>
<p>Implement a custom autograd function. E.g., its forward function can use an <code>allgather</code> to collect <code>stddev</code> from all processes, and its backward function can use another <code>allgather</code> to collect gradients and then extract the part belongs to the local <code>stddev</code> and sum them up.</p>
</li>
<li>
<p>Use <a href="https://pytorch.org/docs/stable/rpc.html" rel="noopener nofollow ugc"><code>torch.distributed.rpc</code></a>. There can be a master and a few workers, where each worker calculates the <code>feats</code> and <code>stddev</code> for its own input, and then the master gathers all <code>feats</code> and <code>stddev</code> to compute the final loss. Some more tutorials for this:<br>
a. <a href="https://pytorch.org/tutorials/intermediate/rpc_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Getting Started with Distributed RPC Framework — PyTorch Tutorials 2.1.1+cu121 documentation</a><br>
b. <a href="https://pytorch.org/tutorials/intermediate/rpc_param_server_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Implementing a Parameter Server Using Distributed RPC Framework — PyTorch Tutorials 2.1.1+cu121 documentation</a></p>
</li>
</ol>

username: hij
cooked: <p>In this case, since there is no batch dependencies in feat, would it be different if it is local feat + global stddev vs global feat + global stddev? Since stddev is concatenated to each feat tensor separately, they should have the same effect?</p>
<p>I will try your suggestions. Thank you!</p>