torch-npu 2.1.0rc1

NPU bridge for PyTorch

PyTorch Ascend Adapter

Overview

This repository develops the PyTorch Ascend Adapter 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.1.0

For x86:

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

2. Install torch-npu dependencies

Run the following command to install dependencies.

pip3 install pyyaml
pip3 install setuptools

3. Install torch-npu

Take Aarch64 architecture and Python 3.8 as an example.

NOTE: It will be supported to install torch_npu directly from pypi through pip recently.

wget https://gitee.com/ascend/pytorch/releases/download/v5.0.rc3-pytorch2.1.0/torch_npu-2.1.0rc1-cp38-cp38-linux_aarch64.whl

pip3 install torch_npu-2.1.0rc1-cp38-cp38-linux_aarch64.whl

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:

1. Clone torch-npu

git clone https://github.com/ascend/pytorch.git -b v2.1.0-5.0.rc3 --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

3. Compile torch-npu

Take Python 3.8 as an example.

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

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)

PyTorch and Python Version Matching Table

PyTorch Version	Python Version
PyTorch1.8.1	Python3.7.x(>=3.7.5),Python3.8.x,Python3.9.x
PyTorch1.11.0	Python3.7.x(>=3.7.5),Python3.8.x,Python3.9.x,Python3.10.x
PyTorch2.0.1	Python3.8.x,Python3.9.x,Python3.10.x
PyTorch2.1.0	Python3.8.x,Python3.9.x,Python3.10.x

Ascend Auxiliary Software

CANN Version	Supported PyTorch Version	Supported Adapter Version	Github Branch	AscendHub Image Version/Name(Link)
7.0.RC1.alpha003	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	-
	1.8.1	1.8.1.post3	v1.8.1.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.8.1.post2	v1.8.1-5.0.rc2	23.0.RC1-1.8.1
	1.8.1	1.11.0.post1	v1.11.0-5.0.rc2	23.0.RC1-1.11.0
CANN 6.3.RC1	1.11.0	1.8.1.post1	v1.8.1-5.0.rc1	-
	1.8.1	1.11.0	v1.11.0-5.0.rc1	-
CANN 6.0.1	1.5.0	1.5.0.post8	v1.5.0-3.0.0	22.0.0
	1.8.1	1.8.1	v1.8.1-3.0.0	22.0.0-1.8.1
	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	22.0.RC3
	1.8.1	1.8.1.rc3	v1.8.1-3.0.rc3	22.0.RC3-1.8.1
	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	22.0.RC2
	1.8.1	1.8.1.rc2	v1.8.1-3.0.rc2	22.0.RC2-1.8.1
CANN 5.1.RC1	1.5.0	1.5.0.post5	v1.5.0-3.0.rc1	22.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	21.0.4
CANN 5.0.3	1.8.1	1.5.0.post3	2.0.3.tr5	21.0.3
CANN 5.0.2	1.5.0	1.5.0.post2	2.0.2.tr5	21.0.2

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	3 months	Develop features.
Maintained	6-12 months	Allow the incorporation of all resolved issues and release the version.
Unmaintained	0-3 months	Allow the incorporation of all resolved issues. No dedicated maintenance personnel are available. No version will be released.
End Of Life (EOL)	N/A	Do not accept any modification to a branch.

Maintenance Status of Existing Branches

Branch Name	Status	Launch Date	Subsequent Status	EOL Date
v2.0.2	EOL	2021/7/29	N/A
v2.0.3	EOL	2021/10/15	N/A
v2.0.4	EOL	2022/1/15	N/A
v3.0.rc1	EOL	2022/4/10	N/A
v3.0.rc2	UnMaintained	2022/7/15	EOL 2023-10-15 estimated
v3.0.rc3	Maintained	2022/10/20	Unmaintained 2023-10-20 estimated
v3.0.0	Maintained	2023/1/18	Unmaintained 2024-1-18 estimated
v5.0.rc1	Maintained	2023/4/19	Unmaintained 2024-4-19 estimated
v5.0.rc2	Maintained	2023/7/19	Unmaintained 2024-7-19 estimated

License

PyTorch Ascend Adapter has a BSD-style license, as found in the LICENSE file.