aclnnLightningIndexerGrad

产品支持情况

[object Object]undefined

功能说明

• 算子功能：训练场景下，实现LightningIndexer反向，其中输入有Query, Key, Weights, Dy, Indices，反向主要利用正向计算的Indices从Key中提取TopK序列从而降低Matmul计算量。

• 计算公式： LightningIndexer反向计算公式如下：

  $$S = Relu(Matmul(Query, Gather(Key, Indices)))$$ $$Y = Dy*Weights$$ $$dW = Reduce(S * dy)$$ $$dQ = Matmul(ReluGrad(Y, S), Gather(Key, Indices))$$ $$dK = ScatterAdd(Matmul(ReluGrad(Y, S), Q), Indices)$$

函数原型

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

aclnnLightningIndexerGradGetWorkspaceSize

• 参数说明：

[object Object]

• 返回值：

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

[object Object]

aclnnLightningIndexerGrad

• 参数说明：

  [object Object]

• 返回值：

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

约束说明

• 确定性计算：

  • aclnnLightningIndexerGrad默认非确定性实现，不支持通过aclrtCtxSetSysParamOpt开启确定性。

• 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。

• inputLayout支持TND/BSND。

• 关于数据shape的约束，以Layout的BSND举例。其中：

  • B（Batchsize）：取值范围为1~1024。
  • N（Head-Num）：取值为64。
  • G（Group）：取值为64。
  • S1（Seq-LengthQ）：取值范围为1~128K。
  • S2（Seq-LengthK）：取值范围为topK~128K。
  • D（Head-Dim）：取值为128。
  • TopK：取值为2048

调用示例

通过aclnn单算子调用示例代码如下，仅供参考，具体编译和执行过程请参考。