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:

pip3 install torch==2.1.0

For x86:

pip3 install torch==2.1.0+cpu  --index-url https://download.pytorch.org/whl/cpu

Install torch-npu dependencies

Run the following command to install dependencies.

pip3 install pyyaml
pip3 install setuptools

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

pip3 install torch-npu==2.1.0.post17

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

git clone https://github.com/ascend/pytorch.git -b 2.1.0-7.2.0 --depth 1

Build Docker Image

cd pytorch/ci/docker/{arch} # {arch} for X86 or ARM
docker build -t manylinux-builder:v1 .

Enter Docker Container

docker run -it -v /{code_path}/pytorch:/home/pytorch manylinux-builder:v1 bash
# {code_path} is the torch_npu source code path

Compile torch-npu

Take Python 3.8 as an example.

cd /home/pytorch
bash ci/build.sh --python=3.8

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

export _GLIBCXX_USE_CXX11_ABI=1

Meanwhile, we support configuring -fabi-version using the following variables，require consistency with the community torch package

export _ABI_VERSION=16

Getting Started

Prerequisites

Initialize CANN environment variable by running the command as shown below.

# Default path, change it if needed.
source /usr/local/Ascend/ascend-toolkit/set_env.sh

Quick Verification

You can quickly experience Ascend NPU by the following simple examples.

import torch
- import torch_npu # No longer needed in torch_npu 2.5.1 and later versions

x = torch.randn(2, 2).npu()
y = torch.randn(2, 2).npu()
z = x.mm(y)

print(z)

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 {PyTorch version}-{Ascend version}, 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 |-----------------------|------ | CANN 8.5.0 | 2.9.0 | | 2.8.0 | | 2.7.1 | | 2.6.0 | CANN 8.3.RC1 | 2.8.0 | | 2.7.1 | | 2.6.0 | | 2.1.0 | CANN 8.2.RC1 | 2.6.0 | | 2.5.1 | | 2.1.0 | CANN 8.1.RC1 | 2.5.1 | | 2.4.0 | | 2.3.1 | | 2.1.0 | CANN 8.0.0 | 2.4.0 | | 2.3.1 | | 2.1.0 | CANN 8.0.RC3 | 2.4.0 | | 2.3.1 | | 2.1.0 | CANN 8.0.RC2 | 2.3.1 | | 2.2.0 | | 2.1.0 | | 1.11.0 | Supported PyTorch Version | Supported Extension Version | Github Branch | ---------------------|-----------------------------|-------------------| | 2.9.0 | v2.9.0-7.3.0 | | 2.8.0.post2 | v2.8.0-7.3.0 | | 2.7.1.post2 | v2.7.1-7.3.0 | | 2.6.0.post5 | v2.6.0-7.3.0 | | 2.8.0 | v2.8.0-7.2.0 | | 2.7.1 | v2.7.1-7.2.0 | | 2.6.0.post3 | v2.6.0-7.2.0 | | 2.1.0.post17 | v2.1.0-7.2.0 | | 2.6.0 | v2.6.0-7.1.0 | | 2.5.1.post1 | v2.5.1-7.1.0 | | 2.1.0.post13 | v2.1.0-7.1.0 | | 2.5.1 | v2.5.1-7.0.0 | | 2.4.0.post4 | v2.4.0-7.0.0 | | 2.3.1.post6 | v2.3.1-7.0.0 | | 2.1.0.post12 | v2.1.0-7.0.0 | | 2.4.0.post2 | v2.4.0-6.0.0 | | 2.3.1.post4 | v2.3.1-6.0.0 | | 2.1.0.post10 | v2.1.0-6.0.0 | | 2.4.0 | v2.4.0-6.0.rc3 | | 2.3.1.post2 | v2.3.1-6.0.rc3 | | 2.1.0.post8 | v2.1.0-6.0.rc3 | | 2.3.1 | v2.3.1-6.0.rc2 | | 2.2.0.post2 | v2.2.0-6.0.rc2 | | 2.1.0.post6 | v2.1.0-6.0.rc2 |

Pytorch

Install / Use

README