How to setup PyTorch and CANN environments on Ascend 910b

The following operations are executed on the Ascend 910b platform. If you are using a different environment, please modify the instructions accordingly to suit your specific setup.

install NPU driver and firmware

Check if you have installed the NPU driver and firmware already, if not, you should install them as the first step.

Please go to firmware-drivers to download the latest NPU Driver and NPU firmware, such as Ascend-hdk-910b-npu-firmware_7.5.0.1.129.run and Ascend-hdk-910b-npu-driver_24.1.rc3_linux-x86-64.run

Execute the following command to install them

./Ascend-hdk-910b-npu-driver_24.1.rc3_linux-x86-64.run --full  --install-for-all

./Ascend-hdk-910b-npu-firmware_7.5.0.1.129.run --full

And you can reference this

install necessary libraries

Execute the following commands to install Python3.10, IPython

apt update
apt install -y software-properties-common wget
add-apt-repository ppa:deadsnakes/ppa
apt update
apt install -y python3.10
apt install -y python3-pip
pip3 install ipython

install PyTorch 2.40

pip3 install numpy==1.26.4
pip3 install torch==2.4.0
# install torch_npu
pip3 install pyyaml
pip3 install setuptools
wget https://gitee.com/ascend/pytorch/releases/download/v6.0.rc3-pytorch2.4.0/torch_npu-2.4.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
pip3 install torch_npu-2.4.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

install CANN

Download necessary cann packages from CANN packages such as Ascend-cann-kernels-910b_8.0.RC3_linux-aarch64.run and Ascend-cann-toolkit_8.0.RC3_linux-aarch64.run

Execute the following commands to install

chmod +x Ascend-cann-toolkit_8.0.RC3_linux-aarch64.run
chmod +x Ascend-cann-kernels-910b_8.0.RC3_linux-aarch64.run

./Ascend-cann-toolkit_8.0.RC3_linux-aarch64.run   --quiet --install --install-for-all
./Ascend-cann-kernels-910b_8.0.RC3_linux-aarch64.run  --quiet --install --install-for-all

# set environments automatically
echo 'source /usr/local/Ascend/ascend-toolkit/set_env.sh' >> ~/.bashrc

Please make sure the NPU driver is installed before you install CANN

And if everything goes well, you can execute the following python codes without any exceptions

import torch
import torch_npu

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

print(z)

Please reference the folowing materials when encountering problems when installing

install