allgather
产品支持情况
产品  | 
是否支持  | 
|---|---|
√  | 
|
√  | 
|
x  | 
|
x  | 
|
√  | 
|
√  | 
 针对
功能说明
集合通信算子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 300I Duo 推理卡,支持数据类型为:int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64。  | 
rank_size  | 
输入  | 
int类型。 group内device的数量。 最大值为32768。  | 
group  | 
输入  | 
String类型,最大长度为128字节,含结束符。 group名称,可以为用户自定义group或者"hccl_world_group"。  | 
fusion  | 
输入  | 
int类型。 AllGather算子融合标识,支持以下取值: 
  | 
fusion_id  | 
输入  | 
int类型。 AllGather算子的融合id。 当“fusion”取值为“2”时,网络编译时对相同fusion_id的AllGather算子进行融合。  | 
返回值
tensor:对输入tensor执行完allgather操作之后的结果tensor。
约束说明
- 调用该接口的rank必须在当前接口入参group定义的范围内,不在此范围内的rank调用该接口会失败。
 
调用示例
以下仅为代码片段,不可执行,调用HCCL Python接口进行集合通信操作的完整示例可参见样例代码。
1 2 3 4  | from npu_bridge.hccl import hccl_ops tensor = tf.random_uniform((1, 3), minval=1, maxval=10, dtype=tf.float32) rank_size = 2 result = hccl_ops.allgather(tensor, rank_size)  |