How does reduce_scatter work?

Vedant_Roy: <p>I’m trying to understand how <code>reduce_scatter</code> works in Pytorch. So far, I have found the following code snippet:</p>
<pre><code class="lang-auto">class AllGatherFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, reduce_dtype: torch.dtype = torch.float32):
        ctx.reduce_dtype = reduce_dtype

        output = list(torch.empty_like(tensor) for _ in range(dist.get_world_size()))
        dist.all_gather(output, tensor)
        output = torch.cat(output, dim=0)
        return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        grad_dtype = grad_output.dtype
        input_list = list(grad_output.to(ctx.reduce_dtype).chunk(dist.get_world_size()))
        grad_input = torch.empty_like(input_list[dist.get_rank()])
        dist.reduce_scatter(grad_input, input_list)
        return grad_input.to(grad_dtype)


def all_gather(tensor):
    return AllGatherFunction.apply(tensor)
</code></pre>
<p>which is from OpenAI’s CLIP: <a href="https://github.com/openai/CLIP/issues/111" class="inline-onebox" rel="noopener nofollow ugc">how to use multiple GPUs,the default is to use the first CUDA device · Issue #111 · openai/CLIP · GitHub</a>.<br>
I’m trying to figure out what <code>reduce_scatter</code> does by staring at this code snippet, and I have no idea. Could someone help me out?</p>

ptrblck: <p><a href="https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/operations.html#reducescatter">These docs</a> might be helpful as they visualize communication patterns.</p>

Vedant_Roy: <p>Thanks! That was a good link. So to clarify, from the docs, the <code>inX</code> refers to the concatenation of the input tensor from each rank. (Specifically a column-wise concat, so the # of rows remains the same but the length of each row increases).</p>
<p>There are also some magic variables, like <code>count</code> and <code>Y</code> that I don’t quite understand.<br>
Currently, writing up some toy code to figure this out.</p>

Vedant_Roy: <p>Sample code for understanding:</p>
<pre><code class="lang-auto">import argparse

import torch as th
import torch.distributed as dist

if __name__ == "__main__":
	argparse = argparse.ArgumentParser()
	argparse.add_argument("rank", type=int)
	args = argparse.parse_args()

	dist.init_process_group(
		backend="nccl", 
		init_method="tcp://localhost:8001",
		rank=args.rank,
		# Testing on 2 GPUs
		world_size=2
	)

	rank = dist.get_rank()
	assert rank == args.rank

	device = th.device(f"cuda:{rank}")

	# For NCCL, the tensors must be on different GPUs
	tensor_len = 3
	vals = [(rank * tensor_len) + i for i in range(tensor_len)]
	tensor = th.IntTensor(vals).to(device=device)
	# tensor = th.full((3,), rank).to(device=device)
	print(f"Rank: {rank}")
	print(f"Start tensor: {tensor}")
	output = list(th.empty_like(tensor) for _ in range(dist.get_world_size()))
	# print(f"All-gather list: {output}")
	dist.all_gather(output, tensor)

	print(f"Gathered: {output}")
	
	# catted = th.cat(output)
	scatter_size = 5
	aggregate = th.IntTensor(list(range(dist.get_world_size() * scatter_size))).to(device=device)
	input_list = list(aggregate.chunk(dist.get_world_size()))
	print(f"Input list: {input_list}")
	rscatter_input = th.empty_like(input_list[dist.get_rank()])
	dist.reduce_scatter(rscatter_input, input_list)
	print(f"Reduce Scatter output: {rscatter_input}")

	aggregate = th.IntTensor([dist.get_rank() * scatter_size + i for i in range(dist.get_world_size() * scatter_size)]).to(device=device)
	input_list = list(aggregate.chunk(dist.get_world_size()))
	print(f"Input list 2: {input_list}")
	rscatter_input = th.empty_like(input_list[dist.get_rank()])
	dist.reduce_scatter(rscatter_input, input_list)
	print(f"Reduce Scatter output 2: {rscatter_input}")
</code></pre>