DistributedDataParallel: model weights and grads not synchronized with multiple forward backward pass

zzzf: <p>To replicate, change only def demo_basic(rank, world_size) in <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/ddp_tutorial.html</a> to the following:</p>
<pre><code class="lang-auto">def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=1)
    
    optimizer.zero_grad()
    
    outputs = {}
    outputs['0'] = ddp_model(torch.rand(20, 10))
    outputs['1'] = ddp_model(torch.rand(20, 10))
    outputs['2'] = ddp_model(torch.rand(20, 10))

    labels = torch.rand(20, 5).to(rank)

    for i in range(3):
        print(f"before {i}, rank: {rank}, weight: {ddp_model.module.net1.weight[0][0]}")
                 
        if i &lt; 2:
            loss_fn(outputs[str(i)], labels).backward(retain_graph=True)
        else:
            loss_fn(outputs[str(i)], labels).backward()

        print(f"after {i}, rank: {rank}, weight: {ddp_model.module.net1.weight[0][0]}, grad: {ddp_model.module.net1.weight.grad[0][0]}")
                
    optimizer.step()
    print(f"last, rank: {rank}, weight: {ddp_model.module.net1.weight[0][0]}, grad: {ddp_model.module.net1.weight.grad[0][0]}")
    
    cleanup()
</code></pre>
<p>and the output is:</p>
<pre><code class="lang-auto">before 0, rank: 0, weight: 0.1450435221195221
before 0, rank: 3, weight: 0.1450435221195221
before 0, rank: 1, weight: 0.1450435221195221
before 0, rank: 2, weight: 0.1450435221195221
after 0, rank: 0, weight: 0.1450435221195221, grad: -0.018003715202212334
before 1, rank: 0, weight: 0.1450435221195221
after 0, rank: 3, weight: 0.1450435221195221, grad: -0.018003715202212334
after 0, rank: 1, weight: 0.1450435221195221, grad: -0.018003715202212334
after 0, rank: 2, weight: 0.1450435221195221, grad: -0.018003715202212334
before 1, rank: 1, weight: 0.1450435221195221
before 1, rank: 3, weight: 0.1450435221195221
before 1, rank: 2, weight: 0.1450435221195221
after 1, rank: 0, weight: 0.1450435221195221, grad: -0.03955963999032974
after 1, rank: 3, weight: 0.1450435221195221, grad: -0.03072114661335945
before 2, rank: 0, weight: 0.1450435221195221
before 2, rank: 3, weight: 0.1450435221195221
after 1, rank: 1, weight: 0.1450435221195221, grad: -0.03775426745414734
before 2, rank: 1, weight: 0.1450435221195221
after 1, rank: 2, weight: 0.1450435221195221, grad: -0.03235533833503723
before 2, rank: 2, weight: 0.1450435221195221
after 2, rank: 0, weight: 0.1450435221195221, grad: -0.06408560276031494
after 2, rank: 3, weight: 0.1450435221195221, grad: -0.04222358390688896
after 2, rank: 1, weight: 0.1450435221195221, grad: -0.056242190301418304
last, rank: 0, weight: 0.20912912487983704, grad: -0.06408560276031494
last, rank: 3, weight: 0.18726710975170135, grad: -0.04222358390688896
last, rank: 1, weight: 0.201285719871521, grad: -0.056242190301418304
after 2, rank: 2, weight: 0.1450435221195221, grad: -0.04413666948676109
last, rank: 2, weight: 0.1891801953315735, grad: -0.04413666948676109
</code></pre>
<p>Weights and grads do not seem to be synchronized.</p>

mrshenli: <p>Hey <a class="mention" href="/u/zzzf">@zzzf</a>,</p>
<p>Does this problem persistent, if you change the flow of fw(1)-&gt;fw(2)-&gt;fw(3)-&gt;bw(1)-&gt;bw(2)-&gt;bw(3) to fw([1, 2, 3]) -&gt; bw([1, 2, 3])?</p>
<p>BTW, which version of PyTorch are you using? If it’s &lt;=v1.6, I would expect it throws an error here:<br>
</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/b31f58de6fa8bbda5353b3c77d9be4914399724d/torch/csrc/distributed/c10d/reducer.cpp#L724" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/b31f58de6fa8bbda5353b3c77d9be4914399724d/torch/csrc/distributed/c10d/reducer.cpp#L724" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/b31f58de6fa8bbda5353b3c77d9be4914399724d/torch/csrc/distributed/c10d/reducer.cpp#L724</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="714" style="counter-reset: li-counter 713 ;">
<li>// want to start performing reductions on `torch.autograd.backward()`.</li>
<li>void Reducer::prepare_for_backward(</li>
<li>    const std::vector&lt;torch::autograd::Variable&gt;&amp; outputs) {</li>
<li>  std::lock_guard&lt;std::mutex&gt; lock(mutex_);</li>
<li>  std::unordered_set&lt;torch::autograd::Node*&gt; seen;</li>
<li>  std::vector&lt;torch::autograd::Node*&gt; queue;</li>
<li>
</li><li>  // Check that any prior reduction has finished.</li>
<li>  // The variable `require_finalize_` is true until all gradients</li>
<li>  // have been computed and reduction of all buckets has been kicked off.</li>
<li class="selected">  if (require_finalize_) {</li>
<li>    TORCH_CHECK(</li>
<li>        false,</li>
<li>        "Expected to have finished reduction in the prior iteration before ",</li>
<li>        "starting a new one. ",</li>
<li>        "",</li>
<li>        "This error indicates that your module has parameters that were ",</li>
<li>        "not used in producing loss. ",</li>
<li>        "",</li>
<li>        "You can enable unused parameter detection by (1) passing the keyword "</li>
<li>        "argument `find_unused_parameters=True` to ",</li>
</ol></code></pre>


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


zzzf: <p>Hi <a class="mention" href="/u/mrshenli">@mrshenli</a>,</p>
<p>Thanks for your reply. My torch version is 1.6.0 and I receive no warnings/errors.</p>
<p>I’ve tried:</p>
<pre><code class="lang-auto">for _ in range(3):
    output = ddp_model(torch.rand(20, 10)
    loss_fn.backward(output, labels)

optimizer.step()
</code></pre>
<p>This works as I expected.</p>
<p>When you said fw([1, 2, 3]) -&gt; bw([1, 2, 3]), do you mean the following?</p>
<pre><code class="lang-auto">outputs = (ddp_model(torch.rand(20, 10)), ddp_model(torch.rand(20, 10)), ddp_model(torch.rand(20, 10)))

for i in range(3):
     (loss_fn(outputs[0], labels).backward(retain_graph=True),  loss_fn(outputs[1], labels).backward(retain_graph=True), loss_fn(outputs[2], labels).backward(retain_graph=True))

optimizer.step()
</code></pre>
<p>This still doesn’t synchronize the weights and doesn’t throw any error.</p>

pritamdamania87: <p><a class="mention" href="/u/mrshenli">@mrshenli</a> This seems to be a gap in DDP where it doesn’t support running backward twice? I couldn’t find any tests that use <code>retain_graph=True</code> with DDP.</p>
<p>The problem seems to be this line: <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L530" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L530</a>, which skips any gradient reduction. This variable is set to false after the first backward is done: <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L1152" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L1152</a> and then never set to True again since <code>prepare_for_backward</code> is not called anymore.</p>
<p><a class="mention" href="/u/zzzf">@zzzf</a> Is the workaround you mentioned in your previous reply sufficient for now?</p>