npu.distribute.all_reduce
功能说明
用于NPU分布式部署场景下，worker间的聚合运算。
函数原型
npu.distribute.all_reduce(values, reduction="mean", fusion=1, fusion_id=-1, group="hccl_world_group")
参数说明
参数名 |
输入/输出 |
描述 |
|---|---|---|
values |
输入 |
TensorFlow的tensor类型。 针对Atlas 训练系列产品，tensor支持的数据类型为int8、int32、float16、float32。 针对Atlas 推理系列产品，tensor支持的数据类型为int8、int16（仅sum支持）、int32、float16、float32。 针对Atlas A2 训练系列产品，tensor支持的数据类型为int8、int32、float16、float32、bfloat16（prod操作不支持）。 |
reduction |
输入 |
String类型。 聚合运算的类型，可以为"mean"、"max"、"min"、"prod"或"sum"。 默认为"mean"。 |
fusion |
输入 |
int类型。 allreduce算子融合标识。
|
fusion_id |
输入 |
int类型。 allreduce算子的融合id。 对相同fusion_id的allreduce算子进行融合。 |
group |
输入 |
String类型，最大长度为128字节，含结束符。 group名称，可以为用户自定义group或者"hccl_world_group"。 |
返回值
对values进行聚合运算后的结果，类型与values一致，值与values输入一一对应。
调用示例
在多卡上聚合某个值：
1 2 3 4 5 |
# 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 |
