Dist.all_gather with uneven tensor sizes

JVP15: <p>The docs for <a href="https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.all_gather" rel="noopener nofollow ugc">torch.distributed.all_gather</a> state that</p>
<blockquote>
<p>Complex and uneven sized tensors are supported.</p>
</blockquote>
<p>but I can’t find any documentation or examples about how to gather unevenly sized tensors. I have a script like this:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

if __name__ == '__main__':
    torch.distributed.init_process_group(backend="cpu:gloo,cuda:nccl")

    world_size = torch.distributed.get_world_size()
    rank = torch.distributed.get_rank()

    # should be [0] for rank 0, [1, 1] for rank 1, [2, 2, 2] for rank 2, etc.
    tensor = torch.ones((rank + 1, )) * rank

    # Gather tensors from all ranks
    gathered_tensors = [torch.empty((i + 1, )) for i in range(world_size)]
    torch.distributed.all_gather(gathered_tensors, tensor)

    if rank == 0:
        print(f"Gathered tensors from all ranks")
        for i, t in enumerate(gathered_tensors):
            print(f"Rank {i}: {t}")
</code></pre>
<p>When I run it with <code>torchrun --nproc-per-node 42 gather_tensor_test.py</code>, I get all kinds of fun errors like:</p>
<blockquote>
<p>[rank0]: Traceback (most recent call last):<br>
[rank0]:   File “/mnt/aged-glade/home/jordan/lab-asr/lab/model/asr_8_rl/all_gather_test.py”, line 14, in <br>
[rank0]:     torch.distributed.all_gather(gathered_tensors, tensor)<br>
[rank0]:   File “/mnt/aged-glade/home/jordan/.conda/envs/asr/lib/python3.12/site-packages/torch/distributed/c10d_logger.py”, line 81, in wrapper<br>
[rank0]:     return func(*args, **kwargs)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/mnt/aged-glade/home/jordan/.conda/envs/asr/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py”, line 3728, in all_gather<br>
[rank0]:     work = group.allgather([tensor_list], [tensor])<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]: RuntimeError: ProcessGroupGloo::allgather: invalid tensor size at index 1 (expected (1), got (2))<br>
[rank1]: Traceback (most recent call last):<br>
[rank1]:   File “/mnt/aged-glade/home/jordan/lab-asr/lab/model/asr_8_rl/all_gather_test.py”, line 14, in <br>
[rank1]:     torch.distributed.all_gather(gathered_tensors, tensor)<br>
[rank1]:   File “/mnt/aged-glade/home/jordan/.conda/envs/asr/lib/python3.12/site-packages/torch/distributed/c10d_logger.py”, line 81, in wrapper<br>
[rank1]:     return func(*args, **kwargs)<br>
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^<br>
[rank1]:   File “/mnt/aged-glade/home/jordan/.conda/envs/asr/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py”, line 3728, in all_gather<br>
[rank1]:     work = group.allgather([tensor_list], [tensor])<br>
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank1]: RuntimeError: ProcessGroupGloo::allgather: invalid tensor size at index 0 (expected (2), got (1))<br>
…</p>
</blockquote>
<p>So how exactly do we use <code>dist.all_gather</code> with uneven tensor sizes?</p>

yf225: <p>I believe this is a GLOO backend limitation, cc. <a class="mention" href="/u/d4l3k">@d4l3k</a></p>
<p>It doesn’t repro with cuda NCCL backend:</p>
<pre><code class="lang-auto">import torch

if __name__ == '__main__':
    torch.distributed.init_process_group(backend="nccl")

    world_size = torch.distributed.get_world_size()
    rank = torch.distributed.get_rank()
    torch.cuda.set_device(rank)

    # should be [0] for rank 0, [1, 1] for rank 1, [2, 2, 2] for rank 2, etc.
    tensor = torch.ones((rank + 1, ), device="cuda") * rank

    # Gather tensors from all ranks
    gathered_tensors = [torch.empty((i + 1), device="cuda") for i in range(world_size)]
    torch.distributed.all_gather(gathered_tensors, tensor)

    if rank == 0:
        print(f"Gathered tensors from all ranks")
        for i, t in enumerate(gathered_tensors):
            print(f"Rank {i}: {t}")
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">Gathered tensors from all ranks
Rank 0: tensor([0.], device='cuda:0')
Rank 1: tensor([1., 1.], device='cuda:0')
Rank 2: tensor([2., 2., 2.], device='cuda:0')
</code></pre>