Ascend for PyTorch 上 `torch.where` 反向传播精度差异分析
**作者**:昇腾实战派
**知识地图**:[https://blog.csdn.net/Lumos\_Lovegood/article/details/161601003](https://blog.csdn.net/Lumos_Lovegood/article/details/161601003)
## 背景概述
在深度学习框架中,`torch.where` 是一个常用的条件选择算子,其正向操作为 select 操作,不涉及算术运算。然而,在反向传播过程中,当输入张量存在广播(broadcasting)时,需要对扩展维度进行 `reduce_sum` 归约操作。由于不同硬件平台(如 CPU 与昇腾 NPU)在浮点数累加顺序上的差异,可能导致反向传播梯度存在微小精度误差。本文以实际开发中遇到的精度检测问题为例,分析该问题的根因,并给出相应的解决方案与精度标准建议。
## 问题代码
使用 atk 工具进行 torch.where 反向传播精度检测,覆盖精度检测、小值域检测、内存检测和确定性检测四个场景:
```yaml
# 精度检测
atk case -f 1000/torch.where_backward.yaml -p 1000/ascend_generate_torch_where.py
atk task -c result/torch.where_backward/json/all_torch.where_backward.json -n 0000/npu_cpu.yaml -to 7200
# 小值域
atk case -f 1000/float_small_nd_torch.where_backward.yaml -p 1000/ascend_generate_torch_where.py
atk task -c result/float_small_nd_torch.where_backward/json/all_float_small_nd_torch.where_backward.json -n 0000/npu_cpu.yaml -to 7200
atk case -f 1000/float_small_torch.where_backward.yaml -p 1000/ascend_generate_torch_where.py
atk task -c result/float_small_torch.where_backward/json/all_float_small_torch.where_backward.json -n 0000/npu_cpu.yaml -to 7200
# 内存检测(align_32B: True)
atk case -f 1000/torch.where_deterministic_backward.yaml -p 1000/ascend_generate_torch_where.py
mssanitizer --cache-size 8192 -t memcheck --log-level=error --log-file=./torch_where/memory.log -- /usr/local/bin/atk task -n 0000/nodes_mem_npu.yaml -c result/torch.where_deterministic_backward/json/all_torch.where_deterministic_backward.json -ms -msl ./torch_where/memory.log -to 7200 -sp
# 确定性(align_32B: True,复用内存检测json)
atk case -f 1000/torch.where_deterministic_backward.yaml -p 1000/ascend_generate_torch_where.py
atk task -c result/torch.where_deterministic_backward/json/all_torch.where_deterministic_backward.json -n 0000/npu_cpu_dc.yaml -to 7200
```
## 报错日志
用例执行报告显示,失败用例采用 `torch.equal` 作为精度标准,要求 NPU 和 CPU 输出二进制一致:
```json
{
"output_grad_1.pt": {"result": false, "error_info": "torch.equal_failed"},
"output_grad_0.pt": {"result": true, "max_diff": 0.0, "diff_percent": 100.0},
"output_grad_2.pt": {"result": true, "max_diff": 0.0, "diff_percent": 100.0}
}
```
## 根因
在昇腾 NPU 上,反向传播时广播场景下需要对扩展维度进行 **reduce_sum 归约**,CPU 和 NPU 对浮点数的 reduce_sum 运算做不到二进制完全一致(浮点数累加顺序不同导致微小差异)。因此对于"正向有广播、反向需归约"的算子输入,都会引入梯度误差。
该问题与昇腾 CANN 软件栈的算子调度策略以及硬件浮点运算单元(AI Core)的累加顺序有关,属于不同计算架构下的固有差异。
## 分析过程
- 实际误差范围极小,且误差**仅**出现在 `other` 的反向梯度上,`input` 的梯度完全一致
- 分析发现失败用例中 `other` 的 shape 全部是需要向 `input` 广播的场景
- torch.where 正向操作为 select 操作,不涉及算术运算,误差不可能来自正向
- 将失败用例的输入 shape 改为一致消除广播后,反向结果完全一致,确认误差由归约引入
## 修复方案
此类场景(正向 select/非算术操作 + 反向广播归约)的精度对比标准不应采用 `torch.equal`(二进制一致),而应参照 reduce_sum 采用双千双万标准。
## 结果
修改输入 shape 消除广播后,torch.where 反向精度检测通过。该问题本质上属于浮点数 reduce_sum 在不同硬件平台上的二进制精度差异,不属于昇腾 NPU 算子实现 bug。建议参照 reduce_sum 采用双千双万标准。
## 总结与建议
- **问题本质**:浮点数 `reduce_sum` 在不同硬件平台上的累加顺序不同,导致二进制结果不完全一致。
- **影响范围**:所有正向为 select/非算术操作、反向需归约的算子输入,都可能引入类似梯度误差。
- **建议**:
- 对于此类场景,建议采用双千双万标准进行精度对比,而非严格的二进制对齐。
- 在测试用例设计时,可考虑消除广播场景以验证算子实现的正确性。
- 昇腾 NPU 用户应了解不同架构下的浮点运算差异,合理设置精度阈值。
通过以上分析,我们明确了 `torch.where` 反向传播精度差异的根因,并给出了可行的解决方案,为后续类似问题的排查和处理提供了参考。
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。
更多推荐
所有评论(0)