Pt_main_thread occasionally gets stuck, leading to training time increased

xgbj: <p>Hello everyone, after expanding our training data scale, we noticed a <strong>significant increase</strong> in the time per iteration. Using the nsys profiler, we found that after loading a large dataset into memory for training, the main training thread would randomly get stuck. You can see the situation in the image below.</p>
<ul>
<li>Initially, we suspected it was a communication issue, as most of the cards were stuck on NCCL communication operators.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2afc851e8d5e13a3b7ec514e48189417e5182bb9.png" data-download-href="https://discuss.pytorch.org/uploads/default/2afc851e8d5e13a3b7ec514e48189417e5182bb9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2afc851e8d5e13a3b7ec514e48189417e5182bb9_2_690x159.png" alt="image" data-base62-sha1="68h4FFYexCfvwuaO0cMnGlAAs7n" width="690" height="159" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2afc851e8d5e13a3b7ec514e48189417e5182bb9_2_690x159.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2afc851e8d5e13a3b7ec514e48189417e5182bb9_2_1035x238.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2afc851e8d5e13a3b7ec514e48189417e5182bb9_2_1380x318.png 2x" data-dominant-color="C9D4D2"><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">4606×1062 425 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></li>
<li>After careful investigation, we discovered that the cause of the NCCL slowdown was due to some GPUs’ <strong><code>pt_main_thread</code> being in an unscheduled state during the iterations where the time increased</strong>. This caused an increase in the time taken for the forward phase. Meanwhile, the other GPUs, which were functioning normally, experienced NCCL slowdown because they had to wait for the GPUs corresponding to the stuck <code>pt_main_thread</code> .<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/c/dcde343f1997eb90a58c18760973a36bb5844595.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/dcde343f1997eb90a58c18760973a36bb5844595" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dcde343f1997eb90a58c18760973a36bb5844595_2_690x253.jpeg" alt="image" data-base62-sha1="vvT72wQIXlXDG44VPOcvVAXI04d" width="690" height="253" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dcde343f1997eb90a58c18760973a36bb5844595_2_690x253.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dcde343f1997eb90a58c18760973a36bb5844595_2_1035x379.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dcde343f1997eb90a58c18760973a36bb5844595_2_1380x506.jpeg 2x" data-dominant-color="DEE4E0"><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">1920×704 121 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></li>
<li>We found that in most cases, the stuck state of the <code>pt_main_thread</code> depended on another thread that was being forked. Generally, once the other thread finished forking, the stuck <code>pt_main_thread</code> would resume scheduling.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/c/aca813f22a992f7cae23e5db8c6b0f3d17013009.png" data-download-href="https://discuss.pytorch.org/uploads/default/aca813f22a992f7cae23e5db8c6b0f3d17013009" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/aca813f22a992f7cae23e5db8c6b0f3d17013009_2_690x243.png" alt="image" data-base62-sha1="oDoegQ0FsnPgfY62ZCkTIv91Ey5" width="690" height="243" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/aca813f22a992f7cae23e5db8c6b0f3d17013009_2_690x243.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/aca813f22a992f7cae23e5db8c6b0f3d17013009_2_1035x364.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/aca813f22a992f7cae23e5db8c6b0f3d17013009_2_1380x486.png 2x" data-dominant-color="D8E3DD"><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">3190×1124 195 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></li>
<li>From the perspective of the nsys profiler, we could only see multiple threads under python3.8, and there were many threads with the same name as “pt_main_thread.” We couldn’t gain any information from the thread names. Additionally, the call stacks of each thread did not provide much valuable information.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/1/315256b99b46f0b1c2e780eff299d154ae12ede6.png" data-download-href="https://discuss.pytorch.org/uploads/default/315256b99b46f0b1c2e780eff299d154ae12ede6" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/315256b99b46f0b1c2e780eff299d154ae12ede6_2_174x500.png" alt="image" data-base62-sha1="72jNYsQgTQN8muu2ZcLCYig2dsW" width="174" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/315256b99b46f0b1c2e780eff299d154ae12ede6_2_174x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/315256b99b46f0b1c2e780eff299d154ae12ede6_2_261x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/315256b99b46f0b1c2e780eff299d154ae12ede6_2_348x1000.png 2x" data-dominant-color="F1F4F2"><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">402×1150 62.5 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></li>
</ul>
<p>This phenomenon did not occur when we used smaller datasets or fewer GPUs for training. This issue is quite peculiar, and we haven’t encountered it before. We would like to ask if anyone has experienced similar bugs or if there are any suggestions on how to approach analyzing this problem more effectively?<br>
We would like to ask for advice on how to further debug this issue and identify the root cause.</p>