Pull Request已成功合入, 合并人@CANN-robot
(感谢 sinobr 的贡献)变更摘要
此 PR 修复了 swiglu_group_grad 反向算子中 weight 参数的形状校验过严的问题:原先要求 weight 与 grad_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_weight 与 weight 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 L2CheckShape)、swiglu_group_grad_infershape.cpp(InferShapeForSwigluGroupGrad)、swiglu_group_grad_regbase_tiling.cpp(ParseOptionalInputs)和swiglu_group_grad_tiling_base.cpp(ParseOptionalInputs)中,全部移除了对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.cpp和swiglu_group_grad_tiling_base.cpp中分别新增ShapeElementNumHost和ShapeElementNum,用于计算任意非零 rank 张量的元素总数;在swiglu_group_grad_infershape.cpp中新增TryGetKnownElementNum,支持含未知维度的 numel 计算。 -
文档全面同步更新:
aclnn_swiglu_group_grad.h、swiglu_group_grad.h、swiglu_group_grad_proto.h、swiglu_group_grad_def.cpp和gen_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 不匹配失败用例。


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
/approveor/lgtm- Commenting
/approveimplies 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. 👍


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.


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


描述
swiglu_group_backward原校验要求weight.dim() == grad_output.dim()且weight.shape[-1] == 1,导致 flat 形态的 weight(如[T]/[B*S])在 autograd 路径被拒绝。kernel 实际按一维行偏移访问 weight(
weightGm_[row]),只依赖元素个数,不感知逻辑 rank。前向swiglu_grouptiling 也只校验numel == bs。反向校验与前向不一致,autograd 桥透传 flat weight 时反向校验失败。修改内容
将反向六层 weight 校验从 rank 对齐 + 末维=1 改为 element num 校验,跟前向 tiling 对齐:
torch_extension/swiglu_group_grad.pytorch_extension/csrc/swiglu_group_grad.cppop_api/aclnn_swiglu_group_grad.cppop_host/swiglu_group_grad_infershape.cppswiglu_group_quant_grad对齐)op_host/arch35/swiglu_group_grad_tiling_base.cppop_host/arch35/swiglu_group_grad_regbase_tiling.cpp命名风格统一跟前向
swiglu_grouptiling 对齐: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生成声明