Loss.backward() hangs at the second iteration when using DDP

Ferd: <p>The network works well when using a single GPU.</p>
<pre><code class="lang-auto">  for cur_it in range(total_it_each_epoch):
      try:
          batch = next(dataloader_iter)
      except StopIteration:
          dataloader_iter = iter(train_loader)
          batch = next(dataloader_iter)
          print('new iters')

      lr_scheduler.step(accumulated_iter)

      if rank == 0:
          print('\n---data passed---')

      try:
          cur_lr = float(optimizer.lr)
      except:
          cur_lr = optimizer.param_groups[0]['lr']

      model.train()
      optimizer.zero_grad()

      loss, tb_dict, disp_dict = model_func(model, batch)

      if rank == 0:
          print('\n---forward passed---')

      print(loss)
      
      loss.backward()

      optimizer.step()

      if rank == 0:
          print('\n---backward passed---')

      time.sleep(0.01)
</code></pre>
<p>Here is the output:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbe2a282b5bd06b92fa6e473d482f8e26a9d4429.png" data-download-href="https://discuss.pytorch.org/uploads/default/bbe2a282b5bd06b92fa6e473d482f8e26a9d4429" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbe2a282b5bd06b92fa6e473d482f8e26a9d4429_2_690x420.png" alt="image" data-base62-sha1="qO6QGFaiMsnoReWJwKAmAafoLTH" width="690" height="420" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbe2a282b5bd06b92fa6e473d482f8e26a9d4429_2_690x420.png, https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbe2a282b5bd06b92fa6e473d482f8e26a9d4429.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbe2a282b5bd06b92fa6e473d482f8e26a9d4429.png 2x" data-dominant-color="F3F3F3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">813×495 60.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>it works well in the first iteration; however, it hangs at the second iteration (sometimes the third iteration).<br>
I am really confused, any help would be appreciated.</p>

agu: <p>Could you provide a minimal script to reproduce the hang? In particular, can you include the model construction (e.g. how are you using <code>DistributedDataParallel</code>)?</p>

Ferd: <p>Thanks for your reply!<br>
I identify the problem by skipping some code blocks. It could be caused by a customed operator (CUDA extension).<br>
But I cannot find the specific reason.</p>

ptrblck: <p>Skipping layers is not allowed without using <code>find_unused_parameters=True</code> since the backward pass will wait for the gradient of all registered parameters to be ready that that it can communicate it.<br>
From the <a href="https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel">docs</a>:</p>
<blockquote>
<p><strong>find_unused_parameters</strong> (<a href="https://docs.python.org/3/library/functions.html#bool"><em>bool</em></a>) – Traverse the autograd graph from all tensors contained in the return value of the wrapped module’s <code>forward</code> function. Parameters that don’t receive gradients as part of this graph are preemptively marked as being ready to be reduced. In addition, parameters that may have been used in the wrapped module’s <code>forward</code> function but were not part of loss computation and thus would also not receive gradients are preemptively marked as ready to be reduced. (default: <code>False</code>)</p>
</blockquote>

Ferd: <p>Thanks. But the problem occurs when all parameters are used to compute loss.<br>
I skip some code blocks and pass <code>find_unused_parameters=True</code>,  it is exactly the method I used to locate the bug.<br>
And I find out the bug is related to a customed operator.</p>