这是一个不常见的通信函数,目前的使用场景笔者仅在 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']