aclnnMoeInitRoutingV2Grad

支持的产品型号

• Atlas A2训练系列产品/Atlas 800I A2推理产品

接口原型

每个算子分为，必须先调用 “aclnnMoeInitRoutingV2GradGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnMoeInitRoutingV2Grad”接口执行计算。

• aclnnStatus aclnnMoeInitRoutingV2GradGetWorkspaceSize(const aclTensor *gradExpandedX, const aclTensor *expandedRowIdx, int32_t k, int32_t dropPadMode, int32_t activeNum, aclTensor *gradXOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnMoeInitRoutingV2Grad(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述

• 算子功能：的反向传播，完成tokens的加权求和。
• 计算公式：$$gradX_i=\sum_{t=0}^{k}gradExpandedX[expandedRowIdx[i * k + t]]$$

aclnnMoeInitRoutingV2GradGetWorkspaceSize

• 参数说明：

  • gradExpandedX（aclTensor*，计算输入）：表示Routing过后的目标张量，要求为一个2D/3D的Tensor，2D shape为[NUM_ROWS, H]，3D shape为[E, C, H]，数据类型支持FLOAT16、BFLOAT16、FLOAT32，要求为ND。
  • expandedRowIdx（aclTensor*，计算输入）：表示token按照专家序排序索引，一维Tensor，shape为[NUM_ROWS*K]；数据类型支持INT32，要求为ND。
  • k（int32_t，计算输入）：topk值，Host侧的整型，必须大于0，且能被expandedRowIdx的0轴大小整除。
  • dropPadMode（int32_t，计算输入）：表示场景是否为Drop类，Host侧整型，取值范围为[0, 1]，0表示Dropless场景，1表示Drop/Pad场景。
  • activeNum（int32_t，计算输入）：表示场景是否为Active场景，Host侧整型，值范围大于等于0，0表示非Active场景，大于0表示Active场景，Active场景下gradExpandedX的0轴大小必须等于activeNum值。
  • gradXOut（aclTensor*，计算输出）：表示Routing反向输出，2D的Tensor，shape为[NUM_ROWS, H]；数据类型支持FLOAT16、BFLOAT16、FLOAT32，输出类型与输入gradExpandedX一致，要求为ND。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。 shape符号说明： B: batch size; S: tokens数量; H: hidden size, 即每个token序列长度; K: 即topk, token被处理的专家数 E: expert num, 即专家数; C: expert capacity, 表示专家处理token数量的能力阈值

• 返回值

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

  [object Object]

aclnnMoeInitRoutingV2Grad

• 参数说明：

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

• 返回值：

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

约束与限制

无。

调用示例

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