torchair:基于 PyTorch 生态的昇腾NPU图模式推理扩展库项目

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

分支29Tags0
文件最后提交记录最后更新时间
20 天前
16 天前
2 个月前
2 个月前
15 天前
3 个月前
16 天前
23 小时前
1 天前
1 小时前
1 个月前
1 天前
29 天前
8 个月前
3 天前
3 天前
1 个月前
4 个月前
15 天前
3 个月前
11 天前
1 个月前
29 天前
26 天前
19 天前
23 小时前
3 个月前
8 天前
16 天前
3 年前
1 个月前
2 个月前
22 天前
3 个月前
3 年前
3 个月前
3 个月前
1 个月前
4 个月前

简介

本项目开发了名为TorchAir(Torch Ascend Intermediate Representation)的扩展库,支持用户基于PyTorch框架和TorchNPU插件在昇腾NPU上使用图模式进行推理。 TorchAir继承自PyTorch框架Dynamo模式,将PyTorch的FX图转换为GE计算图,并提供了GE计算图在昇腾NPU的编译与执行的能力。

  • 如果您想了解如何使用TorchAir,可以优先通过访问TorchAir图模式使用指南,了解关于TorchAir的更多信息。
  • 如果您计划为TorchAir贡献代码,请参考README内容。

安装与卸载

1.安装

  • 安装依赖

在安装TorchAir之前,请参考版本配套表支持型号说明,确保您的硬件能够使用TorchAir,并安装最新昇腾软件栈。 建议:执行ut和简单用例请使用docker/Dockerfile中推荐的版本。

  • 编译准备

克隆TorchAir代码仓

git clone https://gitcode.com/ascend/torchair.git

下载依赖三方库

cd torchair
git submodule update --init --recursive

配置编译环境,执行配置命令

配置只需要进行一次,用于获取pytorch的编译选项(如当前的torch是否开启ABI)及Ascend sdk的安装路径(如果需要在本地CPU上进行调试)。

cd ./torchair
bash ./configure

默认情况下,执行上述命会弹出如下的交互式会话窗口

您的会话可能有所不同。

Please specify the location of python with available torch 2.1.x installed. [Default is /usr/bin/python3]
(You can make this quiet by set env [TARGET_PYTHON_PATH]):

此时,要求您输入安装了 Torch 2.1 版本的python解释器路径,如果默认路径是正确的,直接回车,否则请输入正确的 python 解释器路径。

您可以通过设置 TARGET_PYTHON_PATH 环境变量,来抑制交互式窗口弹出,但是要确保路径是有效的,否则,仍然会要求您输入正确的 python 解释器路径。

键入后,会耗费几秒钟以确保您的输入是有效的,接着,会弹出下面的交互式窗口

Specify the location of ascend sdk for debug on localhost or leave empty.
(You can make this quiet by set env [ASCEND_SDK_PATH]):

如果您不需要在本地CPU模拟执行,可以直接回车跳过。否则,需要输入昇腾处理器开发套件的安装路径(默认安装路径通常为/usr/local/Ascend/ascend-toolkit/latest/)。

您可以通过设置 ASCEND_SDK_PATH 环境变量指定 Ascend CANN SDK 目录,你后续需要执行UT或者简单用例则需要提前设置,仅编译构建可设置 NO_ASCEND_SDK 环境变量指定不需要 SDK 来抑制交互式窗口弹出。

键入后,等待配置完成。

  • 编译

执行以下命令,编译生成TorchAir安装包:

mkdir build
cd build
cmake ..
make torchair -j8
  • 安装

编译完成后,会在build/dist/dist/目录下生成名为torchair-{version}-py3-none-any.whl的安装包文件。

您可以直接使用pip安装该安装包,或者使用make命令安装至您configure时指定的python环境中。

make install_torchair

2.卸载

torchair的卸载只需要执行命令:

pip3 uninstall torchair

如需要保存卸载日志,可在pip3 uninstall命令后面加上参数--log <PATH>,并对您指定的目录<PATH>做好权限管控。

快速上手

本的调试(最小CPU调试):

1、在执行./configure设置ASCEND_SDK_PATH如(ASCEND_SDK_PATH=xx/Ascend/ascend-toolkit/latest/)。

2、执行环境设置

source tools/env.sh

npu设备调试:

1、安装torch_npu

torch_npu版本可以根据环境要求安装,请参照[版本配套表]。

pip3 install numpy
pip3 install pyyaml
pip3 install setuptools
pip3 install torch_npu==2.12 # 通过pip安装torch_npu时会自动安装依赖的torch版本

如果安装了真实torch_npu,需要把torch_npu的stubs模块链接环境去除,否则无法正常使用。 如LD_LIBRARY_PATH中去除/workspace/build/stubs设置,如果没有设置请忽略。

2、安装CANN主线版本

npu调试依赖CANN主线版本。如果您是外部用户,可以联系对应的支撑团队获取安装包。

执行前,需要执行 CANN 安装目录下的 set_env.sh 设置 CANN 相关环境变量

例如CANN安装路径为 /usr/local/Ascend/cann,则执行:

source /usr/local/Ascend/cann/set_env.sh

执行以下python脚本快速验证TorchAir基本功能

import torch
import torchair

config = torchair.CompilerConfig()
npu_backend = torchair.get_npu_backend(compiler_config=config)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        return torch.add(x, y)

model = Model()
model = torch.compile(model, backend=npu_backend, dynamic=False)
x = torch.randn(2, 2)
y = torch.randn(2, 2)
model(x, y)

特性介绍

TorchAir常用特性如下,全量特性介绍请访问docs获取详细教程。

特性功能 功能介绍
日志功能 日志功能
graph dump功能 图dump功能
dynamo export功能 air格式图导出功能
data dump功能 精度数据dump功能
graph fusion功能 用户自定义关闭/开启部分融合算子功能
experimental性能提升功能 试验性质功能,不同功能适用于特定场景,详见参考资料
converter功能拓展 用户自行扩展模型中缺失的converter功能
支持的aten API清单 支持的aten API清单

版本配套表

本项目与PyTorch、TorchNPU、CANN以及Python的版本匹配详情如下:

TorchAir版本 PyTorch版本 TorchNPU版本 CANN版本 Python版本
master(主线) 2.7.1 + 在研版本 在研版本 Python3.10.x
v26.1.0 2.7.1 + v26.1.0 9.1.0 Python3.10.x / Python3.11.x
v26.0.0 2.7.1 + v26.0.0 9.0.0 Python3.10.x / Python3.11.x
7.3.0 2.6.0 / 2.7.1 / 2.8.0 7.3.0 8.5.0 Python3.9.x / Python3.10.x / Python3.11.x
7.3.0 2.9.0 7.3.0 8.5.0 Python3.9.x / Python3.10.x / Python3.11.x / Python3.12.x
7.2.0 2.1.0 7.2.0 8.3.RC1 Python3.8.x / Python3.9.x / Python3.10.x / Python3.11.x
7.2.0 2.6.0 / 2.7.1 / 2.8.0 7.2.0 8.3.RC1 Python3.9.x / Python3.10.x / Python3.11.x
7.1.0 2.1.0 7.1.0 8.2.RC1 Python3.8.x / Python3.9.x / Python3.10.x / Python3.11.x
7.1.0 2.5.1 / 2.6.0 7.1.0 8.2.RC1 Python3.9.x / Python3.10.x / Python3.11.x
7.0.0 2.1.0 / 2.3.1 / 2.4.0 7.0.0 8.1.RC1 Python3.8.x / Python3.9.x / Python3.10.x / Python3.11.x
7.0.0 2.5.1 7.0.0 8.1.RC1 Python3.9.x / Python3.10.x / Python3.11.x

torch_npu的版本选择请参照[版本配套表],下载地址:TorchNPUReleases

示例: https://gitcode.com/Ascend/pytorch/releases/download/v26.1.0-pytorch2.10.0/torch_npu-2.10.0.post4-cp310-cp310-manylinux_2_28_x86_64.whl

支持的型号

  • Ascend 950PR/Ascend 950DT
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品

贡献

如果您计划为TorchAir做出贡献,请参考CONTRIBUTING

联系我们

安全声明

TorchAir安全声明参考SECURITY_README文件。

许可证

TorchAir插件使用BSD许可证。详见LICENSE文件。


免责声明

致TorchAir使用者

  1. TorchAir提供的模型仅供您用于非商业目的。
  2. 对于各模型,TorchAir平台仅提示性地向您建议可用于训练的数据集,华为不提供任何数据集,如您使用这些数据集进行训练,请您特别注意应遵守对应数据集的License,如您因使用数据集而产生侵权纠纷,华为不承担任何责任。
  3. 如您在使用TorchAir模型过程中,发现任何问题(包括但不限于功能问题、合规问题),请在GitCode提交issue,我们将及时审视并解决。

致数据集所有者

如果您不希望您的数据集在TorchAir中的模型被提及,或希望更新TorchAir中的模型关于您的数据集的描述,请在GitCode提交issue,我们将根据您的issue要求删除或更新您的数据集描述。衷心感谢您对TorchAir的理解和贡献。

项目介绍

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

定制我的领域