CANN / tensorflow AllReduce操作

张开发
2026/5/9 12:26:41 15 分钟阅读

分享文章

CANN / tensorflow AllReduce操作
allreduce【免费下载链接】tensorflowAscend TensorFlow Adapter项目地址: https://gitcode.com/cann/tensorflow功能说明集合通信算子AllReduce的操作接口将group内所有节点的输入数据进行归约操作后再把结果发送到所有节点的输出buf其中归约操作类型由reduction参数指定。函数原型def allreduce(tensor, reduction, fusion1, fusion_id-1, grouphccl_world_group)参数说明参数名输入/输出描述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。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数据类型。fusion输入allreduce算子融合标识int类型支持以下取值- 0网络编译时不会对该算子进行融合即该allreduce算子不和其他allreduce算子融合。- 1网络编译时对该算子按照梯度切分策略进行融合。- 2网络编译时对allreduce算子按照相同的fusion_id进行融合即“fusion_id”相同的allreduce算子之间会进行融合。fusion_id输入allreduce算子的融合idint类型。当“fusion”取值为“2”时网络编译时会对相同fusion_id的allreduce的算子进行融合。group输入String类型最大长度为128字节含结束符。group名称可以为用户自定义group或者hccl_world_group。返回值对输入tensor执行完allreduce操作之后的结果tensor。约束说明调用该接口的rank必须在当前接口入参group定义的范围内不在此范围内的rank调用该接口会失败。每个rank只能有一个输入。allreduce上游节点暂不支持variable算子。该接口要求输入tensor的数据量不超过8GB。allreduce算子融合场景只支持reduction操作类型sum。调用示例from npu_bridge.hccl import hccl_ops tensor tf.random_uniform((1, 3), minval1, maxval10, dtypetf.float32) result hccl_ops.allreduce(tensor, sum)【免费下载链接】tensorflowAscend TensorFlow Adapter项目地址: https://gitcode.com/cann/tensorflow创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

更多文章