简介

概述

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

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

TorchAir在Ascend Extension for PyTorch中的架构如图1所示，其实现图模式的方式有如下两种，可通过config.mode切换模式。

max-autotune模式（Ascend IR）：将PyTorch的FX计算图转换为昇腾中间表示（IR，Intermediate Representation），即Ascend IR计算图，并通过GE（Graph Engine，图引擎）实现计算图的编译和执行。
reduce-overhead模式（aclgraph）：采用Capture&Replay方式实现任务一次捕获多次执行，Capture阶段捕获Stream任务到Device侧，暂不执行；Replay阶段从Host侧发出执行指令，Device侧再执行已捕获的任务，从而减少Host调度开销，提升性能。

reduce-overhead模式（aclgraph）又称为捕获模式，该模式通过Runtime提供的aclmdlRICaptureXxx系列接口实现，其原理和接口介绍请参考《CANN ACL应用开发指南 (C&C++)》中“运行时资源管理>基于捕获方式构建模型运行实例”章节。

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

图1 TorchAir逻辑架构图

使用约束

PyTorch图模式功能当前仅支持单进程单卡场景。
目前仅如下产品支持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算子	非原地算子，该类算子保持原始输入数据不变，会创建并返回新对象，带来额外存储开销。

使用说明

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

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