npu.distribute.all_reduce

Applicability

Product	Supported
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	☓
Atlas inference product	☓
Atlas training product	√

Description

Performs reduction-based aggregation across workers in distributed NPU training scenarios. This is the AllReduce collective communication operator interface.

This API can be used together with the Python communicator management APIs provided by HCCL. For details about the HCCL Python APIs, see HCCL Communicator Management APIs (Python).

Prototype

npu.distribute.all_reduce(values, reduction="mean", fusion=1, fusion_id=-1, group="hccl_world_group")

Parameters

Parameter	Input/Output	Description
values	Input	TensorFlow tensor type. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are int8, int32, float16, float32, and bfloat16 (not supported by prod). For the Atlas A2 training product/Atlas A2 inference product, the supported data types are int8, int32, float16, float32, and bfloat16 (not supported by prod). For the Atlas training product, the supported data types are int8, int32, float16, and float32. For the Atlas inference product, the supported data types are int8, int16 (only supported by sum), int32, float16, and float32.
reduction	Input	A string. Aggregation operation type. The value can be mean, max, min, prod, or sum.
fusion	Input	An int. AllReduce operator fusion flag. The value can be one of the following: 0: The AllReduce operator is not fused with other AllReduce operators during network compilation. 1: The AllReduce operator is fused based on the gradient splitting policy during network compilation. 2: AllReduce operators with the same fusion_id are fused during network compilation.
fusion_id	Input	An int. AllReduce operator fusion ID. If fusion is set to 2, AllReduce operators with the same fusion_id are fused during network compilation.
group	Input	A string of up to 128 bytes, including the end character. Group name, which can be a user-defined value or hccl_world_group.

Returns

Result tensor, whose values are consistent with those in the values input with ordering preserved. It has the same type as values.

Example

To aggregate a value on multiple devices:

# rank_id = 0  rank_size = 8
import npu_device as npu
v = tf.constant(1.0)
x = npu.distribute.all_reduce([v], 'sum') # 8.0
y = npu.distribute.all_reduce([v], 'mean') # 1.0

Parent topic: TF Adapter 2.x API Reference