已关闭
swiglu_group_backward weight 形态校验过严导致 autograd 反向失败 #4750
sinobr创建于  27 天前关闭于  26 天前
sinobr
sinobr
27 天前 创建

问题描述

swiglu_group_backward 反向算子的 weight 校验要求 weight.dim() == grad_output.dim()weight.shape[-1] == 1,导致 flat 形态的 weight(如 [T] / [B*S])在 autograd 路径被拒绝。

根因

  • 前向 swiglu_group tiling 只校验 weight.numel() == bs,不限制 rank,flat weight 能过
  • 反向 swiglu_group_backward 校验 weight.dim() == grad_output.dim() + weight.shape[-1] == 1,flat weight 被拦
  • kernel 实际按一维行偏移访问 weight(weightGm_[row]),不感知逻辑 rank,[T][T,1] 跑起来完全一样
  • 前后向对 weight 形态的校验契约不一致,autograd 桥透传 flat weight 时反向校验失败

复现

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 对齐。

likedislike
yuning_chenyuning_chen成员
27 天前 将 qq_51545127 设为负责人
CANN-robotCANN-robot成员
26 天前 关闭了 issue
CANN-robotCANN-robot成员
26 天前 添加了label:resolved
sinobrsinobr
21 天前 关联了pull request:fix: 放宽 swiglu_group_grad 的 grad_y 维度限制并对齐前向 clamp_limit 校验