已合并
fix: 放宽 swiglu_group_grad 的 weight 形态校验为 element num #8541
fix: 放宽 swiglu_group_grad 的 weight 形态校验为 element num #8541
已合并
sinobr创建于 8月11日
sinobr
sinobr
8月11日

描述

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

kernel 实际按一维行偏移访问 weight(weightGm_[row]),只依赖元素个数,不感知逻辑 rank。前向 swiglu_group tiling 也只校验 numel == bs。反向校验与前向不一致,autograd 桥透传 flat weight 时反向校验失败。

修改内容

将反向六层 weight 校验从 rank 对齐 + 末维=1 改为 element num 校验,跟前向 tiling 对齐:

文件 改动
Torch Meta torch_extension/swiglu_group_grad.py rank/逐维/末维 → element num
C++ bridge torch_extension/csrc/swiglu_group_grad.cpp 同上
ACLNN L2 op_api/aclnn_swiglu_group_grad.cpp 同上
infershape op_host/swiglu_group_grad_infershape.cpp 删旧 rank 校验,加 element num 校验(跟 swiglu_group_quant_grad 对齐)
tiling base op_host/arch35/swiglu_group_grad_tiling_base.cpp rank/末维/逐维 → element num
regbase tiling op_host/arch35/swiglu_group_grad_regbase_tiling.cpp 同上

命名风格统一跟前向 swiglu_group tiling 对齐:ShapeElementNum / weightElementNum / totalRows

grad_weight 输出形态保持 weight.new_empty(weight.shape),与输入 weight 形态一致,autograd 形态匹配天然正确。kernel 代码未改动。

关联的Issue

https://gitcode.com/cann/ops-nn/issues/4750

测试

现有正向用例的 weight 形态从 [T,1] / [B,S,1] 改为 flat [T] / [B*S],验证 flat weight 通过各层校验:

  • test_swiglu_group_backward.py:weight (2,4,1)(8,)
  • test_aclnn_swiglu_group_grad.cpp:weight/dWeight {4,1}{4}
  • test_swiglu_group_grad_tiling.cpp:加 weightShape = {8}
  • test_swiglu_group_grad_infershape.cpp:加 weightShape = {4}
    950 环境验证 10 个 autograd 用例全部 PASS(含 flat weight + clamp + bf16 组合)。
    冒烟测试通过

类型标签

AI/Agent生成声明

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 sinobr 的贡献)
sinobrsinobr
8月11日 创建了 pull request,commit 688f0559
atomgit-bot
atomgit-bot
8月11日 评论:

变更摘要

此 PR 修复了 swiglu_group_grad 反向算子中 weight 参数的形状校验过严的问题:原先要求 weightgrad_output rank 相同且末维为 1(如 [T,1] / [B,S,1]),导致前向已支持的 flat weight(如 [T] / [B*S])在反向路径被拒绝。由于 kernel 内部按一维行偏移访问逐 token weight,仅依赖元素个数而非逻辑 rank,因此将校验契约统一为 weight.numel() == prod(grad_output.shape[:-1]),支持 [T][T,1][B,S][B,S,1] 等等价形态,同时保持 grad_weightweight shape 完全一致的输出契约。改动覆盖了 Torch Meta、C++ bridge、ACLNN L2、infer-shape、RegBase/SIMT tiling 全部层级,并同步更新了文档和测试。

主要改动

  • Weight 校验契约从 rank/末维匹配改为 numel 匹配:在 swiglu_group_grad.cpp(C++ bridge)、swiglu_group_grad.py(Python 端校验)、aclnn_swiglu_group_grad.cpp(ACLNN L2 CheckShape)、swiglu_group_grad_infershape.cppInferShapeForSwigluGroupGrad)、swiglu_group_grad_regbase_tiling.cppParseOptionalInputs)和 swiglu_group_grad_tiling_base.cppParseOptionalInputs)中,全部移除了对 weight.dim() == grad_output.dim()weight.shape[-1] == 1 的逐维校验,改为计算 totalRows = prod(grad_output.shape[:-1]) 后比较 weight.numel() == totalRows

  • 引入 ShapeElementNum/ShapeElementNumHost 辅助函数:在 swiglu_group_grad_regbase_tiling.cppswiglu_group_grad_tiling_base.cpp 中分别新增 ShapeElementNumHostShapeElementNum,用于计算任意非零 rank 张量的元素总数;在 swiglu_group_grad_infershape.cpp 中新增 TryGetKnownElementNum,支持含未知维度的 numel 计算。

  • 文档全面同步更新aclnn_swiglu_group_grad.hswiglu_group_grad.hswiglu_group_grad_proto.hswiglu_group_grad_def.cppgen_input_data.py 中的 weight/grad_weight shape 描述均从 (T,1)/(B,S,1) 更新为"非标量 Tensor,元素个数为 T 或 B*S"及"shape 与 weight 相同"的 numel 语义。

  • 新增多层级回归测试:在 Torch Meta 测试(test_swiglu_group_backward.py)中新增 parametrize 覆盖 [8][8,1][2,4][2,4,1] 四种等 numel 形态及 numel 不匹配的错误用例;在 InferShape(test_swiglu_group_grad_infershape.cpp)、Tiling(test_swiglu_group_grad_tiling.cpp)和 ACLNN L2 API(test_aclnn_swiglu_group_grad.cpp)测试中均新增 flat weight、[B,S] 形态通过用例及 numel 不匹配失败用例。

likedislike
不准确?
atomgit-bot
atomgit-bot
8月11日 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
CANN-robotCANN-robot成员
8月11日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
8月11日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here
You can get sig-info at here


PR Approval Progress

Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
quant 汤平川, 王星 (2/2) 汤平川, 王星 (2/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

qq_51545127, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了195条消息 查看更多
CANN-robot
CANN-robot成员
8月13日 评论:

The MR can not be merged, because of CodeReview discussion not resolved

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike
Lliuchuangdev成员
8月13日 解决了最后一个问题
CANN-robotCANN-robot成员
8月13日 关闭了关联的issue
CANN-robotCANN-robot成员
8月13日 合入了pull request
CANN-robot
CANN-robot成员
8月13日 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike