The program hangs when using DDP with multiple GPUs

JohnsonLC: <h3><a name="problem-1" class="anchor" href="#problem-1"></a>Problem:</h3>
<p>My DDP program always hangs when performing backward for some specific inputs.</p>
<h3><a name="background-2" class="anchor" href="#background-2"></a>Background:</h3>
<p>Simply put, I have a CNN network called <strong>‘net_cnn’</strong>, and a MLP network called <strong>‘net_mlp’</strong>.<br>
The inputs are first feed into ‘net_cnn’, generating outputs called <strong>‘out_cnn’</strong>. ‘out_cnn’ is then feed into ‘net_mlp’, generating outputs called <strong>‘out_mlp’</strong>.</p>
<p><strong>‘loss1’</strong> is computed using ‘out_cnn’. <strong>‘loss2’</strong> is computed using ‘out_mlp’.</p>
<p>A sketch map is shown below.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/9/99b828fa17241ec3727ad5611b57488c27c65b2d.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/99b828fa17241ec3727ad5611b57488c27c65b2d" title="1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/99b828fa17241ec3727ad5611b57488c27c65b2d_2_690x225.jpeg" alt="1" data-base62-sha1="lVRC4BmttilwOkpD9a6stCcpREN" width="690" height="225" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/99b828fa17241ec3727ad5611b57488c27c65b2d_2_690x225.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/99b828fa17241ec3727ad5611b57488c27c65b2d_2_1035x337.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/9/99b828fa17241ec3727ad5611b57488c27c65b2d_2_1380x450.jpeg 2x" data-dominant-color="FAFAFA"><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">1</span><span class="informations">1462×477 58.7 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>
<h3><a name="the-question-3" class="anchor" href="#the-question-3"></a>The question:</h3>
<p>For some specific inputs, ‘out_mlp’ may output all zeros, and have nothing to do with ‘out_cnn’ and ‘net_mlp’. This makes ‘loss2’ has no gradients related to the parameters of ‘net_cnn’ and ‘net_mlp’. When using DDP with multiple GPUs, if one of these GPUs has the inputs mentioned above, the others will wait for the gradients when performing backward, which causes the problem of hanging.</p>
<p>So is there any solution for this kind of problem?  Thanks for any suggestions.</p>

H-Huang: <p>Will setting <a href="https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel" rel="noopener nofollow ugc"><code>find_unused_paramters</code></a> to <code>True</code> in DDP initialization fix your issue?</p>

JohnsonLC: <p>Unfortunately, it doesn’t work.</p>
<p>Now I have made the training process feasible by dropping those specific inputs. I detect the validity of ‘out_mlp’, and if there are one or more inputs in the current batch that cause all zeros outputs, I will drop the whole batch (means don’t perform ‘loss.backward()’). ‘torch.all_reduce()’ is employed to do the detection.</p>
<p>Thanks anyway for your kind suggestion! <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=12" title=":grinning:" class="emoji" alt=":grinning:" loading="lazy" width="20" height="20"></p>