npu_allreduce

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")

Option	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 segmentation policy. 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