allgather
功能说明
集合通信算子AllGather的操作接口，将通信域内所有节点的输入按照rank id重新排序，然后拼接起来，再将结果发送到所有节点的输出。
针对AllGather操作，每个节点都接收按照rank id重新排序后的数据集合，即每个节点的AllGather输出都是一样的。
函数原型
def allgather(tensor, rank_size, group = "hccl_world_group", fusion=0, fusion_id=-1)
参数说明
参数名 |
输入/输出 |
描述 |
|---|---|---|
tensor |
输入 |
TensorFlow的tensor类型。 针对Atlas 训练系列产品，支持的数据类型为：int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64。 针对Atlas 300I Duo 推理卡，支持数据类型为：int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64。 针对Atlas A2 训练系列产品，支持的数据类型为：int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64、bfp16。 |
rank_size |
输入 |
int类型。 group内device的数量。 最大值为32768。 |
group |
输入 |
String类型，最大长度为128字节，含结束符。 group名称，可以为用户自定义group或者"hccl_world_group"。 |
fusion |
输入 |
int类型。 allgather算子融合标识，支持以下取值：
默认值为“0”。 |
fusion_id |
输入 |
allgather算子的融合id。 对相同fusion_id的allgather算子进行融合。 |
返回值
tensor：对输入tensor执行完allgather操作之后的结果tensor。
约束说明
调用该接口的rank必须在当前接口入参group定义的范围内，不在此范围内的rank调用该接口会失败。
支持的型号
Atlas 训练系列产品
Atlas 300I Duo 推理卡
Atlas A2 训练系列产品
调用示例
1 2 3 |
from npu_bridge.npu_init import * rank_size = 2 result = hccl_ops.allgather(tensor, rank_size) |
