Is it ok to use dist.barrier() in a custom torch.autograd.Function's forward and backward?

lzhbrian: <p>Is it ok to use dist.barrier() in the forward/backward of a custom pytorch torch.autograd.Fucntion?</p>
<p>I have a custom torch.autograd.Function looks like this (adding dist.barrier() because I am debuggling my code with other errors):<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b502e0bf0fd33b95a366e1014e9fe1a1c5679bb.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/3b502e0bf0fd33b95a366e1014e9fe1a1c5679bb" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b502e0bf0fd33b95a366e1014e9fe1a1c5679bb_2_690x349.jpeg" alt="image" data-base62-sha1="8sHXej6f9cm2gbOQb59dTyyG3wv" width="690" height="349" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b502e0bf0fd33b95a366e1014e9fe1a1c5679bb_2_690x349.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b502e0bf0fd33b95a366e1014e9fe1a1c5679bb_2_1035x523.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b502e0bf0fd33b95a366e1014e9fe1a1c5679bb.jpeg 2x" data-dominant-color="F6F7F8"><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">1342×680 164 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>I am encountering randomly training hangs where the error look like:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/1/719f8145dc4a7c9bb26aceb70d8b42b96ca16086.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/719f8145dc4a7c9bb26aceb70d8b42b96ca16086" title="Screenshot 2024-06-10 at 07.19.22"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/719f8145dc4a7c9bb26aceb70d8b42b96ca16086_2_690x324.jpeg" alt="Screenshot 2024-06-10 at 07.19.22" data-base62-sha1="gd9GU02vbfxRuwkhs3C4ngWkKkS" width="690" height="324" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/719f8145dc4a7c9bb26aceb70d8b42b96ca16086_2_690x324.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/719f8145dc4a7c9bb26aceb70d8b42b96ca16086_2_1035x486.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/719f8145dc4a7c9bb26aceb70d8b42b96ca16086_2_1380x648.jpeg 2x" data-dominant-color="EDEEED"><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">Screenshot 2024-06-10 at 07.19.22</span><span class="informations">1920×903 200 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 looks like rank0 is doing dist.barrier() written by me in the backward function above, while rank1 is doing all_reduce of some gradient by ddp itself.<br>
This error occurs randomly (5 in 200 same experiments with different seeds).</p>
<p>Is it possible that in rank0 and rank1, ddp all_reduce gradients asyncly? (so that rank0 enters my dist.barrier(), and rank1’s ddp is already starting to try to sync some gradients between gpus?)</p>
<p>p.s. some context:</p>
<ol>
<li>env: python=3.9.7, pytorch=2.3.0+cu121. Some flags added for showing the above debuggling msgs: TORCH_CPP_LOG_LEVEL=INFO TORCH_DISTRIBUTED_DEBUG=DETAIL TORCH_SHOW_CPP_STACKTRACES=1</li>
<li>I am using a lot of custom torch.autograd.Function to sync data and gradients between gpus in my training code.</li>
<li>Only rank0 is doing validation every several training iterations (rank1 waits and do nothing, I also add a dist.barrier there.)</li>
</ol>
<p>Thanks!</p>

lzhbrian: <p>update: I remove the dist.barrier() in my CustomFunc.backward() impl, and the problem still exists. The mismatch behavior now lies between: <code>my_custom_backward's all_reduce</code> and <code>ddp's all_reduce</code>.</p>
<p>I think the question could be update to: <strong>Are all kinds of sync/async ops among gpus allowed in custom torch.autograd.Function backward?</strong></p>
<p>In addition, I saw a PR <a href="https://github.com/pytorch/pytorch/pull/96365" class="inline-onebox" rel="noopener nofollow ugc">[DDP] Enable delayed all reduce in DDP by GD06 · Pull Request #96365 · pytorch/pytorch · GitHub</a>, where the motivation is : “Enable the functionality of delaying all reduce in DDP to specify the parameters whose all reduce will be hooked to a specific param. This prevents AllReduce blocking All2All in some recommendation models.”</p>

lzhbrian: <p>Solved, it’s another part of my code accidentally letting rank0 and rank1 computing different graph.</p>