What's the right way to update a field of a DistributedDataParallel object during training?

username: epignatelli
cooked: <p>What is the correct way to update a class variable once the model has been wrapped around a <code>DistributedDataParallel</code>?</p>
<p>In the case below, if we take a snapshot of <code>self.k</code> for each model in each gpu, at the same time, we can get different results.</p>
<p>Any idea why that happens and how to solve that?</p>
<p>I guess <code>loss</code> would be different across all the models?</p>
<pre><code>class Model(torch.nn.Module):
    def __init__(self):
        self.fc = torch.nn.Linear(128, 128)
        self.register_buffer("k", 0)
        self.callback = lambda x: x + 1

    def forward(self, x):
        print(self.k)
        return self.fc(x)

    def training_step(self, x, y):
        y_hat = self(x)
        loss = torch.nn.functional.binary_cross_entropy(y_hat, y)
        loss.backward()
        if loss &gt; 1.:
            self.callback(loss, self.k)</code></pre>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/epignatelli">@epignatelli</a></p>
<p>When did you take the snapshot of <code>self.k</code>? <code>DDP</code> will broadcast all buffers from rank 0 process to other processes right before calling <code>Mode.forward</code>. See the code below. So given the above code, the buffer should be consistent across all processes before the <code>self.callback</code> is launched.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/79cfd85987c128b672850b64d6680719e4eda552/torch/nn/parallel/distributed.py#L695-L713" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/79cfd85987c128b672850b64d6680719e4eda552/torch/nn/parallel/distributed.py#L695-L713" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/79cfd85987c128b672850b64d6680719e4eda552/torch/nn/parallel/distributed.py#L695-L713</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="695" style="counter-reset: li-counter 694 ;">
<li># module buffer sync</li>
<li>if self.broadcast_buffers and len(self.modules_buffers[0]) &gt; 0:</li>
<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>
<li>    # only do intra-node buffer sync for replicated single-device</li>
<li>    # CUDA modules</li>
<li>    if self.device_ids and len(self.device_ids) &gt; 1:</li>
<li>        # intra-node buffer sync</li>
<li>        result = comm.broadcast_coalesced(</li>
<li>            self.modules_buffers[0],</li>
<li>            self.device_ids,</li>
<li>            self.broadcast_bucket_size)</li>
<li>        for tensors, module_buffers in zip(result[1:],</li>
<li>                                           self.modules_buffers[1:]):</li>
<li>            for tensor, buffer in zip(tensors, module_buffers):</li>
<li>                buffer.set_(tensor)</li>
</ol></code></pre>


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


username: mrshenli
cooked: <p>BTW, could you please add a “distributed” tag to distributed-training related questions? So that people working on it can get back to you promptly. Thanks!</p>