ops-transformer 中 moe_finalize_routing_grad 反向算子参数、梯度公式与正反向布局转置全解析【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer在 MoEMixture-of-Experts模型的训练流程中moe_finalize_routing负责将各专家的输出按路由权重缩放、合并并写回原始 token 布局其反向算子moe_finalize_routing_grad则把输出梯度grad_y拆分回专家侧输入expanded_x的梯度与路由权重scales的梯度。本篇以 ops-transformer 仓库中 moe_finalize_routing_grad PyTorch API 文档 为主体完整梳理该算子的功能定义、梯度计算公式、全部参数约束、返回值语义并结合 torch_extension 封装源码、C ACLNN 包装层 与 正向算子的 autograd 绑定实现讲清一个容易踩坑的关键点正反向算子对expanded_row_idx采用了不同的内存布局手动调用反向时必须自行完成布局转置。读完后你将能够正确地在单算子模式与自动反向两种路径下调用该算子并理解每条约束背后的实现依据。产品支持情况与功能定位该算子的产品支持情况如下与文档保持一致产品是否支持Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持moe_finalize_routing_grad是 moe_finalize_routing 的反向算子底层封装aclnnMoeFinalizeRoutingV2Grad两段式接口。其功能是根据正向输出的expanded_row_idx将grad_y正向输出y的梯度按路由权重scales映射回expanded_x的梯度grad_expanded_x和scales的梯度grad_scales。算子的算子级文档见 README 与 aclnnMoeFinalizeRoutingV2Grad 接口文档C 调用样例可参考 test_aclnn_moe_finalize_routing_v2_grad.cpp。梯度计算公式设i为扁平索引范围为[0, R * K)其中R为 token 行数batch * sequenceK为 top-K 专家个数j为 hidden 维索引。公式按scales/bias是否传入分三种情况scales未传入此时 K 必须为 1梯度直接按行索引回写不做缩放grad_expanded_x[expanded_row_idx[i]][j] grad_y[i / K][j]scales传入、bias未传入grad_expanded_x需要乘以路由权重grad_scales由expanded_x与grad_y的内积累加得到grad_expanded_x[expanded_row_idx[i]][j] grad_y[i / K][j] * scales[i / K][i % K] grad_scales[i] Σ_j expanded_x[expanded_row_idx[i]][j] * grad_y[i / K][j]scales和bias均传入grad_expanded_x同上grad_scales中需加上按专家索引查取的 bias 项grad_expanded_x[expanded_row_idx[i]][j] grad_y[i / K][j] * scales[i / K][i % K] grad_scales[i] Σ_j (expanded_x[expanded_row_idx[i]][j] bias[expert_idx[i]][j]) * grad_y[i / K][j]两个重要说明该反向算子仅计算grad_expanded_x和grad_scales不计算x1、x2、bias的梯度自动反向下这些位置的梯度返回None。布局差异expanded_row_idx的排布布局与正向不同。正向moe_finalize_routing使用(K, R)布局即idx[k * R row]本反向算子使用(R, K)布局即idx[row * K k]。通过自动反向调用时框架会自动完成布局转置手动调用时必须自行确保布局正确。函数原型与参数说明cann_ops_transformer.moe_finalize_routing_grad( grad_y: torch.Tensor, expanded_row_idx: torch.Tensor, expanded_x: Optional[torch.Tensor] None, scales: Optional[torch.Tensor] None, expert_idx: Optional[torch.Tensor] None, bias: Optional[torch.Tensor] None, drop_pad_mode: int 0, active_num: int 0, expert_num: int 0, expert_capacity: int 0, ) - Tuple[torch.Tensor, torch.Tensor]输入参数参数名类型可选/必选描述数据类型维度(shape)grad_yTensor必选正向输出y的梯度float16、bfloat16、float32(R, H)expanded_row_idxTensor必选行索引映射采用(R, K)布局注意与正向的(K, R)布局不同。Dropless 场景下值域为[0, R * K - 1]且不重复DropPad 场景下值域为[-1, E * C - 1]除 -1 外不重复-1 表示丢弃。当scales未传入时 K 必须为 1int32(R * K,)expanded_xTensor可选正向输入expanded_x用于计算grad_scales。传入scales时必须同时传入数据类型需与grad_y一致与 grad_y 一致(NUM_ROWS * K, H)或(E, C, H)scalesTensor可选正向输入scales即路由权重系数。传入scales时必须同时传入expanded_xfloat16、bfloat16、float32(R, K)expert_idxTensor可选正向输入expert_idx即专家索引。传入bias时必须同时传入值域为[0, E-1]int32(R, K)biasTensor可选正向输入bias即偏置量与 grad_y 一致(E, H)drop_pad_modeint可选丢弃/填充模式需与正向保持一致。0drop less1drop pad。默认值为 0--active_numint可选Dropless 场景下grad_expanded_x的最大输出行数仅当大于 0 且小于R * K时生效DropPad 场景下不生效。默认值为 0--expert_numint可选专家数。Dropless 场景下不生效DropPad 场景下必须大于 0且当bias传入时必须等于bias的第 0 维大小E。默认值为 0--expert_capacityint可选每个专家能够处理的 tokens 数drop_pad_mode1时必须传入且大于 0。默认值为 0--返回值参数名类型描述数据类型维度(shape)grad_expanded_xTensorexpanded_x的梯度。Dropless 场景下 shape 为(dim0, H)其中dim0为expanded_row_idx.numel()active_num0时或min(active_num, expanded_row_idx.numel())active_num 0时DropPad 场景下 shape 为(expert_num, expert_capacity, H)与 grad_y 一致(dim0, H)或(E, C, H)grad_scalesTensorscales的梯度。shape 为(R, K)其中K为scales的第二维scales传入时或 1scales未传入时。数据类型scales传入时与scales一致否则与grad_y一致float16、bfloat16、float32(R, K)上述输出 shape 的推断逻辑在源码中可以直接对应csrc 包装层 先根据drop_pad_mode预分配gradExpandedXOut与gradScalesOutDropPad 用{expertNum, expertCapacity, hidden}Dropless 用{dim0, hidden}且dim0受activeNum截断再调用ACLNN_CMD(aclnnMoeFinalizeRoutingV2Grad, ...)完成实际计算Python 侧注册的 Meta 实现 也按同样规则推导 shape保证编译期图与运行期一致。约束说明该接口支持训练场景下使用支持单算子模式和图模式。grad_y必须是 2 维张量expanded_row_idx必须是 1 维 int32 张量。drop_pad_mode仅支持取值 0 或 1不支持 2 或 3行排列模式。传入bias时expert_idx必须同时传入传入scales时expanded_x必须同时传入。drop_pad_mode1时expert_num必须大于 0bias传入时必须等于bias的第 0 维大小 Eexpert_capacity必须大于 0。expanded_x、bias的数据类型必须与grad_y一致scales的数据类型在 Atlas A2/A3 系列上必须与grad_y一致在 Ascend 950PR/950DT 上可以与grad_y不一致。该反向算子仅支持常规专家场景不支持 aclnnMoeFinalizeRoutingV4 特有特性正向的x、alpha1、alpha2、v以及有效的zero_expert_range、copy_expert_range、constant_expert_range。当正向使用了这些特性时调用自动反向会抛出NotImplementedError。各参数需与正向调用保持一致。这些约束在 C 包装层源码 中均有对应的TORCH_CHECK前置校验grad_y的 2 维与 float16/float32/bfloat16 类型、expanded_row_idx的 1 维 int32 类型、drop_pad_mode ∈ {0, 1}、bias/expert_idx与scales/expanded_x的成对传入关系、DropPad 模式下expertNum/expertCapacity必须为正均在进入 ACLNN 调用前完成校验报错信息清晰可定位。关于 V4 特性不支持的限制正向算子的 autograd 绑定实现 MoeFinalizeRoutingFn 中有明确实现_check_backward_supported会在ctx.has_v4_extras即提供了x/alpha1/alpha2/v或有效的 expert range或drop_pad_mode不在 (0, 1) 时抛出NotImplementedError其 docstring 也声明了forward - aclnnMoeFinalizeRoutingV4, backward - aclnnMoeFinalizeRoutingV2Grad的组合关系以及 V2Grad 只处理常规专家场景的原因。expanded_row_idx 布局转置反向调用最关键的细节这是该算子与正向最易混淆的一点。正向moe_finalize_routing在drop_pad_mode为 0 或 1 时expanded_row_idx按(K, R)布局排布即idx[k * R row]而反向算子要求(R, K)布局即idx[row * K k]。自动反向下框架在backward中通过_transpose_row_idx完成这一转换见 moe_finalize_routing.pydef _transpose_row_idx(expanded_row_idx, top_k, row_num): Convert expanded_row_idx from forward (K, R) layout to grad op (R, K) layout. if top_k 1: return expanded_row_idx.reshape(top_k, row_num).t().contiguous().reshape(-1) return expanded_row_idx也就是说先把一维索引按(K, R)还原成二维转置为(R, K)后展平正好得到反向算子要求的排布。手动调用反向时需要自己执行同样的 reshape 转置 展平操作。此外backward中还做了两件辅助工作值得注意合成 scales_synthesize_scales在scales为 None 但top_k 1时会构造一个全 1 的(R, K)张量作为有效scales传入从而保证grad_scales形状正确推断 shape 参数_infer_grad_shape根据正向保存的expanded_x形状推断active_num/expert_num/expert_capacityDropPad 模式返回(0, E, C)Dropless 模式返回(expanded_x.shape[0], 0, 0)。自动反向还遵循了文档声明的行为若x1/x2/bias需要梯度会发出警告说明这些位置的梯度不会被计算并返回None建议改用外部残差加法若expanded_x和scales都不需要梯度则直接跳过对反向算子的调用。调用示例单算子模式手动调用反向注意示例中对expanded_row_idx的手动布局转置import torch import torch_npu from cann_ops_transformer.ops import moe_finalize_routing, moe_finalize_routing_grad torch_npu.npu.set_device(0) NUM_ROWS 2 K 2 H 8 E 4 expanded_x torch.randn(NUM_ROWS * K, H, dtypetorch.float32, devicenpu) # 正向expanded_row_idx为(K, R)布局idx[k*R row] forward_row_idx torch.tensor([0, 1, 2, 3], dtypetorch.int32, devicenpu) scales torch.randn(NUM_ROWS, K, dtypetorch.float32, devicenpu) expert_idx torch.tensor([[0, 1], [2, 3]], dtypetorch.int32, devicenpu) bias torch.randn(E, H, dtypetorch.float32, devicenpu) # 正向 y moe_finalize_routing(expanded_x, forward_row_idx, biasbias, scalesscales, expert_idxexpert_idx, kK) # 手动调用反向需将expanded_row_idx从(K, R)布局转置为(R, K)布局 grad_row_idx forward_row_idx.reshape(K, NUM_ROWS).t().contiguous().reshape(-1) grad_y torch.randn_like(y) grad_expanded_x, grad_scales moe_finalize_routing_grad( grad_y, grad_row_idx, expanded_xexpanded_x, scalesscales, expert_idxexpert_idx, biasbias, drop_pad_mode0, ) print(grad_expanded_x.shape) # torch.Size([4, 8]) print(grad_scales.shape) # torch.Size([2, 2])自动反向autograd 自动触发当正向moe_finalize_routing中expanded_x、scales等可微输入的requires_grad为 True且未使用 V4 特有特性x、alpha1、alpha2、v、zero_expert_range等时loss.backward()会自动触发本算子无需手动调用仅在需要显式控制梯度的场景下才保留手动调用路径。import torch import torch_npu from cann_ops_transformer.ops import moe_finalize_routing torch_npu.npu.set_device(0) NUM_ROWS 4 K 2 H 8 expanded_x torch.randn(NUM_ROWS * K, H, dtypetorch.float32, devicenpu, requires_gradTrue) row_idx torch.arange(NUM_ROWS * K, dtypetorch.int32, devicenpu) scales torch.randn(NUM_ROWS, K, dtypetorch.float32, devicenpu, requires_gradTrue) # 正向expanded_x.requires_gradTrue时自动启用autograd y moe_finalize_routing(expanded_x, row_idx, scalesscales, kK) # 反向自动触发moe_finalize_routing_grad loss y.sum() loss.backward() print(expanded_x.grad.shape) # torch.Size([8, 8]) print(scales.grad.shape) # torch.Size([4, 2])确定性计算与配套接口确定性计算本算子默认支持确定性计算与 aclnn 接口文档 中aclnnMoeFinalizeRoutingV2Grad 默认确定性实现的说明一致。配套接口本算子是 moe_finalize_routing 的反向算子二者参数需保持对齐。如需使用原生两段式接口可查阅 aclnnMoeFinalizeRoutingV2Grad 接口文档其第一段aclnnMoeFinalizeRoutingV2GradGetWorkspaceSize负责入参校验并返回 workspace 大小与执行器第二段aclnnMoeFinalizeRoutingV2Grad执行实际计算典型报错包括ACLNN_ERR_PARAM_NULLPTR161001必选输入输出为空指针、ACLNN_ERR_PARAM_INVALID161002数据类型/格式不支持、ACLNN_ERR_INNER_NULLPTR561103shape 或取值不满足要求。最后再强调一次使用要点该反向算子只产出grad_expanded_x与grad_scales两个梯度如果需要对x1/x2求梯度建议使用外部残差加法将其作为正向输入而不是依赖本算子手动调用时务必把expanded_row_idx从正向的(K, R)布局转置为反向的(R, K)布局这是保证梯度正确性的前提。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
ops-transformer 中 moe_finalize_routing_grad 反向算子：参数、梯度公式与正反向布局转置全解析
ops-transformer 中 moe_finalize_routing_grad 反向算子参数、梯度公式与正反向布局转置全解析【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer在 MoEMixture-of-Experts模型的训练流程中moe_finalize_routing负责将各专家的输出按路由权重缩放、合并并写回原始 token 布局其反向算子moe_finalize_routing_grad则把输出梯度grad_y拆分回专家侧输入expanded_x的梯度与路由权重scales的梯度。本篇以 ops-transformer 仓库中 moe_finalize_routing_grad PyTorch API 文档 为主体完整梳理该算子的功能定义、梯度计算公式、全部参数约束、返回值语义并结合 torch_extension 封装源码、C ACLNN 包装层 与 正向算子的 autograd 绑定实现讲清一个容易踩坑的关键点正反向算子对expanded_row_idx采用了不同的内存布局手动调用反向时必须自行完成布局转置。读完后你将能够正确地在单算子模式与自动反向两种路径下调用该算子并理解每条约束背后的实现依据。产品支持情况与功能定位该算子的产品支持情况如下与文档保持一致产品是否支持Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持moe_finalize_routing_grad是 moe_finalize_routing 的反向算子底层封装aclnnMoeFinalizeRoutingV2Grad两段式接口。其功能是根据正向输出的expanded_row_idx将grad_y正向输出y的梯度按路由权重scales映射回expanded_x的梯度grad_expanded_x和scales的梯度grad_scales。算子的算子级文档见 README 与 aclnnMoeFinalizeRoutingV2Grad 接口文档C 调用样例可参考 test_aclnn_moe_finalize_routing_v2_grad.cpp。梯度计算公式设i为扁平索引范围为[0, R * K)其中R为 token 行数batch * sequenceK为 top-K 专家个数j为 hidden 维索引。公式按scales/bias是否传入分三种情况scales未传入此时 K 必须为 1梯度直接按行索引回写不做缩放grad_expanded_x[expanded_row_idx[i]][j] grad_y[i / K][j]scales传入、bias未传入grad_expanded_x需要乘以路由权重grad_scales由expanded_x与grad_y的内积累加得到grad_expanded_x[expanded_row_idx[i]][j] grad_y[i / K][j] * scales[i / K][i % K] grad_scales[i] Σ_j expanded_x[expanded_row_idx[i]][j] * grad_y[i / K][j]scales和bias均传入grad_expanded_x同上grad_scales中需加上按专家索引查取的 bias 项grad_expanded_x[expanded_row_idx[i]][j] grad_y[i / K][j] * scales[i / K][i % K] grad_scales[i] Σ_j (expanded_x[expanded_row_idx[i]][j] bias[expert_idx[i]][j]) * grad_y[i / K][j]两个重要说明该反向算子仅计算grad_expanded_x和grad_scales不计算x1、x2、bias的梯度自动反向下这些位置的梯度返回None。布局差异expanded_row_idx的排布布局与正向不同。正向moe_finalize_routing使用(K, R)布局即idx[k * R row]本反向算子使用(R, K)布局即idx[row * K k]。通过自动反向调用时框架会自动完成布局转置手动调用时必须自行确保布局正确。函数原型与参数说明cann_ops_transformer.moe_finalize_routing_grad( grad_y: torch.Tensor, expanded_row_idx: torch.Tensor, expanded_x: Optional[torch.Tensor] None, scales: Optional[torch.Tensor] None, expert_idx: Optional[torch.Tensor] None, bias: Optional[torch.Tensor] None, drop_pad_mode: int 0, active_num: int 0, expert_num: int 0, expert_capacity: int 0, ) - Tuple[torch.Tensor, torch.Tensor]输入参数参数名类型可选/必选描述数据类型维度(shape)grad_yTensor必选正向输出y的梯度float16、bfloat16、float32(R, H)expanded_row_idxTensor必选行索引映射采用(R, K)布局注意与正向的(K, R)布局不同。Dropless 场景下值域为[0, R * K - 1]且不重复DropPad 场景下值域为[-1, E * C - 1]除 -1 外不重复-1 表示丢弃。当scales未传入时 K 必须为 1int32(R * K,)expanded_xTensor可选正向输入expanded_x用于计算grad_scales。传入scales时必须同时传入数据类型需与grad_y一致与 grad_y 一致(NUM_ROWS * K, H)或(E, C, H)scalesTensor可选正向输入scales即路由权重系数。传入scales时必须同时传入expanded_xfloat16、bfloat16、float32(R, K)expert_idxTensor可选正向输入expert_idx即专家索引。传入bias时必须同时传入值域为[0, E-1]int32(R, K)biasTensor可选正向输入bias即偏置量与 grad_y 一致(E, H)drop_pad_modeint可选丢弃/填充模式需与正向保持一致。0drop less1drop pad。默认值为 0--active_numint可选Dropless 场景下grad_expanded_x的最大输出行数仅当大于 0 且小于R * K时生效DropPad 场景下不生效。默认值为 0--expert_numint可选专家数。Dropless 场景下不生效DropPad 场景下必须大于 0且当bias传入时必须等于bias的第 0 维大小E。默认值为 0--expert_capacityint可选每个专家能够处理的 tokens 数drop_pad_mode1时必须传入且大于 0。默认值为 0--返回值参数名类型描述数据类型维度(shape)grad_expanded_xTensorexpanded_x的梯度。Dropless 场景下 shape 为(dim0, H)其中dim0为expanded_row_idx.numel()active_num0时或min(active_num, expanded_row_idx.numel())active_num 0时DropPad 场景下 shape 为(expert_num, expert_capacity, H)与 grad_y 一致(dim0, H)或(E, C, H)grad_scalesTensorscales的梯度。shape 为(R, K)其中K为scales的第二维scales传入时或 1scales未传入时。数据类型scales传入时与scales一致否则与grad_y一致float16、bfloat16、float32(R, K)上述输出 shape 的推断逻辑在源码中可以直接对应csrc 包装层 先根据drop_pad_mode预分配gradExpandedXOut与gradScalesOutDropPad 用{expertNum, expertCapacity, hidden}Dropless 用{dim0, hidden}且dim0受activeNum截断再调用ACLNN_CMD(aclnnMoeFinalizeRoutingV2Grad, ...)完成实际计算Python 侧注册的 Meta 实现 也按同样规则推导 shape保证编译期图与运行期一致。约束说明该接口支持训练场景下使用支持单算子模式和图模式。grad_y必须是 2 维张量expanded_row_idx必须是 1 维 int32 张量。drop_pad_mode仅支持取值 0 或 1不支持 2 或 3行排列模式。传入bias时expert_idx必须同时传入传入scales时expanded_x必须同时传入。drop_pad_mode1时expert_num必须大于 0bias传入时必须等于bias的第 0 维大小 Eexpert_capacity必须大于 0。expanded_x、bias的数据类型必须与grad_y一致scales的数据类型在 Atlas A2/A3 系列上必须与grad_y一致在 Ascend 950PR/950DT 上可以与grad_y不一致。该反向算子仅支持常规专家场景不支持 aclnnMoeFinalizeRoutingV4 特有特性正向的x、alpha1、alpha2、v以及有效的zero_expert_range、copy_expert_range、constant_expert_range。当正向使用了这些特性时调用自动反向会抛出NotImplementedError。各参数需与正向调用保持一致。这些约束在 C 包装层源码 中均有对应的TORCH_CHECK前置校验grad_y的 2 维与 float16/float32/bfloat16 类型、expanded_row_idx的 1 维 int32 类型、drop_pad_mode ∈ {0, 1}、bias/expert_idx与scales/expanded_x的成对传入关系、DropPad 模式下expertNum/expertCapacity必须为正均在进入 ACLNN 调用前完成校验报错信息清晰可定位。关于 V4 特性不支持的限制正向算子的 autograd 绑定实现 MoeFinalizeRoutingFn 中有明确实现_check_backward_supported会在ctx.has_v4_extras即提供了x/alpha1/alpha2/v或有效的 expert range或drop_pad_mode不在 (0, 1) 时抛出NotImplementedError其 docstring 也声明了forward - aclnnMoeFinalizeRoutingV4, backward - aclnnMoeFinalizeRoutingV2Grad的组合关系以及 V2Grad 只处理常规专家场景的原因。expanded_row_idx 布局转置反向调用最关键的细节这是该算子与正向最易混淆的一点。正向moe_finalize_routing在drop_pad_mode为 0 或 1 时expanded_row_idx按(K, R)布局排布即idx[k * R row]而反向算子要求(R, K)布局即idx[row * K k]。自动反向下框架在backward中通过_transpose_row_idx完成这一转换见 moe_finalize_routing.pydef _transpose_row_idx(expanded_row_idx, top_k, row_num): Convert expanded_row_idx from forward (K, R) layout to grad op (R, K) layout. if top_k 1: return expanded_row_idx.reshape(top_k, row_num).t().contiguous().reshape(-1) return expanded_row_idx也就是说先把一维索引按(K, R)还原成二维转置为(R, K)后展平正好得到反向算子要求的排布。手动调用反向时需要自己执行同样的 reshape 转置 展平操作。此外backward中还做了两件辅助工作值得注意合成 scales_synthesize_scales在scales为 None 但top_k 1时会构造一个全 1 的(R, K)张量作为有效scales传入从而保证grad_scales形状正确推断 shape 参数_infer_grad_shape根据正向保存的expanded_x形状推断active_num/expert_num/expert_capacityDropPad 模式返回(0, E, C)Dropless 模式返回(expanded_x.shape[0], 0, 0)。自动反向还遵循了文档声明的行为若x1/x2/bias需要梯度会发出警告说明这些位置的梯度不会被计算并返回None建议改用外部残差加法若expanded_x和scales都不需要梯度则直接跳过对反向算子的调用。调用示例单算子模式手动调用反向注意示例中对expanded_row_idx的手动布局转置import torch import torch_npu from cann_ops_transformer.ops import moe_finalize_routing, moe_finalize_routing_grad torch_npu.npu.set_device(0) NUM_ROWS 2 K 2 H 8 E 4 expanded_x torch.randn(NUM_ROWS * K, H, dtypetorch.float32, devicenpu) # 正向expanded_row_idx为(K, R)布局idx[k*R row] forward_row_idx torch.tensor([0, 1, 2, 3], dtypetorch.int32, devicenpu) scales torch.randn(NUM_ROWS, K, dtypetorch.float32, devicenpu) expert_idx torch.tensor([[0, 1], [2, 3]], dtypetorch.int32, devicenpu) bias torch.randn(E, H, dtypetorch.float32, devicenpu) # 正向 y moe_finalize_routing(expanded_x, forward_row_idx, biasbias, scalesscales, expert_idxexpert_idx, kK) # 手动调用反向需将expanded_row_idx从(K, R)布局转置为(R, K)布局 grad_row_idx forward_row_idx.reshape(K, NUM_ROWS).t().contiguous().reshape(-1) grad_y torch.randn_like(y) grad_expanded_x, grad_scales moe_finalize_routing_grad( grad_y, grad_row_idx, expanded_xexpanded_x, scalesscales, expert_idxexpert_idx, biasbias, drop_pad_mode0, ) print(grad_expanded_x.shape) # torch.Size([4, 8]) print(grad_scales.shape) # torch.Size([2, 2])自动反向autograd 自动触发当正向moe_finalize_routing中expanded_x、scales等可微输入的requires_grad为 True且未使用 V4 特有特性x、alpha1、alpha2、v、zero_expert_range等时loss.backward()会自动触发本算子无需手动调用仅在需要显式控制梯度的场景下才保留手动调用路径。import torch import torch_npu from cann_ops_transformer.ops import moe_finalize_routing torch_npu.npu.set_device(0) NUM_ROWS 4 K 2 H 8 expanded_x torch.randn(NUM_ROWS * K, H, dtypetorch.float32, devicenpu, requires_gradTrue) row_idx torch.arange(NUM_ROWS * K, dtypetorch.int32, devicenpu) scales torch.randn(NUM_ROWS, K, dtypetorch.float32, devicenpu, requires_gradTrue) # 正向expanded_x.requires_gradTrue时自动启用autograd y moe_finalize_routing(expanded_x, row_idx, scalesscales, kK) # 反向自动触发moe_finalize_routing_grad loss y.sum() loss.backward() print(expanded_x.grad.shape) # torch.Size([8, 8]) print(scales.grad.shape) # torch.Size([4, 2])确定性计算与配套接口确定性计算本算子默认支持确定性计算与 aclnn 接口文档 中aclnnMoeFinalizeRoutingV2Grad 默认确定性实现的说明一致。配套接口本算子是 moe_finalize_routing 的反向算子二者参数需保持对齐。如需使用原生两段式接口可查阅 aclnnMoeFinalizeRoutingV2Grad 接口文档其第一段aclnnMoeFinalizeRoutingV2GradGetWorkspaceSize负责入参校验并返回 workspace 大小与执行器第二段aclnnMoeFinalizeRoutingV2Grad执行实际计算典型报错包括ACLNN_ERR_PARAM_NULLPTR161001必选输入输出为空指针、ACLNN_ERR_PARAM_INVALID161002数据类型/格式不支持、ACLNN_ERR_INNER_NULLPTR561103shape 或取值不满足要求。最后再强调一次使用要点该反向算子只产出grad_expanded_x与grad_scales两个梯度如果需要对x1/x2求梯度建议使用外部残差加法将其作为正向输入而不是依赖本算子手动调用时务必把expanded_row_idx从正向的(K, R)布局转置为反向的(R, K)布局这是保证梯度正确性的前提。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
