How can I receive the outputs from dist.all_gather_object() asynchronously?

Mirian_Hipolito: <p>I’m using NCCL with dist.all_gather_object() to send back the outputs from all workers to rank 0. However the code below work asynchronously, rank 0 can only reads the output until <strong>all the workers</strong>  have sent their response. The script below just hangs while waiting … it doesn’t allow me to access to output. So,  I would like to implement it in an asyncrhonously way, let’s say that if I uncomment the <em>‘time.sleep()’</em>  for rank 2, I can be able to read in worker 0 the output for worker 1 in the meantime, while waiting for the worker 2 to finish.</p>
<pre><code class="lang-auto">if rank == 0:
        output  = {'acc': 1.89, 'loss': 1.2, 'extra_metric':10.0}
        gather_objects = [output for i in range(world)] # any picklable object
        output = [None for _ in gather_objects]
        dist.all_gather_object(output, gather_objects[dist.get_rank()])
        print(f"Rank output {rank} {output}")
    else:
                
        if rank == 2:
            output  = {'acc': 2.89, 'loss': 2.2, 'extra_metric':20.0}
            gather_objects = [output for i in range(world)]
            #time.sleep(5)
        else:
            output = {'acc': 3.89, 'loss': 3.2, 'extra_metric':30.0}
            gather_objects = [output for i in range(world)]

        output = [None for _ in gather_objects]
        dist.all_gather_object(output, gather_objects[dist.get_rank()])
        print(f"Rank output {rank} {output}")
</code></pre>
<p>If I execute the code above, I’m getting the following once all workers have sent back their output:</p>
<pre><code class="lang-auto">Rank output 0 [{'acc': 1.89, 'loss': 1.2, 'extra_metric': 10.0}, {'acc': 3.89, 'loss': 3.2, 'extra_metric': 30.0}, {'acc': 2.89, 'loss': 2.2, 'extra_metric': 20.0}]
Rank output 2 [{'acc': 1.89, 'loss': 1.2, 'extra_metric': 10.0}, {'acc': 3.89, 'loss': 3.2, 'extra_metric': 30.0}, {'acc': 2.89, 'loss': 2.2, 'extra_metric': 20.0}]
Rank output 1 [{'acc': 1.89, 'loss': 1.2, 'extra_metric': 10.0}, {'acc': 3.89, 'loss': 3.2, 'extra_metric': 30.0}, {'acc': 2.89, 'loss': 2.2, 'extra_metric': 20.0}]
</code></pre>
<p>I need to mandatory send a python object (the dictionary) I cannot use tensors. This is the reason I’m using all_gather_object(), also I tried to use gather_object() instead but this one doesn’t work for NCCL.</p>

rvarm1: <p>Hi, currently there is no way to make these APIs asynchronous, but filed a feature request asking for it: <a href="https://github.com/pytorch/pytorch/issues/80417" class="inline-onebox" rel="noopener nofollow ugc">[c10d] Async object-based collectives · Issue #80417 · pytorch/pytorch · GitHub</a></p>
<p>Also, as of PyTorch 1.12, NCCL now supports <code>gather</code> API, so <code>gather_object</code> is enabled. You should be able to use <code>gather_object</code> in nightly builds or building from source, or in 1.12 when it is released.</p>

Mirian_Hipolito: <p>This is great, thank you so much <a class="mention" href="/u/rvarm1">@rvarm1</a> ! However <code>gather_object</code> also uses collective operations for the output list and according to the documentation stills blocking any other call in the meantime <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/9/e94a7368f26d2db353abd66f6921c39a4346e172.png" data-download-href="https://discuss.pytorch.org/uploads/default/e94a7368f26d2db353abd66f6921c39a4346e172" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/9/e94a7368f26d2db353abd66f6921c39a4346e172_2_690x128.png" alt="image" data-base62-sha1="xhMLamFbL8QRGmWJForvWjl60rE" width="690" height="128" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/9/e94a7368f26d2db353abd66f6921c39a4346e172_2_690x128.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/9/e94a7368f26d2db353abd66f6921c39a4346e172_2_1035x192.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/9/e94a7368f26d2db353abd66f6921c39a4346e172.png 2x" data-dominant-color="D7EBF3"><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">1077×200 9.21 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>