☰
TorchAir Eager 模式算子数据 Dump 功能:NPU 图模式推理的输入输出取证与问题定位指南
2026/10/14 20:15:47 网站建设 项目流程
  • 人工智能
  • 模型编译
  • Ascend

【免费下载链接】torchair

TorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。

项目地址:https://gitcode.com/Ascend/torchair
点击查看免费下载

导读

本文介绍 TorchAir(torchair)在Eager 模式下对算子级输入/输出数据执行 dump的调试能力:通过torchair.get_npu_backend的compiler_config一键开启,将图计算过程中各算子的输入、输出张量以npy格式落盘,用于算子运行性能、精度类问题的定位与分析。读完本文你将掌握data_dump.type、data_dump.path、data_dump.filter三个配置项的使用方法、Node.name的获取技巧(借助图结构 dump)、以及 dump 产物文件命名规则与目录组织方式,并理解其底层基于torch.fx.Interpreter的实现原理。

功能简介

在 PyTorch 的torch.compile(model, backend=npu_backend)图编译流程中,TorchAir 先将模型转换为 FX 图,再进一步转换为昇腾 GE 图并执行。当模型出现算子运行性能异常或精度问题(如某个算子输出不符合预期)时,最直接的排查手段就是查看该算子在实际运行时的输入与输出数据。

Eager 模式下的算子数据 dump 功能,正是为这一场景设计:在 Eager 执行阶段,把图计算过程中每个(或经过 filter 筛选的)算子的输入、输出数据导出为npy文件,供后续离线比对与分析。需要特别说明的是,该功能会在 Eager 阶段完成数据采集,因此与“图模式(GE 图编译执行)下的功能配置”互斥,详见下文使用约束。

使用约束

开启本功能后,GE 图模式相关的功能配置均不生效。

从源码实现可以更清楚地理解这一点:在 npu_fx_compiler.py 中,当config.debug.data_dump.enabled为真时,编译器会跳过正常的 NPU 图编译流程,改用_NpuFxDumper(一个基于torch.fx.Interpreter的 FX 图解释器)以 Eager 方式逐节点执行并落盘数据,并打印如下告警日志:

When dumping data of FX Graph, npu run will be skipped, and FALLBACK to EAGER execution, once dump finished, please make sure to disable the data dump config to ensure that the graph is compiled and executed.

因此请牢记:dump 完成后务必关闭该配置,否则模型将一直以 Eager 回退方式执行,无法获得图编译带来的性能收益。同理,debug.fx_summary等其他调试项在开启时也会有类似的执行方式影响,使用调试功能时请注意区分。

使用方法

配置入口

该功能通过 torchair.get_npu_backend 中的compiler_config配置,CompilerConfig.debug.data_dump即数据 dump 配置对象。以下示例仅供参考,不支持直接拷贝运行(依赖具体模型与运行环境):

import torch, torch_npu, torchair config = torchair.CompilerConfig() # Eager模式下数据dump功能 config.debug.data_dump.type = "npy" config.debug.data_dump.path = "./test" # 若只dump 'add_1' 和 'sub' 算子,n.name为Node类实例化的name属性 config.debug.data_dump.filter = lambda n: n if n.name in ['add_1', 'sub'] else None npu_backend = torchair.get_npu_backend(compiler_config=config) opt_model = torch.compile(model, backend=npu_backend)

配置完成后,直接正常调用opt_model(...)即可触发 dump。

参数说明

表 1参数说明

| 参数名 | 参数说明 | | -- | -- | | data_dump.type | 指定 dump 文件类型,字符串类型。默认为None,表示不导出 dump 数据。若设置,当前仅支持npy格式。 | | data_dump.filter | 用户自定义过滤函数,保留满足函数条件的内容。输入:PyTorch 中的Node类的实例n;输出:PyTorch 中的Node类的实例n或者None。默认值为None,表示不过滤任何内容。说明:Node类实例化的各属性(如name、target等)获取方法主要通过图结构 dump 功能获取。以常见的name属性为例,获取方法见下文。 | | data_dump.path | 设置 dump 文件生成的路径,字符串型。可选配置,如果不设置,默认为当前执行路径。请确保参数中指定的路径真实存在,并且运行用户具有读取和写入权限。 |

获取 Node.name 属性(配合图结构 dump)

data_dump.filter回调接收的是torch.fx.Node实例,其name(节点名)与target(算子目标)等属性可以通过图结构 dump 功能获取,具体步骤如下:

  1. 先以py格式 dump 图信息:

    config.debug.graph_dump.type = "py"
  2. 在当前执行路径下会生成dynamo_*.py文件,其中包含 TorchAir 构图信息。搜索关键词FX Code,其后面字段即对应n.name属性信息。示例如下:

    # File "/home/a.py", line 32, in forward x=x+y ## FX Code: **add_1**: torch.float32[s0, s0]npu:0 = torch.ops.aten.add.Tensor(add: torch.float32[s0, s0]npu:0, arg2_1: torch.float32[s0, s0]npu:0) Add_1_0 = ge.Add(Add_0, arg2_1_0, node_name="Add_1") # File "/home/a.py", line 36, in forward x=x-1 ## FX Code: **sub**: torch.float32[s0, s0]npu:0 = torch.ops.aten.sub.Tensor(mul_1: torch.float32[s0, s0]npu:0, 1) Sub_0 = ge.Sub(Cast_1_0, ge.Const(1, dtype=0), node_name="Sub")

    上例中FX Code一行的**add_1**、**sub**即对应filter回调中n.name的取值,可直接用于筛选目标算子。图结构 dump 功能的完整介绍(graph_dump.type的py/txt/pbtxt三种格式与产物样例)参见 图结构 dump 功能。

配置项的源码级细节

从 debug_config.py 可以看到_DataDump类的定义要点:

  • type使用OptionValue(None, ["npy"])定义,默认值为None,可选值当前仅["npy"];若传入不在可选列表中的值,配置框架会直接抛出ValueError(见 _option_base.py);
  • filter默认为None,作为固定属性(_fixed_attrs)参与配置比较与序列化;
  • path通过setter执行os.path.realpath(value)做真实路径归一化,因此配置相对路径时也会被解析为绝对路径后再拼接产物路径;
  • _DebugBase.enabled属性判定type.value is not None,即只要设置了type即视为开启本功能(见 debug_config.py)。

CompilerConfig.debug在 compiler_config.py 中实例化为_DebugConfig(),其中同时包含graph_dump、data_dump、fx_summary、run_eagerly、deadlock_check等调试子项,本文聚焦data_dump。

功能原理:基于 torch.fx.Interpreter 的 Eager 数据采集

理解产物结构与过滤行为,有必要看一下底层实现 fx_dumper.py:

  • _NpuFxDumper继承自torch.fx.Interpreter(fx_dumper.py),以解释执行方式逐个运行 FX 图节点;
  • run_node中先调用用户filter:若返回None则跳过 dump、正常执行该节点;否则以过滤后的节点进入 dump 上下文(fx_dumper.py);
  • _trace_dump装饰器(fx_dumper.py)在节点执行前后分别落盘输入与输出:call_function执行前保存每个输入参数(含 list/tuple 形态的序列化输入,按外层索引i与内层索引j展开),执行后保存每个输出(OUTPUT.{i});
  • 张量统一通过_as_numpy转为numpy.ndarray保存,其中bfloat16张量会先转成float32再落盘(fx_dumper.py),比对数据时需注意 dtype 差异;
  • 每次run(即对编译后模型的一次调用)结束后_step += 1,作为 dump 目录中的 step 标识,从 0 开始递增(fx_dumper.py);
  • 分布式场景下,若torch.distributed已初始化且world_size > 1,目录会按worldsize{world_size}_global_rank{global_rank}分卡组织(fx_dumper.py)。

另外,非 ATen 目标(无对应 GE converter 的节点,如部分自定义 target)会被跳过不参与 dump,日志中会打印Skip dump non aten target级别的调试信息(fx_dumper.py),即data_dump主要覆盖可被 GE 转换的 ATen 算子。

产物说明

开启数据 dump 功能后,得到如下命名的文件:

${op_type}-${aten_ir}.${param_type}${param_idx}${timestamp}.npy

各字段含义:

  • ${op_type}:算子类型,即 FX 节点的name属性(如add_1、sub);
  • ${aten_ir}:ATen 算子名,即节点的target(如aten.add.Tensor、aten.sub.Tensor);
  • ${param_type}:参数类型,INPUT或OUTPUT;
  • ${param_idx}:输入/输出参数的索引,默认从 0 开始;
  • ${timestamp}:时间戳;
  • ${world_size}/${global_rank}:集合通信中的world_size以及global_rank信息(仅分布式多卡场景体现于目录名)。

产物路径默认为当前执行目录,目录结构如下:

├ torch_test.py // PyTorch脚本 │ ├── worldsize${world_size}_global_rank${global_rank} // 若只涉及单卡则表示为"worldsize1_global_rank0" │ ├── graph_1 // ${global_rank}所在进程的子图名称 │ ├── 0 // 首次dump的结果文件,通过step id标记,step id从0开始 │ ├── add_1-aten.add.Tensor.INPUT.0.20251017071211190258.npy // filter过滤函数dump的结果文件 │ ├── add_1-aten.add.Tensor.INPUT.1.20251017071211190258.npy │ ├── add_1-aten.add.Tensor.OUTPUT.0.20251017071211190258.npy │ ├── sub-aten.sub.Tensor.INPUT.0.20251017071211202083.npy │ ├── sub-aten.sub.Tensor.INPUT.1.20251017071211202083.npy │ ├── sub-aten.sub.Tensor.OUTPUT.0.20251017071211202083.npy │ ...... │ ├── n // 第n+1次dump的结果文件,通过step id标记 │ ├── xxx.npy │ ...... │ ├── worldsize${world_size}_global_rank${global_rank} │ ├── graph_1 │ ├── 0 │ ├── xxx.npy

上述组织方式与源码实现一一对应:worldsize${world_size}_global_rank${global_rank}/graph_${id}/${step}/的路径拼接逻辑见 fx_dumper.py,其中graph_${id}中的${id}来自npu_fx_compiler.py中的全局唯一图序号(_next_unique_graph_id()),${step}即每次调用模型产生的 step id。

实操验证与测试参考

仓库内置了数据 dump 的自动化测试,可作为功能开启方式与产物校验的参考:

  • tests/smoke/data_dump_test.py:test_data_dump_with_scope演示了通过torchair.scope.data_dump()上下文对局部算子(add、sub)开启 dump,并断言产物目录worldsize*_global_rank*/0/graph_*/1/0下同时存在 add 与 sub 相关文件;test_aclgraph_data_dump演示了mode = "reduce-overhead"(aclgraph 模式)下基于dump_config.enable_dump的 dump 流程;
  • tests/st/torchair_st.py:test_data_dump_generation直接设置config.debug.data_dump.type = "npy"、config.debug.data_dump.path = "./test_data_dump_generation/dir",随后断言worldsize1_global_rank0目录下生成 3 个 npy 文件(对应某算子的 2 个输入 + 1 个输出);test_data_dump_failpath则验证了当path指向一个已存在的普通文件(非目录)时的失败行为——因此务必保证data_dump.path指向真实存在的可写目录。

典型使用流程总结

  1. 开启图结构 dump(config.debug.graph_dump.type = "py"),确认模型执行,从生成的dynamo_*.py中FX Code关键词处摘录目标算子的name;
  2. 关闭图结构 dump,开启数据 dump:设置config.debug.data_dump.type = "npy"、config.debug.data_dump.path = "<存在的可写目录>",并按需编写data_dump.filter回调只保留目标算子;
  3. 重新执行模型,从产物路径(默认为当前执行目录)下的worldsize1_global_rank0/graph_*/0/中读取对应算子的INPUT/OUTPUTnpy 文件,与期望值或相邻 step 的数据进行比对分析;
  4. 问题定位完成后,务必关闭data_dump配置(将type恢复为None),使模型恢复正常的 GE 图编译执行,避免持续的 Eager 回退带来性能开销。

通过上述流程,即可在 Eager 模式下对任意关注算子的输入输出完成“取证”,为算子性能瓶颈分析与精度问题排查提供可靠的数据基础。

  • 人工智能
  • 模型编译
  • Ascend

【免费下载链接】torchair

TorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。

项目地址:https://gitcode.com/Ascend/torchair
点击查看免费下载

相关推荐

上一篇:Phi-3-medium-128k-instruct部署指南:从云端到边缘计算的完整解决方案
下一篇:大恒相机控制源码+C++/C+显示/控制/存储

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询