def allgather(tensor, rank_size, group = "hccl_world_group", fusion=0, fusion_id=-1)
提供group内的集合通信allgather功能,将所有节点的输入Tensor合并起来。
参数名 |
输入/输出 |
描述 |
---|---|---|
tensor |
输入 |
TensorFlow的tensor类型。 针对Atlas 训练系列产品,支持的数据类型为int8, uint8, 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。
1 2 3 |
from npu_bridge.npu_init import * rank_size = 2 result = hccl_ops.allgather(tensor, rank_size) |
调用该接口的rank必须在当前接口入参group定义的范围内,不在此范围内的rank调用该接口会失败。