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.post6-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.post6-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.post6-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.post6-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.post6-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.post6-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.post6-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.post6-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.post6-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.post6-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.post6-cp313-cp313-manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp313-cp313-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 76a50cb39d61f8bfdc7d7f3c4c5514f418f5f0e7f2f83dfe1f2d754e05d7fd89
MD5 e33ea25984501d36d973f87010d2d209
BLAKE2b-256 affc9875971264d29591bc4c1d0a7a890a5eb88000dd63832c0cce39feb3e5f6

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 e6cf34a35ac49ca241636d118a6b546ba053a1456ae72a2d72c70a7630b67a88
MD5 53a2e7afdc16ac7342bb6a4ab88dba9b
BLAKE2b-256 37b3a00b9f9dc916e596091f5404b21f2ca7012ecc310e40eb7efd436aa76ab2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp312-cp312-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 77c2ee136130dca1b2abd21e353ad28b0bcecee1b288dad35950a8b132a1d277
MD5 e3721b56f377e4ba748f038b444fad85
BLAKE2b-256 1b58506edce1e11826ed59ab14518d8106bcd61fec9d07d3d242e1629762aedd

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 d1301525fea645cb37d0970b5b5d9b1e05c45eec1e4df624ec0dfcc0a83ff044
MD5 f2240cdcb9d23a02fe6162f4205661fd
BLAKE2b-256 c47585afea9597eb50d61d4907de21b76a4387b31396fa3ef5164e088b3c3282

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp311-cp311-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 26645dfce6c4845a9ba7c2302452eba3a36d92e3129325ac94de9ceb880b8598
MD5 ec5b281f990aba2a2f13a5089d7f7091
BLAKE2b-256 4476fcf937104e984d58727c4f06aa5f833a5ff137c2a31f5433f29de9c71e34

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 2f0693c581dc15cb4c575648581bf2d5407812a5b4df5ccabaa8e6e56fbb1ee8
MD5 fb0f11721aba8a17f8829dd1218f1e76
BLAKE2b-256 ede8e88899cf73c70475eb172905e422969a9485876729a6ceb72b2b504aafea

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp310-cp310-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 af5825e58986154e8072e3d704f27946794c47dec4188da5ec42c8f0e072ff86
MD5 4b8bd455edd728b5840c11bbf558098d
BLAKE2b-256 ae125f2c6232e1bb7c7bb99eaac57f713dc15de9e565167ce9f93fabe8b99f59

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 1a02ccabb04835551c8217300475129a165cd82f59f754b2138e1db494c8e387
MD5 a36aa5f4f2abed7d68f8e10dc8af16f3
BLAKE2b-256 2c17e64bfa0fd33ab1d44ec54cb5f46c564cb2fad0cc65af9e263e3c18f9df2d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp39-cp39-manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 adddc417dc33f261a56901dbb36bb098d16da4bb563cc3740a5ff398e6c95951
MD5 af99236a7c9b156b287aff2f2624d5ec
BLAKE2b-256 db24080a1972231d6d5f7020748dd4899a4e821432189daa72736880b65234e9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for torch_npu-2.7.1.post6-cp39-cp39-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 7f49afe57a2eae9707b369c2b61a92f0bfc8a5687d4022b3c15e74d63fbd5d00
MD5 ba2123edb858bce78d7e34ce1a652152
BLAKE2b-256 a225cbfbdb6a3efb8b86e6909f844fc81b3d8f58118e336d69fd80147a533e62

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