aclnnAlltoAllAllGatherBatchMatMul

产品支持情况

产品	是否支持
[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]	×

功能说明

• 算子功能：完成AllToAll、AllGather集合通信与BatchMatMul计算融合、并行。
• 计算公式： 计算逻辑如下，其中y1、y2、 y3为输出$$x1 = AllToAll(x)$$ $$y2 = AllGather(x1)$$ $$y3 = BatchMatMul(y2, weight, bias)$$ $$y1 = 激活函数(y3)$$

函数原型

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

aclnnAlltoAllAllGatherBatchMatMulGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：通信后结果作为BatchMatMul计算的左矩阵。数据类型支持FLOAT16、BFLOAT16。该输入进行AllToAll、AllGather集合通信，必须为3维，undefined支持ND。
  • weight（aclTensor*，计算输入）：BatchMatMul计算的右矩阵。数据类型支持FLOAT16、BFLOAT16，类型与x保持一致，必须为3维，undefined支持ND。
  • biasOptional（aclTensor*，计算输入）：BatchMatMul计算的bias，数据类型支持FLOAT16、FLOAT32。x为FLOAT16时，biasOptional需为FLOAT16；x为BFLOAT16时，biasOptional需为FLOAT32，支持两维或三维，undefined支持ND。支持传入空指针。
  • groupEp（char*,计算输入）：ep通信域名称，专家并行的通信域。字符串长度需大于0，小于128。
  • groupTp（char*,计算输入）：tp通信域名称，Tensor并行的通信域。字符串长度需大于0，小于128。
  • epWorldSize（int64_t，计算输入）：ep通信域size，支持2/4/8/16/32。
  • tpWorldSize（int64_t，计算输入）：tp通信域size，支持2/4/8/16/32。
  • xShardType（int64_t，计算输入）：0表示在H维度（即x的第2维，x为3维，分别为第0维、第1维、第2维）按tp域进行allgather，1表示在C维度（即x的第1维）上按tp域进行allgather。
  • 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保持一致。undefined支持：ND。最终计算结果，如果有激活函数则为激活函数的输出，否则为BatchMatMul的输出。
  • y2OutOptional（aclTensor*，计算输出）：可选输出，数据类型FLOAT16、BFLOAT16，支持3维，数据类型与输入x保持一致。undefined支持ND。AllGather的输出，反向可能需要。空指针表示不需要该输出。
  • y3OutOptional（aclTensor*，计算输出）：可选输出，数据类型FLOAT16、BFLOAT16，支持3维，数据类型与输入x保持一致。undefined支持ND。有激活函数时，BatchMatMul的输出，空指针表示不需要该输出。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnAlltoAllAllGatherBatchMatMul

• 参数说明：

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

• 返回值：

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

约束说明

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

按H轴进行AllGather场景，即xShardType为0场景：

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

按C轴进行AllGather场景，即xShardType为1场景：

• 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]，当xShardType为0时，H是tp的整数倍。
• M/tp的取值范围为：[1, 65535]。
• E/ep的取值范围为：[1, 32]。
• ep、tp均仅支持2、4、8、16、32。
• groupEp和groupTp名称不能相同。
• C必须大于0，上限为算子device内存上限，当xShardType为1时，C是tp的整数倍。
• 通算融合算子不支持并发调用，不同的通算融合算子也不支持并发调用。
• 不支持跨超节点，只支持超节点内。

调用示例

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