Dist.all_gather stuck

111530: <p>I am using the communication hook to implement a simple top-k gradient compression that uses <code>all_gather</code> to gather the indices and values. However, I found that after some iterations, <code>all_gather</code> will hang and throw no errors.</p>
<p>To debug, I removed complicated operations, and only left the async <code>all_gather</code> call as below:</p>
<pre><code class="lang-python">    all_indices = torch.cat(all_indices)

    output_indices = [torch.zeros_like(all_indices, device=device, dtype=all_indices.dtype) for _ in range(world_size)]

    tmp = dist.all_gather(
        output_indices, all_indices, group=group_to_use, async_op=True
    ).get_future()

    def unpack_indices_and_all_gather_values(fut):
        output_indices = fut.value()[0]

        state.maybe_increase_iter(bucket)
        return [input_tensor]
 return (tmp.then(unpack_indices_and_all_gather_values))
</code></pre>
<p>the hanging problem still happens.<br>
I am registering this model to a resnet32 (DDP for sure) on cifar-10. I saw other people with similar problems were not making the output list on different devices, but I believe I’ve set them correctly in my case as in the code snippet.</p>

111530: <p>The pytorch version was 1.9. But I also tried on Pytorch 1.12 (changed some code as the APIs changed), same problem.</p>
<p>Also, as I was thinking if it is the problem of the all_gather primitive, I changed the first <code>all_reduce</code> to an <code>all_gather</code> then take the average, and remove the second <code>all_gather</code>. But it turned out won’t get stuck. So the problem will lie in my code. But I am really confused about what could be the possible reasons that make this code stuck.</p>

kwen2501: <p>Glad you found a workaround.<br>
To your original question – can you check if <code>all_indices</code> have the same size on all ranks? If not, this could easily lead to a hang in all_gather.</p>

111530: <p>Yes thank you! In my previous search, I didn’t know different sizes would lead to a hang too, and I just verified with my original code, yes there was a size mismatch because sometimes top-k is not exactly k values.</p>
<p>By the way, I still have two (unrelated) questions:</p>
<ol>
<li>the GradBucket’s API is changing among versions, in 1.9 there’s a get_tensor_per_parameter, but in 1.12 it’s gone and replaced by parameters() and gradients(), I wonder what is the difference between them? Because as far as I know, all tensors in GradBucket are gradients.</li>
<li>I seem to have spotted a bug in pytorch 1.9 (not 100% sure), I was running resnet32, it’s structure is as below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/f/6fc5dc4570d4753706d51af496ff18d044b5b72a.png" data-download-href="https://discuss.pytorch.org/uploads/default/6fc5dc4570d4753706d51af496ff18d044b5b72a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6fc5dc4570d4753706d51af496ff18d044b5b72a_2_307x500.png" alt="image" data-base62-sha1="fWMUCw3L4XEu8LGa6bsmAvTjTm2" width="307" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6fc5dc4570d4753706d51af496ff18d044b5b72a_2_307x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6fc5dc4570d4753706d51af496ff18d044b5b72a_2_460x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6fc5dc4570d4753706d51af496ff18d044b5b72a_2_614x1000.png 2x" data-dominant-color="40413B"><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">942×1532 157 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><br>
There were two buckets for this model, after the second iteration (I know buckets will be rebuilt after the first iteration, so this is the second iteration), the contents in both buckets are (bias tensor ignored):<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7c8b3078055cbc8961dcde908cb7cdf1ea65182.png" data-download-href="https://discuss.pytorch.org/uploads/default/e7c8b3078055cbc8961dcde908cb7cdf1ea65182" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7c8b3078055cbc8961dcde908cb7cdf1ea65182.png" alt="image" data-base62-sha1="x4si6uDD6IAH0XDpoGA3AhQY62K" width="690" height="392" data-dominant-color="383933"><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">848×482 32.1 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><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/0/c01163c0bff783c97d83fd53fc0c67209621d3fd.png" data-download-href="https://discuss.pytorch.org/uploads/default/c01163c0bff783c97d83fd53fc0c67209621d3fd" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c01163c0bff783c97d83fd53fc0c67209621d3fd_2_296x500.png" alt="image" data-base62-sha1="rp6VKGH3NfdeI1ehpYtCYpHwcy1" width="296" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c01163c0bff783c97d83fd53fc0c67209621d3fd_2_296x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c01163c0bff783c97d83fd53fc0c67209621d3fd_2_444x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c01163c0bff783c97d83fd53fc0c67209621d3fd_2_592x1000.png 2x" data-dominant-color="393A34"><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">770×1300 82.6 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>
</ol>
<p>You can see in bucket 1 the tensors are not reversed, and it falsely shared four [64, 64, 3, 3] tensors with bucket 0. After the third iteration, this problem was gone: the four tensors are unique in bucket 0 and tensors in bucket 1 are correctly reversed. So this only happens in iteration 2. In pytorch 1.12, tensors are arranged correctly in iteration 2 too.</p>

kwen2501: <p>Cc <a class="mention" href="/u/rohan_varma">@Rohan_Varma</a> <a class="mention" href="/u/mrshenli">@mrshenli</a> for the new question raised right above.</p>