reduce_scatter 通信

这是一个不常见的通信函数，目前的使用场景笔者仅在 DeepSpeed 的ZeRO-Stage1策略上看到调用过这个Op，由于好奇ZeRO算法，因此也顺便调研了这个算子。

算子原型

torch.distributed.reduce_scatter(output, input_list, op=<ReduceOp.SUM>, group=None, async_op=False)
- output(tensor): output tensor
- input_list(list[tensor]): list of tensors to reduce and scatter
- group(ProcessGroup, optional): the process group to work on, If None, the default process will be used. async_op
- async_op(bool, optional)

函数功能：

input_list做all_reduce通信，将all_reduce得到的临时输出scatter到各个output中。
即reduce_scatter = all_reduce + scatter

example

以如下代码为例，各进程中的input_list是[0,1,2,3]，
all_reduce完之后是[0,4,8,12]
scatter到各个进程中的output分别是[0,4,8,12]

import os
import torch
import torch.distributed as dist
from torch.multiprocessing import Process

os.environ['MASTER_ADDR']