npu_allreduce

Product	Supported (√/x)
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	X
Atlas inference product	X
Atlas training product	√

Performs AllReduce and update operations on gradients after the gradient computing is complete.

def _npu_allreduce(values, reduction="mean", fusion=1, fusion_id=-1, group="hccl_world_group")

Parameter	Input/Output	Description
values	Input	Tensor list or tensor.
reduction	Input	Op type of Reduce. The value can be sum or mean.
fusion	Input	Operator fusion flag, which is of the int type. 0: disabled. The AllReduce operator is not fused with other AllReduce operators. 1 (default): enabled. The AllReduce operator is fused based on the gradient splitting strategy. 2: enabled. AllReduce operators with the same fusion_id are fused.
fusion_id	Input	Operator fusion index flag. AllReduce operators with the same fusion_id will be fused.
group	Input	Group name, which is of the string type. It can be a user-defined value or hccl_world_group.

Tensor list or tensor, which is consistent with the inputs.

from npu_bridge.npu_init import *
grads = npu_allreduce(tf.gradients(a + b, [a, b], stop_gradients=[a, b]))

Parent topic: npu_bridge.estimator.npu.npu_optimizer