Distributed Collectives

kr12: <p>Hi! I am newer to pytorch and am working with the distributed package.</p>
<p>I am trying to implement some parallelism functionalities but am running into some issues:</p>
<ul>
<li>I want to partition/shard a tensor. I know there’s some things hardcoded in that I also need to figure out, but for some reason dist.scatter is incorrectly producing [[1][3]] and [[2][3]] instead of [[1][3]] and [[2][4]]:</li>
</ul>
<pre><code class="lang-auto">if local_rank == 0:
    weight = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32).cuda(local_rank)
else:
    weight = None

world_size = dist.get_world_size(m_view_2)
chunks = None
if local_rank == 0:
    chunks = list(torch.chunk(weight, world_size, dim=1))

part_tensor = torch.empty((2, 1), dtype=torch.float32).cuda(local_rank) # need to figure out how to dynamically get shape

dist.scatter(part_tensor, scatter_list=chunks, src=0, group=m_view_2, async_op=False)
</code></pre>
<p>On a more general note, what i am trying to do is to take in a sequence of operations and execute those operations. For example:</p>
<ul>
<li>partition input tensor across devices 0 and 1</li>
<li>replicate weights across devices 0 and 1</li>
<li>multiply input tensor and weight</li>
<li>reduce the results onto device 0</li>
<li>replicate the result across devices 0 and 1</li>
</ul>
<p>I know there are some libraries out there that would make what I am outlining relatively simple, but my goal here is to be more explicit with each individual step instead of something like an all_reduce, etc. Not sure if anyone has thoughts on how I should approach this and any considerations I should probably have with it (async vs sync operations is something I likely have to be mindful of).</p>

fduwjj: <p>If you want to do more explicit, you can call reduce instead of all-reduce for " * reduce the results onto device 0". Is this something you want to achieve?</p>

kr12: <p>yes! For the each of the following operations, I believe the following collectives would be:</p>
<p>partitioning/sharding (ex. [[2, 3], [4, 5]] → [[2, 3]] and [[4, 5]]: scatter<br>
replicating (ex. [[2,3]] → [[2, 3]] and [[2, 3]]: broadcast<br>
combining (ex. [[2, 3]] and [[4, 5]] → [[2, 3], [4, 5]]): gather<br>
reduce (ex. [[2, 3]] and [[4, 5]] → [[6, 8]]): reduce</p>
<p>Im still unsure as to why the example from my original post isn’t creating the right output and whether or not the following collective operations will cause issues with autograd/gradient tracking.</p>

kr12: <p>edit: was able to resolve the issue in the original post by adding:</p>
<pre><code class="lang-auto">if local_rank == 0:
    chunks = list(torch.chunk(weight, world_size, dim=1))

    for i in range(len(chunks)):
        chunks[i] = chunks[i].contiguous()
</code></pre>
<p>Not entirely sure what the implications of this would be on gradient tracking though as it seems to be making a new tensor.</p>

fduwjj: <p>I got confused, why do you use scatter? Is it because originally your tensor is only in rank0 so you will need to call a scatter? If so I think your understanding is correct.</p>

fduwjj: <p>because the input of collective needs to be contiguous, while when you do split or chunk, each split might still keep the stride of the original tensor</p>