CANN/ops-transformer MoE门控Top-K反向算子

MoeGatingTopKBackward

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

产品支持情况

功能说明

• 算子功能：完成MoE（Mixture of Experts）门控Top-K选择的反向梯度计算。该算子是MoeGatingTopK的反向算子，根据前向算子输出的归一化得分（xNorm）、上游梯度（gradY）和专家索引（expertIdx），计算输入得分矩阵的梯度（gradX）。支持sigmoid模式（normType=1）。

• 计算公式（sigmoid模式，normType=1）：

  1. 缩放梯度

  $$ gradYScaled_{ip} = routedScalingFactor \cdot gradY_{ip} $$

  2. 正向renorm的反向传播

  $$ wPrime_{ip} = xNorm_{i,\ expertIdx_{ip}} $$

  $$ D_i = \sum_{p} wPrime_{ip} + eps $$

  $$ w_{ip} = \frac{wPrime_{ip}}{D_i} $$

  $$ beta_i = \sum_{p} w_{ip} \cdot gradYScaled_{ip} $$

  $$ gradWPrime_{ip} = \frac{gradYScaled_{ip} - beta_i}{D_i} $$

  3. 散射到完整维度

  $$ gradNormX_{ij} = \sum_{p:\ expertIdx_{ip}=j} gradWPrime_{ip} $$

  4. Sigmoid反向传播

  $$ gradX_{ij} = xNorm_{ij} \cdot (1 - xNorm_{ij}) \cdot gradNormX_{ij} $$

参数说明

约束说明

无

调用说明

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer