Distributed: scatter list of tensors of different sizes

username: ArchieGertsman
cooked: <p>Say I have two processes running, and that I have two tensors <code>t0</code> and <code>t1</code>, with sizes 13 and 6, respectively, stored in the process with rank zero. How can I scatter the tensor list <code>[t0, t1]</code> to ranks zero and one? When I try using <code>dist.scatter</code> I get the error</p>
<pre><code class="lang-auto">RuntimeError: ProcessGroupGloo::scatter: invalid tensor size at index 1 (expected (13), got (6))
</code></pre>
<p>so it seems that <code>dist.scatter</code> expects the output tensor size to be fixed among all the processes.</p>

username: H-Huang
cooked: <p>Since t0 and t1 are stored on rank 0, you wouldn’t need to scatter t0 to itself right? In this case, I would just <a href="https://pytorch.org/docs/master/distributed.html#torch.distributed.send" rel="noopener nofollow ugc">send</a> from t1 from rank 0 and <a href="https://pytorch.org/docs/master/distributed.html#torch.distributed.recv" rel="noopener nofollow ugc">recv</a> t1 from rank 1.</p>
<p>Regarding the documentation, I added an issue to track this to make it clearer for <code>scatter()</code> <a href="https://github.com/pytorch/pytorch/issues/84566" class="inline-onebox" rel="noopener nofollow ugc">Update dist.scatter() documentation to provide clearer arg description · Issue #84566 · pytorch/pytorch · GitHub</a></p>

username: ArchieGertsman
cooked: <p>You’re right, but in reality I have more than two processes; I simplified it to two for the sake of this post. Thanks for opening the issue.</p>

username: kumpera
cooked: <p>Scatter expects all tensors to be of the same size. This can be addressed by padding and some additional comms.</p>
<p>The trick is to broadcast the max tensor size and then</p>
<p>You can solve this by broadcasting the max size among tensor sizes. Something like the following:</p>
<pre><code class="lang-auto">    is_rank0 = dist.get_rank() == 0
    tensors = [torch.rand(2), torch.rand(3), torch.rand(4), torch.rand(5)] if is_rank0 else None

    # find out the largest tensor that will be used for scattering
    max_len = max(t.numel() for t in tensors) if is_rank0 else 0
    print(f"{dist.get_rank()} local max len {max_len}")

    # broadcast the receive size
    tensor_len = torch.tensor([max_len])
    dist.broadcast(tensor_len, src=0)
    print(f"{dist.get_rank()} max len {tensor_len}")

    # scatter each rank's size
    tensors_len = [torch.tensor([t.numel()], dtype=torch.int64) for t in tensors] if is_rank0 else None
    my_len = torch.zeros(1, dtype=torch.int64)
    dist.scatter(my_len, tensors_len)
    print(f"{dist.get_rank()} my len is {my_len.item()}")

    # zero pad input tensors
    if tensors is not None:
        tensors = [pad(t, (0, max_len - t.numel()), "constant", 0) for t in tensors]

    receive_tensor = torch.empty(tensor_len)
    dist.scatter(receive_tensor, tensors)
    receive_tensor = receive_tensor[0: my_len]
    print(f"{dist.get_rank()} -&gt; final len: {receive_tensor.numel()}")
</code></pre>

username: kumpera
cooked: <p>This usage pattern is common enough that PyTorch should directly support it. Filed a feature request for it: <a href="https://github.com/pytorch/pytorch/issues/84593" class="inline-onebox" rel="noopener nofollow ugc">Uneven and/or Dynamically sized collectives · Issue #84593 · pytorch/pytorch · GitHub</a></p>