Ascend / pytorch

Ascend Extension for PyTorch Overview This repository develops the **Ascend Extension for PyTorch** named **torch_npu** to adapt **Ascend NPU** to **PyTorch** so that developers who use the **PyTorch** can obtain powerful compute capabilities of **Ascend AI Processors**. Ascend is a full-stack AI computing infrastructure for industry applications and services based on Huawei Ascend processors and software. For more information about Ascend, see Ascend Community. Installation From Binary Provide users with wheel package to quickly install **torch_npu**. Before installing **torch_npu**, complete the installation of **CANN** according to Ascend Auxiliary Software. To obtain the **CANN** installation package, refer to the CANN Installation. • **Install PyTorch** Install **PyTorch** through pip. **For Aarch64:** **For x86:** • **Install torch-npu dependencies** Run the following command to install dependencies. If the installation fails, use the download link or visit the PyTorch official website to download the installation package of the corresponding version. | OS arch | Python version | link | |---------|----------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | x86 | Python3.8 | link | | x86 | Python3.9 | link | | x86 | Python3.10 | link | | x86 | Python3.11 | link | | aarch64 | Python3.8 | link | | aarch64 | Python3.9 | link | | aarch64 | Python3.10 | link | | aarch64 | Python3.11 | link | • **Install torch-npu** From Source In some special scenarios, users may need to compile **torch-npu** by themselves.Select a branch in table Ascend Auxiliary Software and a Python version in table PyTorch and Python Version Matching Table first. The docker image is recommended for compiling torch-npu through the following steps(It is recommended to mount the working path only and avoid the system path to reduce security risks.), the generated .whl file path is ./dist/. Note that gcc version has the following constraints if you try to compile without using docker image: we recommend to use gcc 10.2 for ARM and gcc 9.3.1 for X86. • **Clone torch-npu** • **Build Docker Image** • **Enter Docker Container** • **Compile torch-npu** Take **Python 3.8** as an example. **Tips** If you would like to compile with new C++ ABI, then first run this command, at this point, the recommended compilation environment is same to community torch package: glibc 2.28, gcc 11.2.1 Meanwhile, we support configuring -fabi-version using the following variables，require consistency with the community torch package Getting Started Prerequisites Initialize **CANN** environment variable by running the command as shown below. Quick Verification You can quickly experience **Ascend NPU** by the following simple examples. User Manual Refer to API of Ascend Extension for PyTorch for more detailed information. PyTorch and Python Version Matching Table | PyTorch Version | Python Version | |-----------------|:----------------------------------------------------------| | PyTorch1.11.0 | Python3.7.x(>=3.7.5),Python3.8.x,Python3.9.x,Python3.10.x | | PyTorch2.1.0 | Python3.8.x,Python3.9.x,Python3.10.x,Python3.11.x | | PyTorch2.2.0 | Python3.8.x,Python3.9.x,Python3.10.x | | PyTorch2.3.1 | Python3.8.x,Python3.9.x,Python3.10.x,Python3.11.x | | PyTorch2.4.0 | Python3.8.x,Python3.9.x,Python3.10.x,Python3.11.x | | PyTorch2.5.1 | Python3.9.x,Python3.10.x,Python3.11.x | | PyTorch2.6.0 | Python3.9.x,Python3.10.x,Python3.11.x | | PyTorch2.7.1 | Python3.9.x,Python3.10.x,Python3.11.x | | PyTorch2.8.0 | Python3.9.x,Python3.10.x,Python3.11.x | Ascend Auxiliary Software **PyTorch Extension** versions follow the naming convention , where the former represents the PyTorch version compatible with the **PyTorch Extension**, and the latter is used to match the CANN version. The detailed matching is as follows: | CANN Version | Supported PyTorch Version | Supported Extension Version | Github Branch | |-----------------------|---------------------------|-----------------------------|-------------------| | CANN 8.5.0 | 2.9.0 | 2.9.0 | v2.9.0-7.3.0 | | | 2.8.0 | 2.8.0.post2 | v2.8.0-7.3.0 | | | 2.7.1 | 2.7.1.post2 | v2.7.1-7.3.0 | | | 2.6.0 | 2.6.0.post5 | v2.6.0-7.3.0 | | CANN 8.3.RC1 | 2.8.0 | 2.8.0 | v2.8.0-7.2.0 | | | 2.7.1 | 2.7.1 | v2.7.1-7.2.0 | | | 2.6.0 | 2.6.0.post3 | v2.6.0-7.2.0 | | | 2.1.0 | 2.1.0.post17 | v2.1.0-7.2.0 | | CANN 8.2.RC1 | 2.6.0 | 2.6.0 | v2.6.0-7.1.0 | | | 2.5.1 | 2.5.1.post1 | v2.5.1-7.1.0 | | | 2.1.0 | 2.1.0.post13 | v2.1.0-7.1.0 | | CANN 8.1.RC1 | 2.5.1 | 2.5.1 | v2.5.1-7.0.0 | | | 2.4.0 | 2.4.0.post4 | v2.4.0-7.0.0 | | | 2.3.1 | 2.3.1.post6 | v2.3.1-7.0.0 | | | 2.1.0 | 2.1.0.post12 | v2.1.0-7.0.0 | | CANN 8.0.0 | 2.4.0 | 2.4.0.post2 | v2.4.0-6.0.0 | | | 2.3.1 | 2.3.1.post4 | v2.3.1-6.0.0 | | | 2.1.0 | 2.1.0.post10 | v2.1.0-6.0.0 | | CANN 8.0.RC3 | 2.4.0 | 2.4.0 | v2.4.0-6.0.rc3 | | | 2.3.1 | 2.3.1.post2 | v2.3.1-6.0.rc3 | | | 2.1.0 | 2.1.0.post8 | v2.1.0-6.0.rc3 | | CANN 8.0.RC2 | 2.3.1 | 2.3.1 | v2.3.1-6.0.rc2 | | | 2.2.0 | 2.2.0.post2 | v2.2.0-6.0.rc2 | | | 2.1.0 | 2.1.0.post6 | v2.1.0-6.0.rc2 | | | 1.11.0 | 1.11.0.post14 | v1.11.0-6.0.rc2 | | CANN 8.0.RC1 | 2.2.0 | 2.2.0 | v2.2.0-6.0.rc1 | | | 2.1.0 | 2.1.0.post4 | v2.1.0-6.0.rc1 | | | 1.11.0 | 1.11.0.post11 | v1.11.0-6.0.rc1 | | CANN 7.0.0 | 2.1.0 | 2.1.0 | v2.1.0-5.0.0 | | | 2.0.1 | 2.0.1.post1 | v2.0.1-5.0.0 | | | 1.11.0 | 1.11.0.post8 | v1.11.0-5.0.0 | | CANN 7.0.RC1 | 2.1.0 | 2.1.0.rc1 | v2.1.0-5.0.rc3 | | | 2.0.1 | 2.0.1 | v2.0.1-5.0.rc3 | | | 1.11.0 | 1.11.0.post4 | v1.11.0-5.0.rc3 | | CANN 6.3.RC3.1 | 1.11.0 | 1.11.0.post3 | v1.11.0-5.0.rc2.2 | | CANN 6.3.RC3 | 1.11.0 | 1.11.0.post2 | v1.11.0-5.0.rc2.1 | | CANN 6.3.RC2 | 2.0.1 | 2.0.1.rc1 | v2.0.1-5.0.rc2 | | | 1.11.0 | 1.11.0.post1 | v1.11.0-5.0.rc2 | | | 1.8.1 | 1.8.1.post2 | v1.8.1-5.0.rc2 | | CANN 6.3.RC1 | 1.11.0 | 1.11.0 | v1.11.0-5.0.rc…