Ascend Extension for PyTorch
Overview
This repository develops the PyTorch Extension 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.post13
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 the 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 v2.1.0-7.1.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 8.5.0
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=13
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
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 |
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 | Supported PyTorch Version | Supported Extension Version | Github Branch |
|---|---|---|---|
| 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.rc1 |
| 1.8.1 | 1.8.1.post1 | v1.8.1-5.0.rc1 | |
| CANN 6.0.1 | 1.5.0 | 1.5.0.post8 | v1.5.0-3.0.0 |
| 1.8.1 | 1.8.1 | v1.8.1-3.0.0 | |
| 1.11.0 | 1.11.0.rc2（beta) | v1.11.0-3.0.0 | |
| CANN 6.0.RC1 | 1.5.0 | 1.5.0.post7 | v1.5.0-3.0.rc3 |
| 1.8.1 | 1.8.1.rc3 | v1.8.1-3.0.rc3 | |
| 1.11.0 | 1.11.0.rc1（beta) | v1.11.0-3.0.rc3 | |
| CANN 5.1.RC2 | 1.5.0 | 1.5.0.post6 | v1.5.0-3.0.rc2 |
| 1.8.1 | 1.8.1.rc2 | v1.8.1-3.0.rc2 | |
| CANN 5.1.RC1 | 1.5.0 | 1.5.0.post5 | v1.5.0-3.0.rc1 |
| 1.8.1 | 1.8.1.rc1 | v1.8.1-3.0.rc1 | |
| CANN 5.0.4 | 1.5.0 | 1.5.0.post4 | 2.0.4.tr5 |
| CANN 5.0.3 | 1.8.1 | 1.5.0.post3 | 2.0.3.tr5 |
| CANN 5.0.2 | 1.5.0 | 1.5.0.post2 | 2.0.2.tr5 |
Hardware support
The Ascend training device includes the following models, all of which can be used as training environments for PyTorch models
| Product series | Product model |
|---|---|
| Atlas Training series products | Atlas 800（model: 9000） |
| Atlas 800（model：9010） | |
| Atlas 900 PoD（model：9000） | |
| Atlas 300T（model：9000） | |
| Atlas 300T Pro（model：9000） | |
| Atlas A2 Training series products | Atlas 800T A2 |
| Atlas 900 A2 PoD | |
| Atlas 200T A2 Box16 | |
| Atlas 300T A2 |
The Ascend inference device includes the following models, all of which can be used as inference environments for large models
| Product series | Product model |
|---|---|
| Atlas 800I A2 Inference product | Atlas 800I A2 |
Suggestions and Communication
Everyone is welcome to contribute to the community. If you have any questions or suggestions, you can submit Github Issues. We will reply to you as soon as possible. Thank you very much.
Branch Maintenance Policies
The version branches of AscendPyTorch have the following maintenance phases:
| Status | Duration | Description |
|---|---|---|
| Planning | 1-3 months | Plan features. |
| Development | 6-12 months | Develop new features and fix issues, regularly release new versions. Different strategies are adopted for different versions of PyTorch, with a regular branch development cycle of 6 months and a long-term support branch development cycle of 12 months. |
| Maintained | 1 year/3.5 years | Regular Release branch for 1 year, Long Term Support branch maintenance for 3.5 years. Fix major issues, do not incorporate new features, and release patch versions based on the impact of fixed bugs. |
| End Of Life (EOL) | N/A | Do not accept any modification to a branch. |
PyTorch Maintenance Policies
| PyTorch | Maintenance Policies | Status | Launch Date | Subsequent Status | EOL Date |
|---|---|---|---|---|---|
| 2.6.0 | Regular Release | Development | 2025/07/25 | Expected to enter maintenance status from January 25, 2026 | - |
| 2.5.1 | Regular Release | Development | 2024/11/08 | Expected to enter maintenance status from August 8, 2025 | |
| 2.4.0 | Regular Release | Maintained | 2024/10/15 | Expected to enter maintenance free status from June 15, 2026 | |
| 2.3.1 | Regular Release | Maintained | 2024/06/06 | Expected to enter maintenance free status from June 7, 2026 | |
| 2.2.0 | Regular Release | Maintained | 2024/04/01 | Expected to enter maintenance free status from September 10, 2025 | |
| 2.1.0 | Long Term Support | Development | 2023/10/15 | Expected to enter maintenance status from September 15, 2025 | |
| 2.0.1 | Regular Release | EOL | 2023/7/19 | 2024/3/14 | |
| 1.11.0 | Long Term Support | Maintained | 2023/4/19 | Expected to enter maintenance free status from September 10, 2025 | |
| 1.8.1 | Long Term Support | EOL | 2022/4/10 | 2023/4/10 | |
| 1.5.0 | Long Term Support | EOL | 2021/7/29 | 2022/7/29 |
Reference Documents
For more detailed information on installation guides, model migration, training/inference tutorials, and API lists, please refer to the PyTorch Ascend Adapter on the HiAI Community.
| Document Name | Document Link |
|---|---|
| Installation Guide | link |
| Network Model Migration and Training | link |
| Operator Adaptation | link |
| API List (PyTorch and Custom Interfaces) | link |
License
Ascend Extension for Pytorch has a BSD-style license, as found in the LICENSE file.
2 maintainers
Metadata
2 maintainers
Release files for torch-npu 2.1.0.post13
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distributions (wheels)
Total release size: 174.8 MB
Release files / torch_npu-2.1.0.post13-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
| Download URL | torch_npu-2.1.0.post13-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl |
|---|---|
| Size | 23.8 MB |
| Tags | CPython 3.11 Linux glibc 2.17+ x86-64 |
|
SHA-256 checksum How to use checksums |
3bc35098e1f91fc9912e5413a5e81433b2b178e3b3fad424af9207c1f8b80e62
|
|
BLAKE2b-256 checksum How to use checksums |
8b6b411d49b30bb124389ba6dbfb2584c2988f67580d68e5de44ea5feddb06b7
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
Release files / torch_npu-2.1.0.post13-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
| Download URL | torch_npu-2.1.0.post13-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl |
|---|---|
| Size | 19.9 MB |
| Tags | CPython 3.11 Linux glibc 2.17+ ARM64 |
|
SHA-256 checksum How to use checksums |
5ecb8b1cb4fe132f39d9555b641a5604cf8731eeba6b08d217c16b07749bdfd8
|
|
BLAKE2b-256 checksum How to use checksums |
1eb07a19c8f8d7b6b24fe7e4925862466d5524b15341ff4b09b795fd530c81c5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
Release files / torch_npu-2.1.0.post13-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
| Download URL | torch_npu-2.1.0.post13-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl |
|---|---|
| Size | 23.8 MB |
| Tags | CPython 3.10 Linux glibc 2.17+ x86-64 |
|
SHA-256 checksum How to use checksums |
c94e5229425b08ebc36c6992f76984de18b7421605697cda63a82601c07ec076
|
|
BLAKE2b-256 checksum How to use checksums |
f363cf0c17b4ae78eb49ecf46f2d5277345cb506c629aae20b249e8bb64850f8
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
Release files / torch_npu-2.1.0.post13-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
| Download URL | torch_npu-2.1.0.post13-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl |
|---|---|
| Size | 19.9 MB |
| Tags | CPython 3.10 Linux glibc 2.17+ ARM64 |
|
SHA-256 checksum How to use checksums |
aa13f32b2701b3fe1ba3e3de98a178e68519778b6be41805e5618e7fd0fd8104
|
|
BLAKE2b-256 checksum How to use checksums |
37b59cc9b346dda3c1a3cdbccb8610c3b56618053df304a2bd1bdeb1149110f2
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
Release files / torch_npu-2.1.0.post13-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
| Download URL | torch_npu-2.1.0.post13-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl |
|---|---|
| Size | 23.8 MB |
| Tags | CPython 3.9 Linux glibc 2.17+ x86-64 |
|
SHA-256 checksum How to use checksums |
68a44e92b377be5767b1e2f4285a98db0b52da3f3e385695dedc04b49839a98b
|
|
BLAKE2b-256 checksum How to use checksums |
c87144371e172886f0de0a48ffbe9a7b7275559aa275110d43558712005460f9
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
Release files / torch_npu-2.1.0.post13-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
| Download URL | torch_npu-2.1.0.post13-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl |
|---|---|
| Size | 19.9 MB |
| Tags | CPython 3.9 Linux glibc 2.17+ ARM64 |
|
SHA-256 checksum How to use checksums |
d9729232582b00f056794515215bf28393d61382b3b43d7befc335b0ef94eb76
|
|
BLAKE2b-256 checksum How to use checksums |
ea0ebee524554bdb18fd0865576ee2796f7e570a3be872cd4efa34e0df339679
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
Release files / torch_npu-2.1.0.post13-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
| Download URL | torch_npu-2.1.0.post13-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl |
|---|---|
| Size | 23.8 MB |
| Tags | CPython 3.8 Linux glibc 2.17+ x86-64 |
|
SHA-256 checksum How to use checksums |
da3c8434aa3479cac493c8e80127526f2ba8c402cc566cf2ac30e1897efd0bab
|
|
BLAKE2b-256 checksum How to use checksums |
1bafa5489ca3ed1fa11bac1cbf951e6ab2e0d8f069c079572285bb09d3a13f6f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
Release files / torch_npu-2.1.0.post13-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
| Download URL | torch_npu-2.1.0.post13-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl |
|---|---|
| Size | 19.9 MB |
| Tags | CPython 3.8 Linux glibc 2.17+ ARM64 |
|
SHA-256 checksum How to use checksums |
1d2084b9fde4fc2bfec72a2f51f9f4842ae948e01564cd7153b1f958876ecbf4
|
|
BLAKE2b-256 checksum How to use checksums |
90507100ea7f37b00b6957fd1ecb189aba915dfd71c2e92596b9494e37a94b26
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.0.0 CPython/3.10.10
|
