aclnnBatchMatMulReduceScatterAlltoAll

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	×
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能：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)$$

函数原型

每个算子分为undefined，必须先调用“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)

aclnnBatchMatMulReduceScatterAlltoAllGetWorkspaceSize

• 参数说明：

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

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

  [object Object]

aclnnBatchMatMulReduceScatterAlltoAll

• 参数说明：

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

• 返回值：

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

约束说明

因为集合通信及BatchMatMul计算所需，输入输出shape需满足以下数学关系：（其中ep=epWorldSize，tp=tpWorldSize）

• 按H轴进行ReduceScatter场景，即yShardType为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场景，即yShardType为1场景：

  • x: (E/ep, ep*tp*C/tp, M/tp)
  • weight：(E/ep, M/tp, H)
  • biasOptional：非空指针情况下，三维时为(E/ep, 1, H)，两维时为(E/ep, H)
  • y：(E, C/tp, H)

• 数据关系说明：

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

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考undefined。