aclnnAlltoAllAllGatherBatchMatMul

支持的产品型号

当前版本不支持该接口。

接口原型

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

• aclnnStatus aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize(const aclTensor* x, const aclTensor* weight, const aclTensor* biasOptional, const char* groupEp, const char* groupTp, int64_t epWorldSize, int64_t tpWorldSize, int64_t xShardType, int64_t actType, aclTensor* y1Out, aclTensor* y2OutOptional, aclTensor* y3OutOptional, uint64_t* workspaceSize, aclOpExecutor** executor)
• aclnnStatus aclnnAlltoAllAllGatherBatchMatMul(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

功能描述

• 算子功能：完成AllToAll、AllGather集合通信与BatchMatMul计算融合、并行。

• 计算公式：

• 计算逻辑如下，其中y1 y2 y3为输出

  $$x1 = AllToAll(x)$$ $$y2 = AllGather(x1)$$ $$y3 = BatchMatMul(y2, weight, bias)$$ $$y1 = 激活函数(y3)$$

aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：通信后结果作为BatchMatMul计算的左矩阵。数据类型支持float16，bfloat16。该输入进行AllToAll、AllGather集合通信，必须为3维，数据格式支持ND。
  • weight（aclTensor*，计算输入）：BatchMatMul计算的右矩阵。数据类型支持float16, bfloat16，类型与x保持一致，必须为3维，数据格式支持ND。
  • biasOptional（aclTensor*，计算输入）：BatchMatMul计算的bias，数据类型支持float16, float32。x为float16时，bias需为float16；x为bfloat16时，bias需为float32，支持两维或三维，数据格式支持ND。支持传入空指针。
  • groupEp（char*，计算输入）：ep通信域名称，专家并行的通信域。字符串长度需大于0，小于128。
  • groupTp（char*，计算输入）：tp通信域名称，Tensor并行的通信域。字符串长度需大于0，小于128。
  • epWorldSize（int64_t，计算输入）：ep通信域size，支持2/4/8/16。
  • tpWorldSize（int64_t，计算输入）：tp通信域size，支持2/4/8/16。
  • xShardType（int64_t，计算输入）：0表示在H维度（即x的第2维，x为3维，分别为第0维、第1维、第2维）按tp域进行allgather，1表示在C维度（即x的第1维）上按tp域进行allgather。当前仅支持xShardType等于1的场景。
  • actType（int64_t，计算输入）：激活函数类型，支持0/1/2/3/4的输入，0表示无激活函数，对应关系为[0：None，1：GELU，2：Silu，3：Relu，4：FastGELU]。
  • y1Out（aclTensor*，计算输出）：数据类型float16, bfloat16，支持3维，数据类型与输入x保持一致。数据格式支持：ND。最终计算结果，如果有激活函数则为激活函数的输出，否则为BatchMatMul的输出。
  • y2OutOptional（aclTensor*，计算输出）：可选输出，数据类型float16, bfloat16，支持3维，数据类型与输入x保持一致。数据格式支持：ND。AllGather的输出，反向可能需要。空指针表示不需要该输出。
  • y3OutOptional（aclTensor*，计算输出）：可选输出，数据类型float16, bfloat16，支持3维，数据类型与输入x保持一致。数据格式支持：ND。有激活函数时，BatchMatMul的输出，空指针表示不需要该输出。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

aclnnAlltoAllAllGatherBatchMatMul

• 参数说明：

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

• 返回值：

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

约束与限制

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

• x: (E,C/tp,H)
• weight：(E/ep,H,M/tp)
• biasOptional：非空指针情况下，三维时为(E/ep,1,M/tp)，两维时为(E/ep,M/tp)
• y1Out：(E/ep,ep*tp*C/tp,M/tp)
• y2OutOptional：(E/ep,ep*tp*C/tp,H)
• y3OutOptional：(E/ep,ep*tp*C/tp,M/tp) 数据关系说明：
• 比如x.size(0)等于E，weight.size(0)等于E/ep，则表示，x.size(0) = ep*weight.size(0)，x.size(0)是ep的整数倍；其他关系类似。
• E的取值范围为[2, 512]，且E是ep的整数倍。
• H的取值范围为：[1, 65535]。
• M/tp的取值范围为：[1, 65535]。
• E/ep的取值范围为：[1, 32]。
• ep、tp均仅支持2、4、8、16。
• groupEp和groupTp名称不能相同。
• C必须大于0，上限为算子device内存上限。
• 通算融合算子不支持并发调用，不同的通算融合算子也不支持并发调用。
• 不支持跨超节点，只支持超节点内，ep域AlltoAll支持超节点内跨节点，tp域AllGather仅支持超节点内单一节点。

调用示例

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