npu.distribute.broadcast

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

Synchronizes variables across workers in distributed NPU training scenarios. This is the operation interface of the Broadcast collective communication operator.

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.broadcast(values, root_rank, fusion=2, fusion_id=0, group="hccl_world_group")

Parameters

Parameter	Input/Output	Description
values	Input	A TensorFlow variable or variable set. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are int8, int32, float16, float32, int64, uint64, and bfloat16. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are int8, int32, float16, float32, int64, uint64, and bfloat16. For the Atlas training product, the supported data types are int8, int32, float16, float32, int64, and uint64. For the Atlas inference product, the supported data types are int8, int16, int32, float16, float32, int64, and uint64.
root_rank	Input	An int. Rank ID of the root rank. Must be a rank ID in the group.
fusion	Input	An int. Broadcast operator fusion flag. The value can be one of the following: 0: The Broadcast operator is not fused with other Broadcast operators during network compilation. 2: Broadcast operators with the same fusion_id are fused during network compilation.
fusion_id	Input	An int. Broadcast operator fusion ID. If fusion is set to 2, Broadcast 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

None

Example

To broadcast the variables on device 0 to the rest devices:

# rank_id = 0  rank_size = 8
import npu_device as npu
x = tf.Variable(tf.random.normal(shape=()))
print("before broadcast", x)
npu.distribute.broadcast(x, root_rank=0)
print("after_broadcast", x)

Before the broadcast

After the broadcast

Parent topic: TF Adapter 2.x API Reference