昇腾TF适配器ReduceScatter算子
reduce_scatter【免费下载链接】tensorflowAscend TensorFlow Adapter项目地址: https://gitcode.com/cann/tensorflow功能说明集合通信算子ReduceScatter的操作接口将通信域内所有rank的输入数据均分成rank size份然后分别取每个rank的rank size之一份数据进行归约操作如sum、prod、max、min。最后将结果按照编号分散到各个rank的输出buffer。函数原型def reduce_scatter(tensor, reduction, rank_size, grouphccl_world_group, fusion0, fusion_id-1)参数说明参数名输入/输出描述tensor输入TensorFlow的tensor类型。针对Ascend 950PR/Ascend 950DT支持数据类型int8、int16、int32、int64、uint64、float16、float32、float64、bfp16。针对int64、uint64、float64当前仅支持节点内通信。针对Atlas A3 训练系列产品/Atlas A3 推理系列产品支持数据类型int8、int16、int32、int64、float16、float32、bfp16。针对Atlas A2 训练系列产品/Atlas A2 推理系列产品支持数据类型int8、int16、int32、int64、float16、float32、bfp16。需要注意针对int64数据类型性能会有一定的劣化。针对Atlas 训练系列产品支持数据类型int8、int32、int64、float16、float32。针对Atlas 300I Duo 推理卡支持数据类型int8、int16、int32、float16、float32。需要注意tensor的第一个维度的元素个数必须是rank size的整数倍。reduction输入归约操作类型String类型。针对Ascend 950PR/Ascend 950DT支持的操作类型为sum、max、min。针对Atlas A3 训练系列产品/Atlas A3 推理系列产品支持的操作类型为sum、max、min、prod当前版本“prod”操作不支持int16、bfp16数据类型。针对Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的操作类型为sum、max、min、prod当前版本“prod”操作不支持int16、bfp16数据类型。针对Atlas 300I Duo 推理卡支持的操作类型为sum、max、min、prod当前版本“max”、“min”、“prod”操作不支持int16数据类型。rank_size输入group内device的数量int类型。最大值32768。group输入String类型最大长度为128字节含结束符。group名称可以为用户自定义group或者hccl_world_group。fusion输入reducescatter算子融合标识int类型支持以下取值- 0网络编译时不会对该算子进行融合即该reducescatter算子不和其他reducescatter算子融合。- 2网络编译时会对reducescatter算子按照相同的fusion_id进行融合即“fusion_id”相同的reducescatter算子之间会进行融合。fusion_id输入reducescatter算子的融合idint类型。当“fusion”取值为“2”时网络编译时会对相同fusion_id的reducescatter算子进行融合。返回值对输入tensor执行完reducescatter操作之后的结果tensor。约束说明调用该接口的rank必须在当前接口入参group定义的范围内不在此范围内的rank调用该接口会失败。该接口要求输入tensor的数据量不超过8GB。reducescatter算子融合场景只支持reduction操作类型sum。调用示例from npu_bridge.hccl import hccl_ops tensor tf.random_uniform((2, 3), minval1, maxval10, dtypetf.float32) rank_size 2 result hccl_ops.reduce_scatter(tensor, sum, rank_size)【免费下载链接】tensorflowAscend TensorFlow Adapter项目地址: https://gitcode.com/cann/tensorflow创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考