swiglu_group_backward 反向算子的 weight 校验要求 weight.dim() == grad_output.dim() 且 weight.shape[-1] == 1,导致 flat 形态的 weight(如 [T] / [B*S])在 autograd 路径被拒绝。
swiglu_group_backward
weight.dim() == grad_output.dim()
weight.shape[-1] == 1
[T]
[B*S]
swiglu_group
weight.numel() == bs
weightGm_[row]
[T,1]
import torch, torch_npu, cann_ops_nn.ops x = torch.randn((8, 256), dtype=torch.float16, device="npu:0", requires_grad=True) weight = torch.randn((8,), dtype=torch.float32, device="npu:0", requires_grad=True) # flat [T] y = torch.ops.cann_ops_nn.swiglu_group(x, weight=weight, clamp_limit=0.0) y.backward(torch.ones_like(y)) # 报错: weight rank must equal grad_output rank
PR: https://gitcode.com/cann/ops-nn/pull/8541
将反向六层 weight 校验从 rank 对齐 + 末维=1 改为 element num 校验,跟前向 tiling 对齐。
问题描述
swiglu_group_backward反向算子的 weight 校验要求weight.dim() == grad_output.dim()且weight.shape[-1] == 1,导致 flat 形态的 weight(如[T]/[B*S])在 autograd 路径被拒绝。根因
swiglu_grouptiling 只校验weight.numel() == bs,不限制 rank,flat weight 能过swiglu_group_backward校验weight.dim() == grad_output.dim()+weight.shape[-1] == 1,flat weight 被拦weightGm_[row]),不感知逻辑 rank,[T]和[T,1]跑起来完全一样复现
import torch, torch_npu, cann_ops_nn.ops x = torch.randn((8, 256), dtype=torch.float16, device="npu:0", requires_grad=True) weight = torch.randn((8,), dtype=torch.float32, device="npu:0", requires_grad=True) # flat [T] y = torch.ops.cann_ops_nn.swiglu_group(x, weight=weight, clamp_limit=0.0) y.backward(torch.ones_like(y)) # 报错: weight rank must equal grad_output rank修复
PR: https://gitcode.com/cann/ops-nn/pull/8541
将反向六层 weight 校验从 rank 对齐 + 末维=1 改为 element num 校验,跟前向 tiling 对齐。