简介

概述

TorchAir（Torch Ascend Intermediate Representation）是昇腾为Ascend Extension for PyTorch（torch_npu）提供的图模式能力扩展库，支持用户使用PyTorch和torch_npu在昇腾设备上进行图模式的训练和推理。TorchAir提供了昇腾设备的图模式编译后端，对接PyTorch的Dynamo特性，将PyTorch的FX（Functionalization）计算图转换为昇腾的GE（Graph Engine，图引擎）计算图，并提供在昇腾设备上启动GE图编译和执行的能力。

TorchAir继承了Dynamo的大部分特性，如动态shape图功能，在PyTorch基础上又新增了图相关能力，如离线场景导图、模型编译时间优化、集合通信算子入图等功能，详细介绍参见功能介绍。

在当前版本中，TorchAir仅作为一项试用特性，此功能在后续版本中可能会有所调整或改进。请用户在使用过程中关注后续版本的迭代。
对PyTorch模型使能图模式前，请先熟悉如下手册，了解基本概念、模型迁移等基础知识。
- 参考《产品概述》，了解Ascend Extension for PyTorch的概念和作用。
- 参考《PyTorch 训练模型迁移调优指南》的“模型脚本迁移”章节，了解如何将模型迁移至昇腾NPU上等。
Eager（Eager Execution，即时执行）模式与图模式区别：前者是每个计算操作下发后立即执行，后者是所有计算操作构造成一张图，以图的粒度下发执行。

使用约束

对于以下单算子场景，暂不支持图模式：

单进程多线程多卡并发执行场景
单进程单线程多卡执行场景

使用说明

首次阅读本文档时，建议先阅读表1，帮助您快速获取TorchAir安装、功能使用、API参考等指导。

表1 使用向导
使用场景	操作指引	说明
1. TorchAir使用前如何准备环境	安装	阐述了TorchAir包的获取、安装、版本约束以及依赖的三方库等信息。
2. TorchAir图模式如何快速配置	快速上手	提供了TorchAir如何使能图模式的代码示例，同时介绍了图模式下支持配置的config功能。
3. TorchAir提供了哪些关键能力	模型缓存编译功能	提供了模型缓存编译方案，通过优化dynamo和GE编译耗时，加速torch.compile图模式启动时间。
	集合通信算子入图	针对PyTorch 2.1中不支持入图的集合通信算子提供补丁函数，实现部分集合通信算子入图。
	DFX功能	提供了模型训练/推理时可能会用到的一系列问题分析和定位功能，具体如下：日志功能图dump功能算子信息导出功能图模式下数据dump功能 Eager模式下数据dump功能图内Tensor打印功能
	更多功能	提供了能进一步提升模型执行性能的功能。此类功能适用于某些特定场景，用户需要根据实际情况配置对应功能项。 dynamo导图功能动态shape图分档执行功能图单流执行功能图编译多级优化选项算子融合功能固定权重类输入地址功能图模式编译节点遍历选项计算与通信并行功能算子在线编译选项使能RefData类型转换功能 Tiling调度优化功能 View类算子优化功能动静子图拆分场景性能优化
4. 自定义算子如何通过TorchAir入图	自定义算子插件化入图	针对自定义算子，如何在PyTorch框架中以图模式调用，TorchAir提供了一种插件化注册方式。
5. TorchAir支持哪些Aten API	支持的Aten API清单	提供了当前支持使能图模式的PyTorch原生算子Aten API。
6. 图模式下大模型推理和迁移案例	大模型图模式推理案例	LLaMA2模型以图模式在昇腾NPU进行Deepspeed分布式推理的样例。 LLaMA2模型分离部署迁移样例。
7. 更多功能配置相关的API	API参考	提供了基础API、GE相关API、推理常用API以及LLM-DataDist（模型分布式集群和数据加速组件）相关API介绍。
8. 图模式下的常见问题有哪些	FAQ	罗列了一些常见问题，提供问题分析和处理方法建议，供用户参考。