MoeTokenUnpermuteWithRoutingMapGrad

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 算子功能：aclnnMoeTokenUnpermuteWithRoutingMap的反向传播。

• 计算公式：

  (1) probs非None：

  $$permutedTokensGrad[outIndex[i]]=unpermutedTokensGrad[permuteTokenId[i]]$$

  $$permutedProbsGrad=permutedTokensGrad\ast permutedTokensOptional$$

  $$probsGradExpertOrder=\sum _{j=0}^{hidden\_size}(permutedProbsGra{d}_{i,j})$$

  • dropAndPad为false时

  $$probsGradOut=masked\_scatter(routingMapOptiona{l}^T,probsGradExpertOrder)$$

  $$permutedProbs=probsOptiona{l}^T.masked\_select(routingMapOptiona{l}^T)$$

  $$permutedTokensGradOut=permutedProbs.unsqueeze(-1)\ast permutedTokensGrad$$

  • dropAndPad为true时

  $$probsGradOut[permuteTokenId[i],outIndex[i]/capacity]=probsGradExpertOrder[outIndex[i]]$$

  $$permutedProbs[outIndex[i]]=probsOptional.view(1)[i]$$

  $$permutedTokensGradOut=permutedProbs\ast permutedTokensGrad$$

  (2) probs为None：

  $$permutedTokensGradOut[outIndex[i]]=unpermutedTokensGrad[permuteTokenId[i]]$$

  1. hidden_size指unpermutedTokensGrad的第1维大小。
  2. dropAndPad等于true时，每个专家固定能够处理capacity个token。输入routingMapOptional的第1维是experts_num，即专家个数，输入outIndex的第0维是experts_num * capacity，根据这两个维度可以算出capacity。
  3. dropAndPad等于false时，每个token能被小于等于topK_num个专家处理。输入unpermutedTokensGrad的第0维是tokens_num，即token的个数，输入outIndex的第0维是tokens_num * topK_num，根据这两个维度可以算出topK_num。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
unpermutedTokensGrad	输入	Device侧的aclTensor。计算公式中的unpermutedTokensGrad，代表正向输出unpermutedTokens的梯度。	BFLOAT16、FLOAT16、FLOAT	ND
outIndex	输入	Device侧的aclTensor。计算公式中outIndex，代表输出位置索引。	INT32	ND
permuteTokenId	输入	Device侧的aclTensor。计算公式中的permuteTokenId，代表输入permutedTokens每个位置对应的Token序号。	INT32	ND
routingMapOptional	可选输入	Device侧的aclTensor，可选输入，当输入probsOptional为空指针时不需要此输入，应该传入空指针。计算公式中的routingMapOptional，代表对应位置的Token是否被对应专家处理。	INT8	ND
permutedTokensOptional	可选输入	Device侧的aclTensor，可选输入，当输入probsOptional为空指针时不需要此输入，应该传入空指针。	BFLOAT16、FLOAT16、FLOAT	ND
probsOptional	可选输入	Device侧的aclTensor，可选输入，当不需要时为空指针。	BFLOAT16、FLOAT16、FLOAT	ND
dropAndPad	属性	host侧的BOOL。true表示开启dropAndPad，false表示关闭dropAndPad。	BOOL	-
restoreShapeOptional	属性	host侧的aclIntArray。	INT64	-
permutedTokensGradOut	输出	输入permutedTokens的梯度，要求是一个2D的Tensor。	BFLOAT16、FLOAT16、FLOAT	ND
probsGradOutOptional	可选输出	当不需要时为空指针。输入probs的梯度。	BFLOAT16、FLOAT16、FLOAT	ND

约束说明

• 当输入probsOptional非空，且dropAndPad为false时
  • 要求topK_num <= 512且topK_num <= experts_num。
  • 要求experts_num满足(ubSize - (probTypeLen + 1) * numExpertAlign-(tokenTypeLen + 8) * 256) / (6 * tokenTypeLen + 12) >= 1，其中ubSize是芯片ub空间大小，probTypeLen是输入probsOptional的数据类型对应的字节数，tokenTypeLen是输入unpermutedTokensGrad的数据类型对应的字节数，numExpertAlign是experts_num对32做向上对齐的结果。
• 当输入probsOptional非空，且dropAndPad为true时
  • 要求capacity <= tokens_num。
  • 要求hidden_size <= 256 * (ubSize - 2080) / (8 + tokenTypeLen)，其中ubSize是芯片ub空间大小，tokenTypeLen是输入unpermutedTokensGrad的数据类型对应的字节数。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_moe_token_unpermute_with_routing_map_grad	通过aclnnMoeTokenUnpermuteWithRoutingMapGrad接口方式调用MoeTokenUnpermuteWithRoutingMapGrad算子。