aclnnBatchMatMulReduceScatterAlltoAll

支持的产品型号

当前版本不支持该接口。

接口原型

每个算子分为两段式接口，必须先调用“aclnnBatchMatMulReduceScatterAlltoAllGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnBatchMatMulReduceScatterAlltoAll”接口执行计算。

• aclnnStatus aclnnBatchMatMulReduceScatterAlltoAllGetWorkspaceSize(const aclTensor* x, const aclTensor* weight, const aclTensor* biasOptional, const char* groupEp, const char* groupTp, int64_t epWorldSize, int64_t tpWorldSize, int64_t yShardType, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)
• aclnnStatus aclnnBatchMatMulReduceScatterAlltoAll(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

功能描述

• 算子功能：BatchMatMulReduceScatterAllToAll是通算融合算子，实现BatchMatMul计算与ReduceScatter、AllToAll集合通信并行的算子。

• 计算公式：大体计算流程为：BatchMatMul计算-->转置（yShardType等于0时需要）-->ReduceScatter集合通信-->Add-->AllToAll集合通信。

• 计算逻辑如下，其中y为输出：

$$temp1 = BatchMatMul(x，weight)$$ $$temp2 = ReduceScatter(temp1)$$ $$temp3 = Add(temp2, bias)$$ $$y = AllToAll(temp3)$$

aclnnBatchMatMulReduceScatterAlltoAllGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：BatchMatMul计算的左矩阵。数据类型支持float16，bfloat16，必须为3维。数据格式支持：ND。
  • weight（aclTensor*，计算输入）：BatchMatMul计算的右矩阵。数据类型支持：float16, bfloat16，必须为3维，类型与x保持一致。数据格式支持：ND。
  • biasOptional（aclTensor*，可选输入）：BatchMatMul计算的bias，可以为空。数据类型支持：float16, float32。x为float16时，bias需为float16；x为bfloat16时，bias需为float32。支持两维或三维。数据格式支持：ND。(由于要进行ReduceScatter通信，因此需要在通信之后再Add)。
  • groupEp（const char*，计算输入）：专家并行的通信域名。字符串长度需大于0，小于128。
  • groupTp（const char*，计算输入）：Tensor并行的通信域名。字符串长度需大于0，小于128。
  • epWorldSize（int64_t，计算输入）：ep通信域size，支持2/4/8/16。
  • tpWorldSize（int64_t，计算输入）：tp通信域size，支持2/4/8/16。
  • yShardType（int64_t，可选输入）：默认值为0。0表示在H维度（即BatchMatMul计算结果的第2维，计算结果共3维，分别为第0维、第1维、第2维）按tp进行ReduceScatter，1表示在C维度（即BatchMatMul计算结果的第1维）按tp进行ReduceScatter。当前仅支持yShardType等于1的场景。
  • out（aclTensor*，计算输出）：Device侧的aclTensor，batch_matmul计算+reduce_scatter计算+all_to_all通信的结果。数据类型支持：float16, bfloat16，必须为3维。类型与输入x保持一致。数据格式支持：ND。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

  返回aclnnStatus状态码，具体参见undefined。

aclnnBatchMatMulReduceScatterAlltoAll

• 参数说明：

  • workspace（void*，入参）：在Device侧申请的workspace内存地址。
  • workspaceSize（uint64_t，入参）：在Device侧申请的workspace大小，由第一段接口aclnnBatchMatMulReduceScatterAlltoAllGetWorkspaceSize获取。
  • executor（aclOpExecutor*，入参）：op执行器，包含了算子计算流程。
  • stream（aclrtStream，入参）：指定执行任务的AscendCL stream流。

• 返回值：

  返回aclnnStatus状态码，具体参见undefined。

约束与限制

因为集合通信及BatchMatMul计算所需，输入输出shape需满足以下数学关系：（其中ep=epWorldSize，tp=tpWorldSize） 按H轴进行ReduceScatter场景，即shard_type为0场景（暂不支持该场景）：

• x: (E/ep, ep*C, M/tp)
• weight：(E/ep, M/tp, H)
• biasOptional：(E/ep, 1, H/tp) 两维时为(E/ep, H/tp)
• y：(E, C, H/tp)

按C轴进行ReduceScatter场景，即shard_type为1场景：

• x: (E/ep, eptpC/tp, M/tp)
• weight：(E/ep, M/tp, H)
• biasOptional：(E/ep, 1, H) 两维时为(E/ep, H)
• y：(E, C, H)

数据关系说明：

• 比如x.size(0)等于E/tp，y.size(0)等于E，则表示，y.size(0) = ep*x.size(0)，y.size(0)是ep的整数倍；其他关系类似。
• E的取值范围为[2, 2048]，且E是ep的整数倍。
• H的取值范围为：[1, 65535]。
• M/tp的取值为：[1, 65535]。
• ep、tp均仅支持2、4、8、16。
• C大于0，上限为算子device内存上限。
• 通算融合算子不支持并发调用，不同的通算融合算子也不支持并发调用。

调用示例

示例代码如下，仅供参考：