[Bug]: triton_experimental 原生 TorchBench BERT 训练精度异常及编译失败-pytorch-AtomGit

已关闭

[Bug]: triton_experimental 原生 TorchBench BERT 训练精度异常及编译失败 #3813

htchu创建于 8月7日关闭于 27 天前

htchu成员

8月7日创建

在提交新问题之前，请确保您已经在社区中搜索过相关问题，并使用了社区中提供的资源/工具后，仍未找到满意的解决方式。

⚠️ 安全信息提醒：请仔细检查提供的文本内容，确保其不包含敏感数据信息，包括但不限于：

API 令牌或密钥
密码或身份验证凭证
私有网址或接口地址
个人或机密数据
...

在分享配置信息或代码示例时，请将敏感信息脱敏处理，或使用 <TOKEN> 等占位符替代原有内容。

环境信息

NPU：Ascend 910B2（A2）
CANN：9.1.0
PyTorch：2.13 分支自编包
torch_npu：2.13 分支自编包
后端：triton_experimental
测试集：原生 TorchBench BERT 系列源码模型
模式：FP32 training accuracy

🐛 问题描述

按照 benchmarks/torchbench/README.md 安装 TorchBench 源码包后，五个原生 BERT 系列模型均可完成 eager 训练，但 Inductor + triton_experimental 初始测试只有 BERT_pytorch 通过：

模型	初始结果	主要现象
`BERT_pytorch`	`pass_accuracy`	无异常
`hf_Bert`	`fail_accuracy`	第 3 步 loss 突变，第 4 步起 NaN
`hf_DistilBert`	`fail_to_run`	`Unregistered range symbol: r0_1`
`hf_Albert`	`fail_to_run`	reduction/permute autotune 无可运行 tiling，当前版本已不再复现
`fastNLP_Bert`	`fail_to_run`	graph break 后出现 237/474 长度不一致

其中 hf_Bert、hf_DistilBert 和 fastNLP_Bert 分别对应三个确定性的后端缺陷。

1. `hf_Bert`：不受支持的 OUTER split-reduction 破坏融合输出

失败 backward kernel 同时物化：

一个 2048x768 pointwise clone/permute 输出；
一个 1x768 reduction 输出；
一个包含多个 reduction-tree 节点的 [4, 512] reduction。

现有 _npu_rsplit_outer_applicable() 只检查 reduction 类型、输出 x 规模和 reduction 规模，没有限制 live output 数量和 reduction-tree 结构。partial+combine 重写随后按输出顺序选择 output_buffers[0]，将 2048x768 pointwise store 重定向到仅 48x768 的 partial workspace，造成越界写和错误 combine。同时，嵌套 reduction 循环没有应用 r0_lo/r0_hi，每个 core 会重复处理完整 reduction 范围。

表现为：前向输出仍正确，但首步已有大量梯度失配，连续训练第 3 步突变、第 4 步产生 NaN。

2. `hf_DistilBert`：dual-view reduction fold 创建未注册的同名 Symbol

backward 的 buf352 输出 shape 为 [768]，reduction ranges 为 [4, 512]。codegen 在推断 load block shape 时失败：

AssertionError: Unregistered range symbol: r0_1

诊断显示 range_tree_nodes 中实际存在打印名称为 r0_1 的 key。失败地址原本只包含 alias chain：

r0_2 + 512 * r0_3

_fold_dualview_reduction_index() 将它折回 flat alias 时，在 flat symbol 不在原表达式中的 fallback 分支创建了：

sympy.Symbol(inner.name)

range-tree symbol 带 integer/nonnegative assumptions。新建的普通 Symbol 虽然也打印为 r0_1，但不等于注册表中的正式迭代 Symbol，导致字典查找失败。

3. `fastNLP_Bert`：in/out boundary downcast 丢失输入内容

fastNLP 预处理在 .item() graph break 后报错：

The expanded size of the tensor (237) must match the existing size (474)

独立预处理最小复现结果：

eager_piece    [474, 474, 474, 474]
compiled_piece [237, 237, 237, 237]
eager_seq      [474, 474, 474, 474]
compiled_seq   [474, 474, 474, 474]

生成 kernel 将运行时 int64 in_out_ptr0 下转换为 *i32，并通过 mutated_arg_names 标记其写入语义。launcher wrapper 将纯 out_ptr 和 in_out_ptr 一律替换为未初始化的 int32 torch.empty_like()。该 in_out_ptr0 会先读取 advanced-index 结果，再执行 masked-fill 写回；使用空临时区会丢失输入内容，后续 reduction 因此错误。

复现步骤

export ASCEND_RT_VISIBLE_DEVICES=7
source env.sh keep

python -u benchmarks/torchbench/torchbench.py \
  --accuracy --cold-start-latency --train --float32 \
  --backend inductor --npu-backend triton \
  --only hf_Bert --iterations 5 --accu-summary

python -u benchmarks/torchbench/torchbench.py \
  --accuracy --cold-start-latency --train --float32 \
  --backend inductor --npu-backend triton \
  --only hf_DistilBert --iterations 1 --accu-summary

python -u benchmarks/torchbench/torchbench.py \
  --accuracy --cold-start-latency --train --float32 \
  --backend inductor --npu-backend triton \
  --only fastNLP_Bert --iterations 1 --accu-summary

期望行为

OUTER split-reduction 只接受当前 partial+combine 重写能够正确表达的单一 live reduction 输出和单一扁平 reduction-tree 节点。
dual-view reduction fold 必须复用 range-tree 已注册、assumptions 完整的迭代 Symbol。
boundary downcast 对纯输出可以使用空临时区，对 in/out 或 mutated 参数必须保留原始输入内容，并在 kernel 后写回原 dtype buffer。
五个原生 TorchBench BERT 模型的 prediction、loss、梯度、参数更新和 buffers 均通过精度检查。

欢迎加入社区，感谢您对社区的贡献 🎉!