megengine.functional.distributed.all_reduce_sum

all_reduce_sum(inp, group=WORLD, device=None)[源代码]

在指定组中以求和操作来对张量进行规约操作。

注解

inp 张量在整个组的所有进程中必须形状相同。

参数

inp (Tensor) – 需要规约的张量

关键字参数
  • group (Group or sequence of ints) – 需要处理的组。默认值:“WORLD”. WORLD 组包含所有可用的进程。可以以进程序号构建一个列表以创建一个新的组。

  • device (Tensor.device) – 执行此操作的设备。默认值: None. None 表示以 inp 的设备作为执行设备。特别地, GPU 设备可以通过在设备名后的冒号后面添加一个数字来分配要执行的 cuda 流,而 :0 表示 GPU 的默认流,不指定流则表示使用默认流。

返回类型

Tensor

返回

对组中的的值进行求和操作所得到的张量。输出张量的形状必须与 inp 相同,并且输出张量在整个组的所有进程中都是完全相同的。

实际案例

>>> # We execute all_reduce_sum on rank 0 and rank 1
>>> input = F.arange(2) + 1 + 2 * rank 
>>> input 
Tensor([1. 2.], device=xpux:0) # Rank 0
Tensor([3. 4.], device=xpux:0) # Rank 1
>>> F.distributed.all_reduce_sum(input, group=[0, 1]) 
Tensor([4. 6.], device=xpux:0) # Rank 0
Tensor([4. 6.], device=xpux:0) # Rank 1
>>> # We execute all_reduce_sum with on gpu0 with cuda stream 1
>>> megengine.set_default_device("gpu0") 
>>> input = F.arange(2) + 1 + 2 * rank 
>>> input  
Tensor([1. 2.], device=gpu0:0) # Rank 0
Tensor([3. 4.], device=gpu0:0) # Rank 1
>>> F.distributed.all_reduce_sum(input, device="gpu0:1") 
Tensor([4. 6.], device=gpu0:0) # Rank 0
Tensor([4. 6.], device=gpu0:0) # Rank 1