[RFC]: 删除 NPU Sanitizer 中不可达的 Autograd Dispatch Key 兼容 Patch-pytorch-AtomGit

已关闭

[RFC]: 删除 NPU Sanitizer 中不可达的 Autograd Dispatch Key 兼容 Patch #3669

No_neck创建于 8月5日关闭于 8月7日

8月5日创建

项目	内容
状态	Reviewing
作者	No_neck
创建日期	2026-08-05
更新日期	2026-08-05
所属模块	Core / NPU Sanitizer / Stream Check
相关 Issue	待创建或由维护者指定
相关 PR	Ascend/pytorch#43468
历史引入 PR	Ascend/pytorch#10733

1. 概述

1.1 简介

本文提议删除 torch_npu/npu/_stream_check.py 中一段 NPU Sanitizer 私有的 Autograd Dispatch Key 兼容逻辑。该逻辑在 2024 年首次引入 NPU Sanitizer 时加入，通过五个公开算子名称组成的白名单，调用 PyTorch 私有接口修改当前线程的 Dispatch TLS 状态。

待删除代码包括：

self.npu_adjust_autograd 白名单；
enable_autograd() 方法；
aten_api = func.__name__.split(".")[0] 临时变量；
self.enable_autograd(aten_api) 调用；
test/npu/test_stream_check.py 中仅验证上述旧实现的测试和字段断言。

本提议不删除 NPU Sanitizer，不改变其跨 Stream 数据竞争、Event 同步、内存生命周期和 record_stream 检测能力。变更的目标仅是消除当前 master / PyTorch 2.13 组合中真实算子路径无法命中的历史兼容 Patch。

1.2 背景与动机

NPU Sanitizer 是运行时检查工具。Stream Check 模式通过 TorchDispatchMode 拦截 ATen 算子，收集算子输入、输出对应的内存地址以及当前 NPU Stream，并结合 Event、Stream、内存分配释放和 record_stream 回调，检查潜在的跨 Stream 未同步访问。

历史兼容代码如下：

self.npu_adjust_autograd = [
    "adaptive_avg_pool2d", "batch_norm",
    "log_softmax", "nll_loss", "to"
]

def enable_autograd(self, aten_api):
    if aten_api in self.npu_adjust_autograd:
        torch._C._dispatch_tls_set_dispatch_key_excluded(
            torch._C.DispatchKey.AutogradFunctionality,
            False,
        )

该调用中 False 的准确含义不是“排除 Autograd”，而是将 AutogradFunctionality 从当前线程 TLS 的 excluded 集合中移除。对 NPU Tensor 而言，其主要效果是允许 AutogradPrivateUse1 再次参与算子分发。

这段代码存在以下维护问题：

依赖 torch._C._dispatch_tls_set_dispatch_key_excluded 私有接口，跨 PyTorch 版本不保证稳定。
白名单使用公开 Python API 名称，而 PyTorch 2.13 实际进入 TorchDispatchMode 的是内部 ATen 名称，名称无法精确匹配。
当前真实路径仍会对每个算子执行字符串解析和函数调用，但 setter 所在真分支始终不执行。
白名单没有记录明确的故障版本、原始复现、上游修复提交和删除条件。
长期保留不可达逻辑会增加 PyTorch 升级、代码审查和测试维护成本。

1.3 历史追溯

该兼容逻辑由以下提交引入：

Commit：a6e615ee59b3018457cb34e2f0d433c90eae1b4d
日期：2024-04-29
标题：Enable NPU Sanitizer
关联 PR：Ascend/pytorch#10733

从引入至本次分析时，相关白名单和 enable_autograd() 代码未发生实质修改。PyTorch 官方 CUDA Sanitizer 中也不存在相同的算子白名单或 TLS setter 调用，因此该逻辑属于 torch_npu 下游私有兼容 Patch，而不是需要持续与 PyTorch CUDA Sanitizer 同步的通用实现。

1.4 目标

删除 PyTorch 2.13 真实算子路径中不可达的 Autograd Dispatch Key 兼容逻辑。
移除对 PyTorch Dispatch TLS 私有 setter 的无效依赖。
保证删除前后 Sanitizer 拦截到的算子、顺序和次数一致。
保证五类历史关联算子的输出、Autograd 图、反向梯度和状态 Tensor 一致。
保证 Stream Check、Event、内存生命周期和 record_stream 官方测试无回退。
明确结论的版本边界、风险和回滚条件。

1.5 非目标

不修改 PyTorch Dispatcher 或 Autograd 的实现。
不修改任何算子的梯度公式或 NPU kernel。
不删除其他 NPU Sanitizer C++ 埋点和 trace callback。
不重构 _kernel_check.py。
不改变 Tensor.is_cuda 等其他兼容逻辑。
不改变 callback 注册和生命周期管理方式。
不改变跨 Stream 冲突及 record_stream 的判定规则。
不通过本 RFC 推断 PyTorch 2.10～2.12 等旧维护分支也可直接删除该逻辑。
不顺带调整异常日志、无关 import 或代码风格。

2. 用例分析

2.1 Sanitizer 的用户入口和调用链

用户可以显式调用 enable_npu_sanitizer() 启用 Sanitizer；当进程环境中存在 TORCH_NPU_SANITIZER 时，torch_npu 的可选功能初始化也会导入 Sanitizer 模块并调用该入口。Stream Check 模式下，核心调用链如下：

用户启用 NPU Sanitizer
    ↓
enable_npu_sanitizer()
    ↓
全局 npu_sanitizer.enable()
    ↓
NPUSanitizer.enable_stream_check()
    ├─ 创建 NPURecordStreamHandler
    ├─ 创建 NPUSanitizerDispatchMode
    ├─ 调用 dispatch.__enter__()，激活 TorchDispatchMode
    └─ 注册 Event / Stream / Memory / record_stream 回调
    ↓
用户执行 NPU ATen 算子
    ↓
NPUSanitizerDispatchMode.__torch_dispatch__()
    ├─ 解析输入内存访问
    ├─ 调用真实算子 func(*args, **kwargs)
    ├─ 解析输出内存访问
    └─ check_errors() 将访问记录交给 EventHandler

历史兼容 Patch 位于 __torch_dispatch__() 的输入解析之前：

__torch_dispatch__()
    ↓
从 func.__name__ 提取 aten_api
    ↓
enable_autograd(aten_api)
    ↓
仅当名称精确命中五项白名单时修改 TLS excluded

2.2 类和函数职责

符号	调用时机	作用	本次是否修改
`NPUSanitizer`	导入模块时创建全局实例；启用 Sanitizer 时调用其方法	管理 Sanitizer 模式、EventHandler、DispatchMode、trace callback 和启用状态	否
`NPUSanitizer.enable()`	用户启用 Sanitizer 时	根据配置选择 Kernel Check 或 Stream Check，并激活 NPU trace	否
`NPUSanitizer.enable_stream_check()`	未配置 Kernel Check 路径时	创建 Stream Check Handler 和 DispatchMode，并注册回调	否
`NPUSanitizerDispatchMode`	Stream Check 启用时实例化	使用 `TorchDispatchMode` 拦截 NPU ATen 算子	仅删除历史字段和方法
`__torch_dispatch__()`	Mode 激活后，每个被拦截的 ATen 算子执行时	收集输入输出访问，执行算子并触发冲突检查	仅删除无效兼容调用点
`enable_autograd()`	历史实现中，每次普通算子进入 `__torch_dispatch__()` 时	名称命中白名单后，从 TLS excluded 移除 `AutogradFunctionality`	删除
`check_errors()`	算子执行且获得当前 Stream 后	将读写地址、输出、schema 和 Stream 交给 Handler 检查	否

2.3 五类历史白名单用例

PyTorch 2.13 环境中，公开 API 与实际进入 DispatchMode 的内部 ATen 名称不一致：

公开用例	必须观察到的内部 ATen 名称	基线中进入 `enable_autograd()`
`adaptive_avg_pool2d`	`_adaptive_avg_pool2d`、`_adaptive_avg_pool2d_backward`	40 次
`batch_norm`	`native_batch_norm`、`native_batch_norm_backward`	103 次
`log_softmax`	`_log_softmax`、`_log_softmax_backward_data`	40 次
`nll_loss`	`nll_loss_forward`、`nll_loss_backward`	37 次
`to`	`_to_copy`	43 次

必须区分“进入函数”和“执行真分支”：

五个真实用例均进入了 enable_autograd()，因此覆盖了待删函数和调用点。
传入的内部 ATen 名称均未精确命中历史公开名称白名单，因此真实路径没有执行 TLS setter。
直接分支探针向 enable_autograd() 传入五个历史名称时，均观察到 before_excluded=True、after_excluded=False，证明旧代码本身具有 TLS 副作用，只是该副作用在当前真实路径不可达。

2.4 用户可见行为要求

删除后，用户启用 NPU Sanitizer 执行上述算子时，应满足：

前向执行成功；
输出 shape、dtype、device 和数值不变；
需要梯度的输出仍具有正确的 requires_grad 和 grad_fn；
backward() 成功；
输入及参数梯度存在且数值一致；
BatchNorm 的 running_mean、running_var 等状态一致；
Sanitizer 拦截算子的有序序列和各算子次数不变；
Sanitizer 仍能检测真实的跨 Stream 访问问题。

3. 方案设计

3.1 总体方案

采用“完整删除不可达兼容逻辑，不增加替代分支”的方案：

从 NPUSanitizerDispatchMode.__init__() 删除 npu_adjust_autograd 白名单。
删除 enable_autograd() 方法。
从 __torch_dispatch__() 删除 aten_api 临时变量和方法调用。
从 test/npu/test_stream_check.py 删除仅验证旧白名单、旧方法和旧字段的测试内容。
保持 Sanitizer 其余输入解析、算子执行、输出解析、Stream 获取和错误检查逻辑不变。

3.2 删除前后代码

删除前：

class NPUSanitizerDispatchMode(TorchDispatchMode):

    def __init__(self, event_handler: csan.EventHandler):
        super().__init__()
        self.event_handler = event_handler
        self.args_handler = None
        self.npu_adjust_autograd = [
            "adaptive_avg_pool2d", "batch_norm",
            "log_softmax", "nll_loss", "to"
        ]

    def enable_autograd(self, aten_api):
        if aten_api in self.npu_adjust_autograd:
            torch._C._dispatch_tls_set_dispatch_key_excluded(
                torch._C.DispatchKey.AutogradFunctionality,
                False,
            )

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        # ...
        self.args_handler = NPUArgumentHandler()
        aten_api = func.__name__.split(".")[0]
        self.enable_autograd(aten_api)
        self.parse_inputs(func._schema, args, kwargs, is_factory=is_factory)
        outputs = func(*args, **kwargs)
        # ...

删除后：

class NPUSanitizerDispatchMode(TorchDispatchMode):

    def __init__(self, event_handler: csan.EventHandler):
        super().__init__()
        self.event_handler = event_handler
        self.args_handler = None

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        # ...
        self.args_handler = NPUArgumentHandler()
        self.parse_inputs(func._schema, args, kwargs, is_factory=is_factory)
        outputs = func(*args, **kwargs)
        # ...

对应最小差异：

 class NPUSanitizerDispatchMode(TorchDispatchMode):
     def __init__(self, event_handler: csan.EventHandler):
         super().__init__()
         self.event_handler = event_handler
         self.args_handler = None
-        self.npu_adjust_autograd = [
-            "adaptive_avg_pool2d", "batch_norm",
-            "log_softmax", "nll_loss", "to"
-        ]
-
-    def enable_autograd(self, aten_api):
-        if aten_api in self.npu_adjust_autograd:
-            torch._C._dispatch_tls_set_dispatch_key_excluded(
-                torch._C.DispatchKey.AutogradFunctionality,
-                False,
-            )
 
     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         # ...
         self.args_handler = NPUArgumentHandler()
-        aten_api = func.__name__.split(".")[0]
-        self.enable_autograd(aten_api)
         self.parse_inputs(func._schema, args, kwargs, is_factory=is_factory)

3.3 方案选型

方案	优点	缺点或风险	结论
A. 完整删除旧兼容逻辑	改动最小；消除死代码、字符串热路径和私有接口依赖	不能自动覆盖仍可能需要 Patch 的旧版本	采用；当前 PR 面向 master / PyTorch 2.13，且已有同版本消融证据
B. 保留现状	不改变源码	保留不可达代码、私有接口和维护负担	不采用
C. 将白名单改成内部 ATen 名称	可让真分支重新可达	会重新引入未经必要性证明的 TLS 副作用，可能造成额外 Autograd Dispatch	不采用
D. 根据 PyTorch 版本保留分支	理论上可兼容新旧版本	尚无准确故障边界；增加条件分支和长期维护成本	本次不采用；旧分支应独立验证
E. 对所有算子重新启用 Autograd	无需维护白名单	影响面过大，可能导致重复分发、递归或性能问题	不采用

3.4 技术依据

3.4.1 TLS setter 的语义

PyTorch Dispatcher 的当前线程状态包含 included 和 excluded DispatchKeySet。算子最终可参与选择的 Key 可简化理解为：

(Tensor 自带的 Dispatch Keys | TLS included) - TLS excluded

历史代码将 AutogradFunctionality 的 excluded 状态设置为 False，即从排除集合中移除该功能键。该操作会改变当前线程后续 Dispatch 行为，因此只有在存在明确的实际故障和版本边界时才应保留。

3.4.2 当前真实路径不可达

历史实现用 func.__name__.split(".")[0] 得到内部算子名，再与公开 API 名称白名单精确比较。在 PyTorch 2.13 中，五类用例进入 Mode 的名称为 _adaptive_avg_pool2d、native_batch_norm、_log_softmax、nll_loss_forward、_to_copy 等，均不等于白名单项。

因此当前行为实际等价于：

def enable_autograd(aten_api):
    return None

删除该调用不会改变当前真实路径的 TLS 状态。

3.4.3 现有 Autograd 能力

当前 PyTorch / torch_npu 组合已通过正常 Dispatcher、代码生成的 Autograd wrapper、AutogradPrivateUse1 注册及对应 fallback 处理 NPU Autograd。五类算子在删除兼容代码后仍能建立计算图并完成反向传播，消融结果与该机制一致。

3.5 功能、性能和 DFX 设计

3.5.1 功能兼容性

不新增、删除或修改任何用户可见 API。
不改变 enable_npu_sanitizer() 的调用方式。
不改变 Sanitizer 报告格式和错误判定规则。
不改变五类算子的数学结果、梯度公式和设备实现。

3.5.2 性能

变更删除了每个普通 Dispatch 中的一次 func.__name__ 字符串切分、一次 Python 方法调用和一次五项列表成员判断。预期不会产生性能回退；本 RFC 不宣称可观的整体性能收益。

3.5.3 可靠性

删除无效 TLS 修改入口，降低未来算子命名变化意外触发历史副作用的风险。
不修改 Sanitizer 的核心读写集合、当前 Stream 获取及 EventHandler 逻辑。
通过基线 wheel 与删除 wheel 的独立环境对照，避免同一进程状态污染。

3.5.4 可维护性

移除未记录版本边界的白名单。
移除对 PyTorch 私有 setter 的依赖。
测试关注用户行为和 Sanitizer 能力，不继续绑定已删除的内部实现。

3.5.5 可测试性

五类历史算子均提供真实前向、反向路径验证。
记录 Mode 观察到的完整有序 Dispatch 序列和每算子次数。
记录 Sanitizer 启用前、启用后和用例结束后的 TLS 状态。
对输出、梯度和状态 Tensor 进行逐元素比较。

3.5.6 安全与隐私

本变更不新增网络访问、文件写入、遥测、日志字段或用户数据处理，不涉及新增安全和隐私风险。

3.6 影响文件

RFC 设计范围仅包含：

文件	修改内容
`torch_npu/npu/_stream_check.py`	删除白名单、`enable_autograd()`、临时变量和调用点
`test/npu/test_stream_check.py`	删除仅验证旧字段和旧函数的测试内容

异常日志格式、其他 import 清理、代码风格调整和其他 Sanitizer 模块均不属于本 RFC，若确有必要应另行提交，避免扩大 Patch 削减 PR 的审查范围。

3.7 接口与资料设计

接口变更：不涉及客户可见接口。
ABI 变更：不涉及。
资料变更：不涉及用户文档；需要在 PR/RFC 中记录内部实现删除原因和验证边界。
配置变更：不涉及。

4. 测试设计

4.1 验证环境

项目	配置
PyTorch	`2.13.0a0+gitfad7424`
torch_npu	master，实验固定提交 `19cd51d0eb647f9bdcc3d1624de20a9481ab8bc2`
torch_npu 版本	`2.13.0+git19cd51d`
设备	Ascend 910B2
基线	保留历史兼容代码的 wheel
对照	仅删除本 RFC 目标代码的 wheel

两个 wheel 从同一提交构建，分别安装到独立虚拟环境，并使用相同 PyTorch、设备和随机种子执行用例。

4.2 直接分支探针

目的：证明历史代码本身可以修改 TLS，而当前真实路径只是无法命中白名单，避免将“不可达”误写为“setter 失效”。

步骤：

预先将 AutogradFunctionality 的 excluded 状态设置为 True。
分别向 enable_autograd() 直接传入五个历史公开名称。
检查调用前后状态。

结果：五项均为 before_excluded=True、after_excluded=False，符合预期。

4.3 五类真实算子消融测试

用例	前向/反向检查	Dispatch 检查	删除前后结果
`adaptive_avg_pool2d`	输出、`grad_fn`、输入梯度	内部前向/反向算子，40 / 40	逐元素一致
`batch_norm`	输出、输入/weight/bias 梯度、running stats	103 / 103	逐元素一致，状态一致
`log_softmax`	输出、`grad_fn`、输入梯度	40 / 40	逐元素一致
`nll_loss`	loss、backward、输入梯度	37 / 37	逐元素一致
`to`	dtype 转换、`grad_fn`、梯度回传	43 / 43	逐元素一致，输出保持 float16

每个输出、梯度和状态 Tensor 检查：

shape；
dtype；
device；
是否全部为有限值；
展平后的逐元素值。

总比较结果：

success=true
errors=[]

额外一致性结果：

五个用例删除前后的有序 Dispatch 序列完全相同；
每个 operator 的 Dispatch 次数完全相同；
Sanitizer 启用前、启用后和用例结束后的 TLS 状态完全相同；
基线真实路径的 state_change_count 全部为 0；
两个 wheel 分别确认旧代码存在和不存在。

4.4 torch_npu 官方 Sanitizer 单元测试

测试文件	基线 wheel	删除 wheel	说明
`test/npu/test_stream_check.py`	4/4	3/3	删除 1 个仅测试旧函数的过期用例
`test/npu/test_sanitizer.py`	5/5	5/5	Sanitizer 启用和 trace callback 注册
`test/npu/test_sanitizer_record_stream.py`	18/18	18/18	`record_stream` 检测能力
`test/test_npu_sanitizer.py`	10/10	10/10	Sanitizer 综合路径
`test/test_sanitizer_pluggable_allocator.py`	1/1	1/1	可插拔分配器相关路径

基线共执行 38 个官方测试，删除版本共执行 37 个官方测试。数量差 1 是删除了直接调用旧 enable_autograd() 的实现绑定型测试，其余官方 Sanitizer 测试均通过。

4.5 PyTorch 社区参考用例

PyTorch 社区的 test/test_cuda_sanitizer.py 用于参考上游 Sanitizer 的测试组织和行为关注点。当前实验使用 CPU PyTorch 构建，torch.cuda.is_available() == False，该文件运行 0 个测试并记录为 SKIP_CUDA_UNAVAILABLE。

因此，该项只作为社区实现参考，不作为本次 NPU Patch 删除成功的证据。删除结论以 NPU 实机消融和 torch_npu 官方 Sanitizer 测试为准。

4.6 验收标准

只有同时满足以下条件才允许合入：

五个真实用例均进入待删函数，但历史白名单命中全部为 0。
直接探针能够识别旧 setter 的真实 TLS 副作用。
删除 wheel 中不存在白名单、enable_autograd() 和调用点。
五类用例删除前后均通过，输出、梯度和状态 Tensor 一致。
有序 Dispatch 序列和各算子 Dispatch 次数一致。
Sanitizer 启用前后及用例结束后的 TLS 状态一致。
torch_npu 官方 Sanitizer 测试除删除的过期实现测试外全部通过。
git diff --check、代码检查和相关 CI 通过。
PR 最终差异不夹带本 RFC 范围外的日志、import 或代码风格调整。

4.7 测试产物

聚焦测试：ablation/test_dispatch_key_reduction_213.py
自动构建与测试：ablation/run_master_213_focused_ablation.sh
删除补丁：ablation/remove_autograd_workaround.patch
官方 UT 更新补丁：ablation/update_stream_check_tests_after_removal.patch
消融报告：ablation/npu_sanitizer_dispatch_key_reduction_report_pt213_master.md
本地结果：results/master_213_focused_20260730_013025/
官方测试汇总：results/master_213_official_sanitizer_tests_20260730/summary_final.tsv
远端完整产物：/home/w50063847/sanitizer_master_213_focused_run3/

构建产物校验值：

baseline wheel SHA256:
291fb20a5c64cd64ebf39a9b6bab5b903bb36617558a2b3d4b060819487663ae

ablated wheel SHA256:
086d304ce788066a289c47d617307823a011cb0ce2cdcfee133d7d5b832827e4

5. 缺点和风险

风险	影响	缓解措施
旧 PyTorch 版本仍可能依赖历史逻辑	在旧维护分支删除后，个别算子可能无法正确建立 Autograd 图	本 RFC 和 PR 仅面向 master / PyTorch 2.13；不自动回移到旧分支，旧分支需独立消融
当前只覆盖五类历史白名单算子	其他未记录算子可能存在不同问题	保留官方 Sanitizer 综合测试；新问题应以独立复现和最小修复处理
算子内部名称未来再次变化	未来路径可能恰好命中或产生新差异	删除基于名称的副作用逻辑，并长期关注用户行为而非内部名称
当前未定位原始上游修复提交	难以严格给出“从哪个版本开始修复”	将结论限定为已验证组合，不将未知历史推断为事实
删除实现型 UT 后测试数量减少 1	审查者可能误认为覆盖回退	明确该测试只验证已删函数；行为覆盖由五算子消融和官方综合测试承担
PR 混入无关改动	扩大审查范围，难以证明风险边界	最终提交仅保留 RFC 影响文件内与兼容 Patch 删除直接相关的差异

5.1 回滚方案

若合入后发现与本删除直接相关的回归：

使用具体 PyTorch、torch_npu、CANN、设备和算子组合建立最小复现。
确认失败原因确实为 AutogradFunctionality 被排除，而不是其他 Dispatcher、kernel 或 Sanitizer 问题。
回滚本次删除提交，恢复白名单、函数和调用点以快速止损。
保留触发回归的新增测试。
后续修复必须限定真实受影响的版本和算子，记录删除条件；不得仅根据公开 API 名称恢复无边界白名单。

6. 现有技术

6.1 PyTorch `TorchDispatchMode`

PyTorch 使用 TorchDispatchMode 在 Python 层拦截 ATen 算子。Mode 内部再次调用 func(*args, **kwargs) 时，Dispatcher 会结合 Tensor Dispatch Keys 与当前线程 TLS included/excluded 状态选择后续 kernel。

本 RFC 不改变该机制，仅删除 torch_npu 在 Mode 中额外修改 Autograd TLS 的下游逻辑。

6.2 PyTorch CUDA Sanitizer

PyTorch 官方 CUDA Sanitizer 同样在 Python 层解析算子输入和输出，并结合 Stream、Event 与内存生命周期信息进行检查。官方实现未维护本 RFC 所述五项 Autograd 白名单，也未在 Sanitizer DispatchMode 中调用 _dispatch_tls_set_dispatch_key_excluded(AutogradFunctionality, False)。

这说明历史逻辑不是 Sanitizer 通用设计要求；是否保留应由 NPU 当前真实故障和验证结果决定。

6.3 torch_npu 当前实现

torch_npu 通过 PrivateUse1 / AutogradPrivateUse1 机制接入 PyTorch Dispatcher 和 Autograd。当前 master / PyTorch 2.13 组合中，五类历史算子无需 Sanitizer 内部手工修改 TLS 即可完成前向和反向。