megengine.functional.distributed.reduce_scatter_sum

reduce_scatter_sum(inp, group=WORLD, device=None, axis=0)[源代码]

通过求和规约指定组中的张量,并在第一维度将其拆分。

参数
  • inp (Tensor) – 输入张量

  • group (Optional[Group]) – 需要处理的组,默认为包含所有进程的 WORLD 组。你可以使用进程序号来创建新的组并使用,例如 [1,3,5] 。

  • device (Optional[str]) – 执行此操作的设备。默认为输入张量所在的设备。可以通过指定设备为 ”gpu0:1“ 以在不同的 cuda 流上执行此操作,其中1是 cuda 流的编号,默认 cuda 流编号为0。

  • axis – 集合通信结果的分割轴默认为0,数据将在0轴上分割

返回类型

Tensor

返回

分割张量

实际案例

input = Tensor([0 1])
# Rank 0 # input: Tensor([0 1])
# Rank 1 # input: Tensor([0 1])
output = reduce_scatter_sum(input)
# Rank 0 # output: Tensor([0])
# Rank 1 # output: Tensor([2])

input = Tensor([0 1])
group = Group([1, 0])
output = reduce_scatter_sum(input, group)
# Rank 0 # output: Tensor([2])
# Rank 1 # output: Tensor([0])