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 的默认流，不指定流则表示使用默认流。
- 返回类型
- 返回
对组中的的值进行求和操作所得到的张量。输出张量的形状必须与
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