Skip to main content

NPU bridge for PyTorch

Project description

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.

  1. Install PyTorch

Install PyTorch through pip.

For Aarch64:

pip3 install torch==2.7.1

For x86:

pip3 install torch==2.7.1+cpu  --index-url https://download.pytorch.org/whl/cpu
  1. 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.9 link
x86 Python3.10 link
x86 Python3.11 link
aarch64 Python3.9 link
aarch64 Python3.10 link
aarch64 Python3.11 link
  1. Install torch-npu
pip3 install torch-npu==2.7.1

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 11.2 for both ARM and X86.

  1. Clone torch-npu

    git clone https://github.com/ascend/pytorch.git -b 2.7.1-7.2.0 --depth 1
    
  2. Build Docker Image

    cd pytorch/ci/docker/{arch} # {arch} for X86 or ARM
    docker build -t manylinux-builder:v1 .
    
  3. 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
    
  4. Compile torch-npu

    Take Python 3.9 as an example.

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

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 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.1.RC1.alpha002 2.6.0 2.6.0rc1 v2.6.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

Pipeline Status

Due to the asynchronous development mechanism of upstream and downstream, incompatible modifications in upstream may cause some functions of torch_npu to be unavailable (only upstream and downstream development branches are involved, excluding stable branches). Therefore, we built a set of daily tasks that make it easy to detect relevant issues in time and fix them within 48 hours (under normal circumstances), providing users with the latest features and stable quality.

OS CANN Version(Docker Image) Upstream Branch Downstream Branch Period Status
openEuler 22.03 SP2 CANN 7.1 main master UTC 1200 daily Ascend NPU

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.9.0 Regular Release Development 2026/01/15 Expected to enter maintenance status from July 15, 2026 -
2.8.0 Regular Release Development 2025/10/15 Expected to enter maintenance status from March 15, 2026 -
2.7.1 Long Term Support Development 2025/10/15 Expected to enter maintenance status from October 15, 2026
2.6.0 Regular Release Development 2025/07/25 Expected to enter maintenance status from January 15, 2026 -
2.5.1 Regular Release Maintained 2024/11/08 Expected to enter maintenance free status from August 8, 2026
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 EOL 2024/04/01 2025/10/14
2.1.0 Long Term Support Maintained 2023/10/15 Expected to enter maintenance free status from December 30, 2026
2.0.1 Regular Release EOL 2023/7/19 2024/3/14
1.11.0 Long Term Support EOL 2023/4/19 2025/10/25
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 Ascend Extension for PyTorch 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.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

torch_npu-2.7.1.post4-cp313-cp313-manylinux_2_28_x86_64.whl (34.4 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.28+ x86-64

torch_npu-2.7.1.post4-cp313-cp313-manylinux_2_28_aarch64.whl (31.9 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

torch_npu-2.7.1.post4-cp312-cp312-manylinux_2_28_x86_64.whl (34.4 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.28+ x86-64

torch_npu-2.7.1.post4-cp312-cp312-manylinux_2_28_aarch64.whl (31.9 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_x86_64.whl (34.4 MB view details)

Uploaded CPython 3.11manylinux: glibc 2.28+ x86-64

torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_aarch64.whl (31.9 MB view details)

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

torch_npu-2.7.1.post4-cp310-cp310-manylinux_2_28_x86_64.whl (34.4 MB view details)

Uploaded CPython 3.10manylinux: glibc 2.28+ x86-64

torch_npu-2.7.1.post4-cp310-cp310-manylinux_2_28_aarch64.whl (31.9 MB view details)

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

torch_npu-2.7.1.post4-cp39-cp39-manylinux_2_28_x86_64.whl (34.4 MB view details)

Uploaded CPython 3.9manylinux: glibc 2.28+ x86-64

torch_npu-2.7.1.post4-cp39-cp39-manylinux_2_28_aarch64.whl (31.9 MB view details)

Uploaded CPython 3.9manylinux: glibc 2.28+ ARM64

File details

Details for the file torch_npu-2.7.1.post4-cp313-cp313-manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp313-cp313-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 64f871513941308a429650c19df17d36d11d9e15179d4eac89b96900cf18c0d0
MD5 e28318d4b960c7d115bb96e30d2276ab
BLAKE2b-256 8c3b40ffc656ad7c1b616c14cee9b83d41eb3013a2d31aafbaa4e1af143aacfa

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp313-cp313-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 b3b7f02e3488b629e429cbef0b6db021e82b1bcdff96dfe2ae4ca14f43ba567b
MD5 4095532e30c4b699ae4e300067775f32
BLAKE2b-256 fa677ba2a1cc298f0aa76f8b6222772b7a804f9b0236c4a70dc62599cb40b2ff

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp312-cp312-manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp312-cp312-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 c9ca92f7980c46dfedfbd4c18a9e012003faac3c0cd776ce4d7217b913b3f49b
MD5 f1dee25dd7414a5a4080d0de28d01888
BLAKE2b-256 3e7c794c8f0b42affe8046317f125644de40fcb22806d6112a1a8a25237af13a

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp312-cp312-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 2a86cbfa41f576833e50ecbc81d0c2166dfb8a333ba98b1abe1921a2f7730875
MD5 6e50366ae67c91429bf3f6ae17054a02
BLAKE2b-256 9af279c831743266c2698a22095da7f541ffd3d70f0d8793e3166a99abbbab5a

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 b395f58095185f7aabf6c65c96a858e0c9181625716a5482da22658564231e51
MD5 9012bc6725c23b66898d04cc5c195a2b
BLAKE2b-256 cd81a1707135ed61a7daefd0324dc1d7cf4618d1058d16e995b18f19dc5ae215

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 8520a8214966e1af1205a1d5717a8ae53e6bde28731228c8562396d4947bd2d6
MD5 3ab211c203d6cf7af57f3359e09fa1f7
BLAKE2b-256 fccfc4814bc501b149ad492d0b36406e489092d57df13f741df5897e59fc3d60

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp310-cp310-manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp310-cp310-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 d223cb616341fc73302a05135db5e0e092899a14ea55419e4c63b82db7e85379
MD5 e0d105297975d6404e78bdfaed01a9c2
BLAKE2b-256 5975d6df4b81473a731bdb3993b8516aa49a1e56b3d602c6520b9ad36bce2b5d

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp310-cp310-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 9d87039f3e796415d759a04c1afbd34745293b900231d93c734849f7ad70ec84
MD5 12d86ad5980630999197032a920d0751
BLAKE2b-256 cebf932156949114a5c3e3d54d673efdfc1d4df2caaa2f4e4acd734ee64efa7d

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp39-cp39-manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp39-cp39-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 749c32be3ddee1186a08ff7148b9905cde43d122e6b3b24f80a45fe918221784
MD5 ed1121e82cac5be1dc67700187ce493a
BLAKE2b-256 7d8c2efd4bdcd5feaaa20fc61e348c8e569b9577cef14e63f174811f19ef6bdd

See more details on using hashes here.

File details

Details for the file torch_npu-2.7.1.post4-cp39-cp39-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post4-cp39-cp39-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 de4b95482eddb39f14d5d450ca8d69bf887802e4d8961ef58e027edf7b53bb0d
MD5 70ad97b28b9ca66bb0098d4257c200dd
BLAKE2b-256 84bb307a7f6abc8b4b5e7e15a0a337b54dbd48c250175f37067bf7bb6aefd152

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page