简介

概述

TorchAir（Torch Ascend Intermediate Representation）是昇腾Ascend Extension for PyTorch（torch_npu）的图模式能力扩展库，提供了昇腾设备亲和的torch.compile图模式后端，实现了PyTorch网络在昇腾NPU上的图模式推理加速以及性能优化。

当前版本中，TorchAir为beta特性，主要关注推理场景模型优化。
PyTorch网络使能图模式前，建议先了解Ascend Extension for PyTorch相关知识：
- 参考《Ascend Extension for PyTorch 快速入门》，了解Ascend Extension for PyTorch的概念和作用。
- 参考《PyTorch 训练模型迁移调优指南》的“模型脚本迁移”章节，了解如何将模型迁移至昇腾NPU上等。

TorchAir在Ascend Extension for PyTorch中的架构如图1所示，提供两种工作模式（通过config.mode切换）：

reduce-overhead模式（aclgraph）：采用Capture&Replay方式实现任务一次捕获多次执行，Capture阶段捕获Stream任务到Device侧，暂不执行；Replay阶段从Host侧发出执行指令，Device侧再执行已捕获的任务，从而减少Host调度开销，提升性能。

reduce-overhead模式（aclgraph）又称为捕获模式，该模式通过Runtime提供的aclmdlRICaptureXxx系列接口实现，其原理和接口介绍请参考《CANN 应用开发指南 (C&C++)》中“运行时管理>基于捕获方式构建模型运行实例”章节。
max-autotune模式（Ascend IR）：将PyTorch的FX计算图转换为昇腾中间表示（IR，Intermediate Representation），即Ascend IR计算图，并通过GE（Graph Engine，图引擎）实现计算图的编译和执行。

为了实现图模式下模型推理加速和优化，TorchAir不仅继承了大部分PyTorch原生Dynamo特性（如动态shape图功能等），还在此基础上新增了其他图优化和图定位能力，如离线场景下导图、集合通信算子入图等，详细介绍请参见基础功能、max-autotune模式功能、reduce-overhead模式功能。

图1 TorchAir逻辑架构图

使用约束

PyTorch图模式支持单进程和多进程，每个进程只支持使用1张NPU卡，不支持使用多张NPU卡。
目前仅如下产品支持aclgraph（reduce-overhead模式）执行模式：
- Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件
- Atlas A3 训练系列产品/Atlas A3 推理系列产品

常用概念

名称	说明
Eager模式	PyTorch支持的单算子执行模式（未使用torch.compile），特点如下，单击Link可获取PyTorch官网介绍。即时执行：每个计算操作在定义后立即执行，无需构建计算图。动态计算图：每次运行可能生成不同的计算图。
图模式	一般指使用torch.compile加速的模型执行方式。图模式执行特点如下：延迟执行：所有计算操作先构成一张计算图，再在会话中下发执行。静态计算图：计算图在运行前固定。
TorchAir图模式	PyTorch图模式的一种实现，指使用torch.compile，同时指定TorchAir为其backend的执行方式。
ATen	全称为A Tensor Library，是PyTorch张量计算的底层核心函数库，这些函数通常称为ATen算子，负责所有张量操作（如加减乘除、矩阵运算、索引等）的C++实现。详细介绍可单击Link获取PyTorch官网详情。
FX图	Functionality Graph，是PyTorch中用于表示模型计算流程的中间层数据结构。通过符号化追踪代码生成计算图，将Python代码转为中间表示（IR，Intermediate Representation），实现计算图的动态调整和优化（如量化、剪枝等），详细介绍可单击Link获取torch.fx详情。
GE	Graph Engine，图引擎。它是计算图编译和运行的控制中心，提供图优化、图编译管理以及图执行控制等功能。 GE通过统一的图开发接口提供多种AI框架的支持，不同AI框架的计算图可以实现到Ascend IR图的转换，详细介绍可单击Link获取详情。
Pass	在深度学习框架（如PyTorch）和编译器（如TVM）中，Compiler Passes（编译器传递）和Partitioners（分区器）是优化图执行的关键技术，用于性能优化、硬件适配和计算图转换等，而Pass则是指在这些计算图上执行的特定变换操作。常见的Pass操作包括常量折叠、算子融合、内存优化等，详细介绍可单击Link获取PyTorch官网详情。 FX Pass是指对计算图（torch.fx.Graph）进行遍历、分析和转换等一系列操作，类似于传统编译器中的优化步骤（如常量折叠、算子融合）。
In-place算子	原地算子，该类算子可直接修改输入数据，不创建新的存储空间。从而节省内存，避免复制数据的开销。
Out-of-place算子	非原地算子，又称“非In-place算子”，该类算子保持原始输入数据不变，会创建并返回新对象，带来额外存储开销。
算子Schema	在PyTorch中，算子Schema（Operator Schema）定义了算子的输入、输出、属性以及行为规范，确保算子在正向传播（Forward）和反向传播（Backward）时能正确执行。PyTorch使用Schema来注册算子，并在编译或运行时进行验证。算子Schema主要通过修改native_functions.yaml文件实现，该文件位于PyTorch源码的aten/src/ATen目录下，用于声明算子的名称、参数类型、返回值类型及设备端实现函数。原生算子（Native Operators）：PyTorch内置算子（如 torch.add、torch.matmul），其Schema在C++中定义。通常在 native_functions.yaml 或通过 TORCH_LIBRARY 宏定义。例如，torch.add 的 Schema 可能如下： - func: add(Tensor self, Tensor other, , Scalar alpha=1) -> Tensor // 算子名，定义输入/输出类型 python_module: nn // Python模块名（如torch.nn） variants: function, method // 支持的类型（function、method） dispatch: // 指定后端实现 CPU: add_cpu 自定义算子（Custom Operators）：通过torch.library或torch.ops注册的算子，可以在Python或C++中定义Schema。 import torch.library # 定义Schema torch.library.define( "mylib::myop", # 算子名（格式：<namespace>::<op>） "(Tensor x, float scale) -> Tensor" # 算子输入 ) # 实现算子 @torch.library.impl("mylib::myop", "cpu") def myop_cpu(x, scale): return x scale
Ascend C	CANN针对算子开发场景推出的编程语言，原生支持C和C++标准规范，兼具开发效率和运行性能。基于Ascend C编写的算子程序，通过编译器编译和运行时调度，运行在昇腾AI处理器上。开发的算子简称为Ascend C算子，其调用方式一般为aclnnXxx的C接口形式，具体介绍请参考《CANN Ascend C算子开发指南》。
OpPlugin	Ascend Extension for Pytorch（torch_npu）算子插件，为使用PyTorch框架的开发者提供便捷的NPU算子库调用能力，具体介绍参考Ascend/OpPlugin仓，而算子适配开发过程参考《PyTorch 框架特性指南》中的“基于OpPlugin算子适配开发”章节。

使用说明

首次阅读本文档时，建议先阅读表1，以帮助您快速获取TorchAir安装方法、快速上手示例、关键功能项说明、API参考等。

表1 使用向导
使用场景	操作指引	说明
1. TorchAir使用前如何准备环境	安装	阐述了TorchAir包的获取、安装、版本约束以及依赖的三方库等信息。
2. TorchAir图模式如何快速配置	快速上手	提供了TorchAir配置图模式的示例，并介绍图模式相关的config配置。
3. TorchAir提供了哪些图功能	基础功能	提供了简单的基础功能，不区分max-autotune或reduce-overhead图执行模式。日志功能集合通信入图图结构dump功能算子输入输出dump功能（Eager模式） run-eagerly功能
	reduce-overhead模式功能	提供了reduce-overhead模式（aclgraph）下支持配置的功能。 reduce-overhead模式配置 FX Pass配置功能模型编译缓存功能图内多流表达功能 aclgraph间内存复用功能静态Kernel编译配置
	max-autotune模式功能	提供了max-autotune模式（Ascend IR）下支持配置的功能。 Dynamo导图功能模型编译缓存功能图内Tensor打印功能图内多流表达功能算子Converter支持度导出功能算子输入输出dump功能（图模式）动态shape图分档执行功能图编译统计信息导出功能图单流执行功能图编译多级优化选项算子融合规则配置功能（fusion_switch_file）算子融合规则配置功能（optimization_switch）固定权重类输入地址功能图模式编译节点遍历选项计算与通信并行功能算子在线编译选项 RefData类型转换功能 Tiling调度优化功能 View类算子优化功能动静子图拆分场景性能优化图内标定SuperKernel范围图内设置AI Core和Vector Core核数
4. 支持入图的ATen API	支持的ATen API清单	提供了能在昇腾NPU上以图模式执行的Torch API列表。
5. 自定义算子如何入图	自定义算子入图	针对自定义算子，TorchAir提供了插件化入图方式。
6. TorchAir提供了哪些API	API参考	提供了图模式下推理场景常用的API以及LLM-DataDist（模型分布式集群和数据加速组件）相关API介绍。
7. 常用定位方法和案例	常见定位方法和案例	提供了LLaMA模型以图模式在昇腾NPU进行推理的案例。提供了入图问题的定位流程、精度比对和性能分析的方法及案例。 FAQ：罗列使用TorchAir过程中可能遇到的系统或环境问题，提供了分析和处理建议。