| Fix moe unpermute grad probs dtype guard
Co-authored-by: wang_ziqi<wangziqi4@huawei.com>
# message auto-generated for no-merge-commit merge:
!4846 merge fix/grad-probs-dtype-26.0.0 into 26.0.0
Fix moe unpermute grad probs dtype guard
Created-by: wang-ziqi-code
Commit-by: wang_ziqi
Merged-by: ascend-robot
Description: # 【合入来源】
- [ ] 需求
- [x] 问题单
- [ ] issue/工单
- [ ] 重构优化
- [ ] 资料更新
# 【修改方案】修复 npu_moe_token_unpermute_grad 在 probs=None 场景下的回归问题
1. **问题现象**:在合入 grad_probs dtype 修复后,npu_moe_token_unpermute 的反向路径在 probs=None、即 topk=1 场景下出现报错:
The tensor has a non-zero number of elements, but its data is not allocated yet.
2. **问题根因**:
- grad_probs 的 size 推导逻辑使用的是 probs.has_value() && probs.value().defined();
- 但 grad_probs 的 dtype 推导逻辑只判断了 probs.has_value();
- 当 Python 侧传入 None,在该链路下可能表现为“optional 有值但 tensor 未 defined”,导致 size 侧将 grad_probs 视为无效输出,而 dtype 侧仍尝试按有效输出处理,最终引发输出分配状态不一致,反向执行报错。
3. **修改内容**:
- 修改 op_plugin/config/op_plugin_functions.yaml
- 将 grad_probs 的 dtype 判定从:
probs.has_value() ? probs->scalar_type() : grad_unpermuted_tokens.scalar_type()
修正为:
probs.has_value() && probs->defined() ? probs->scalar_type() : grad_unpermuted_tokens.scalar_type()
- 使 dtype 推导逻辑与 size 推导逻辑、以及工程内 optional tensor 的通用判定方式保持一致,避免在 probs=None 场景下误访问未定义 tensor。
4. **修改范围说明**:
- 本次为最小修复,仅修正生成侧 dtype 判定条件;
- 不调整 Python meta 层 probs is None -> grad_probs is None 的现有语义;
- 不修改算子 schema、derivatives 配置和接口行为。
# 【资料变更】
不涉及。
# 【接口变更】
不涉及。本次修复仅纠正 grad_probs 输出 dtype 的判定条件,不改变算子签名和对外行为。
# 【功能验证】
1. 代码级验证:
- 确认 grad_probs 的 dtype 判定与 size 判定统一为“has_value() && defined() 才视为有效 optional tensor”
2. 场景验证:
- probs=None / topk=1 场景下,npu_moe_token_unpermute 反向路径不再因未定义 tensor 参与 dtype 推导而报错
- 普通 probs 存在场景下,grad_probs 仍保持使用 probs 的 dtype
3. 回归风险:
- 修改仅影响 npu_moe_token_unpermute_grad 的 grad_probs dtype 推导条件,影响面小
# 【CheckList】
- [x] 代码注释完备,正确记录错误日志
- [x] 代码实现进行了返回值、空指针等校验
- [x] PR标题正确使用类型标签,如:feat、fix、refactor、docs、test等
- [x] PR持续集成流水线(CI)执行通过,代码检查无异常
See merge request: Ascend/op-plugin!4846 | 1 个月前 |