def reduce_scatter(tensor, reduction, rank_size, group = "hccl_world_group", fusion=0, fusion_id=-1)
提供group内的集合通信reducescatter功能,reduce操作由reduction参数指定。
参数名 |
输入/输出 |
描述 |
---|---|---|
tensor |
输入 |
TensorFlow的tensor类型。 针对Atlas 训练系列产品,tensor支持的数据类型为int8, int32, int64, float16, float32。 针对Atlas A2 训练系列产品,tensor支持的数据类型为int8, int16, int32, int64, float16, float32, bfp16。 tensor的第一个维度的元素个数必需是rank size的整数倍。 |
reduction |
输入 |
String类型。 reduce的op类型,可以为”max”,”min”,”prod”和”sum”。
说明:
针对Atlas A2 训练系列产品,当前版本“prod”操作不支持int16、bfp16数据类型。 |
rank_size |
输入 |
int类型。 group内device的数量。 最大值:32768。 |
group |
输入 |
String类型,最大长度为128字节,含结束符。 group名称,可以为用户自定义group或者"hccl_world_group"。 |
fusion |
输入 |
int类型。 reducescatter算子融合标识,支持以下取值:
默认值为“0”。 |
fusion_id |
输入 |
reducescatter算子的融合id。 对相同fusion_id的reducescatter算子进行融合。 |
tensor:对输入tensor执行完reducescatter操作之后的结果tensor。建议返回tensor的数据量为32Byte对齐,否则会导致性能下降。
1 2 3 |
from npu_bridge.npu_init import * rank_size = 2 result = hccl_ops. reduce_scatter(tensor, "sum", rank_size) |