aclnnFusedLinearOnlineMaxSum

产品支持情况

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

功能说明

• 接口功能：功能等价Megatron的matmul与fused_vocab_parallel_cross_entropy的实现，支持vocabulary_size维度切卡融合matmul与celoss，中间根据通信拆分为undefined、undefined，需要依次调用实现完整功能。
• 计算公式：$$vocabParallelLogitsOutOptional = input @ weight^T$$

$$logitsMaxLocalOut = max(vocabParallelLogitsOutOptional, dim=-1)$$ $$subRes = vocabParallelLogitsOutOptional - logitsMaxLocalOut$$ $$sumExpLogitsLocalOut = sum(exp(subRes), dim=-1)$$ $$targetMask = (target < vocabStartIndex) | (target > vocabEndIndex)$$ $$maskedTargetOut[b] = \begin{cases} 0 & \text{targetMask[b]=true}\\ target[b] - vocabStartIndex & \text{targetMask[b]=false} \end{cases}$$ $$predictedLogitsLocalOut[b] = \begin{cases} 0 & \text{targetMask[b]=true}\\ subRes[b][maskedTargetOut[b]] & \text{targetMask[b]=false} \end{cases}$$ $$alignNum = (input.size(0) + 7) / 8 * 8\\ maskBit[p] = \begin{cases} uint8(targetMask[p]) & \text{p < input.size(0)}\\ 1 & \text{input.size(0) <= p < alignNum} \end{cases} \\ targetMaskOut[k] = 0b(maskBit[8*k:8*k+8])$$

其中$0 \le b \lt input.size(0), 0 \le p \lt alignNum, 0 \le k \lt alignNum / 8$。

函数原型

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

• aclnnStatus aclnnFusedLinearOnlineMaxSumGetWorkspaceSize(const aclTensor* input, const aclTensor* weight, const aclTensor* target, int64_t vocabStartIndex, int64_t vocabEndIndex, aclTensor* logitsMaxLocalOut, aclTensor* sumExpLogitsLocalOut, aclTensor* predictedLogitsLocalOut, aclTensor* targetMaskOut, aclTensor* maskedTargetOut, aclTensor* vocabParallelLogitsOutOptional, uint64_t* workspaceSize, aclOpExecutor** executor)
• aclnnStatus aclnnFusedLinearOnlineMaxSum(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnFusedLinearOnlineMaxSumGetWorkspaceSize

• 参数说明：

  • input(aclTensor*, 计算输入)：表示matmul计算的左矩阵，公式中的input，Device侧的aclTensor。数据类型支持BFLOAT16、FLOAT16，shape支持2维，支持undefined，支持空tensor，undefined支持ND。
  • weight(aclTensor*, 计算输入)：表示matmul计算的右矩阵，公式中的weight，Device侧的aclTensor。数据类型支持BFLOAT16、FLOAT16，数据类型需要与logits一致，shape支持2维且weight.size(1)需要与input.size(1)一致，支持undefined，支持空tensor，undefined支持ND。
  • target(aclTensor*, 计算输入)：表示目标索引，公式中的target，Device侧的aclTensor。数据类型支持INT32、INT64，shape支持1维且target.size(0)需要与input.size(0)一致，支持undefined，支持空tensor，undefined支持ND。
  • vocabStartIndex(int64_t, 计算输入)：表示分到本卡上的开始索引，公式中的vocabStartIndex，Host侧的int64常量。取值范围为[0, max(target) - 1]
  • vocabEndIndex(int64_t, 计算输入)：表示分到本卡上的结束索引，公式中的vocabEndIndex，Host侧的int64常量。取值范围为[vocabStartIndex, min(vocabStartIndex + weight.size(0), max(target) - 1)]
  • logitsMaxLocalOut(aclTensor*, 计算输出)：表示matmul计算后各行的最大值，公式中的logitsMaxLocalOut，Device侧的aclTensor。数据类型支持FLOAT32，shape支持1维且logitsMaxLocalOut.size(0)需要与input.size(0)一致，支持undefined，undefined支持ND。
  • sumExpLogitsLocalOut(aclTensor*, 计算输出)：表示matmul计算结果与其各行的最大值作差后经过exp计算后的结果，公式中的sumExpLogitsLocalOut，Device侧的aclTensor。数据类型支持FLOAT32，shape支持1维且sumExpLogitsLocalOut.size(0)需要与input.size(0)一致，支持undefined，undefined支持ND。
  • predictedLogitsLocalOut(aclTensor*, 计算输出)：表示matmul计算结果与其各行的最大值作差后maskedTargetOut筛选后的结果，公式中的predictedLogitsLocalOut，Device侧的aclTensor。数据类型支持FLOAT32，shape支持1维且predictedLogitsLocalOut.size(0)需要与input.size(0)一致，支持undefined，undefined支持ND。
  • targetMaskOut(aclTensor*, 计算输出)：表示用于筛选词表的mask，公式中的targetMaskOut，Device侧的aclTensor。数据类型支持UINT8，shape支持1维，shape为[(input.size(0) + 7) / 8]，支持undefined，undefined支持ND。
  • maskedTargetOut(aclTensor*, 计算输出)：表示target经过targetMaskOut过滤后的结果，公式中的maskedTargetOut，Device侧的aclTensor。数据类型支持INT32、INT64，且数据类型需要与target一致，shape支持1维且maskedTargetOut.size(0)需要与input.size(0)一致，支持undefined，undefined支持ND。
  • vocabParallelLogitsOutOptional(aclTensor*, 计算输出)：表示matmul计算结果，可选输出，公式中的vocabParallelLogitsOutOptional，Device侧的aclTensor。数据类型支持BFLOAT16、FLOAT16，且数据类型需要input一致，shape支持2维，shape为[input.size(0), weight.size(0)]，支持undefined，undefined支持ND。当vocabParallelLogitsOutOptional为nullptr时为省显存分支，否则为高性能分支，当前仅支持高性能分支。
  • workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnFusedLinearOnlineMaxSum

• 参数说明：

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

• 返回值：

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

约束说明

当前仅支持高性能分支。

调用示例

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