容器环境安装PyTorch Ascend NPU

[Yikun]

作者：@Yikun

本文介绍如何在容器内安装CANN、PyTorch、PyTorch Ascend Adapter。

1. 构建基础容器

FROM ubuntu:20.04

ENV DEBIAN_FRONTEND noninteractive
ENV DEBCONF_NONINTERACTIVE_SEEN true

RUN sed -i 's/ports.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list && \
    apt update && \
    apt install -y adduser sudo vim gcc g++ cmake make gdb git tmux openssh-server net-tools iputils-ping python3 python3-venv python3-setuptools gcc python3-dev patchelf tree && \
    apt remove -y cmake

1.1 构建并进入容器

# 构建容器
docker build -t ubuntu-20.04-yikun -f /data/disk3/yikun/Dockerfile .
# 启动容器
docker run --network host --name yikun --device /dev/davinci5 --device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ -v /data/disk3/yikun:/opt/yikun -ti ubuntu-20.04-yikun:latest bash
# 退出容器
ctrl+p q
# 通过exec进入容器
docker exec -ti yikun bash
export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/common/:/usr/local/Ascend/driver/lib64/driver/:$LD_LIBRARY_PATH
npu-smi info

1.2 容器内安装CANN

wget https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/Milan-ASL/Milan-ASL%20V100R001C13SPC701/Ascend-cann-nnae_7.0.RC1.alpha001_linux-aarch64.run

chmod +x Ascend-cann-nnae_7.0.RC1.alpha001_linux-aarch64.run

Ascend-cann-nnae_7.0.RC1.alpha001_linux-aarch64.run --install

1.3 环境变量加到bashrc

echo """
source /usr/local/Ascend/ascend-toolkit/set_env.sh
export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/common/:/usr/local/Ascend/driver/lib64/driver/:\$LD_LIBRARY_PATH
source ~/.llm-venv/bin/activate
""" >> ~/.bashrc

2. 初始化PyThon环境(venv为例)

python3 -m venv ~/.llm-venv
source ~/.llm-venv/bin/activate

# 换PyPI源
mkdir ~/.pip/
echo """[global]
index-url = https://mirrors.aliyun.com/pypi/simple/
[install]
trusted-host=[mirrors.aliyun.com](http://mirrors.aliyun.com/)
""" > ~/.pip/pip.conf

# 安装PyTorch基础依赖
pip install cmake ninja

3. 安装PyTorch

git clone https://github.com/pytorch/pytorch.git
cd pytorch
git reset 5913437a40a6e45ab7e164afb7c6ec930dd40b2f --hard
export USE_CUDA=0
export USE_XNNPACK=0
pip install -r ./requirements.txt
python setup.py develop

# 如果遇到国内网络无法访问
export http_proxy=http://127.0.0.1:1234;export https_proxy=$http_proxy;no_proxy=127.0.0.1
git submodule update --init --recursive
unset http_proxy;unset https_proxy;

4. 安装PyTorch NPU

git clone https://gitee.com/ascend/pytorch.git ascend-pytorch
git reset  81ece7b664adc28e698044d0e9091d39a1a6dfa6 --hard
pip3 install pyyaml wheel decorator snippy
# 确保python3.8是venv里的python，否则`cd ~/.llm-venv/bin/;ln -sf python3 python3.8`
bash ci/build.sh --python=3.8
# dist下面会生成whl
pip3 install --upgrade dist/torch_npu-2.1.0-cp38-cp38-linux_aarch64.whl

# 进入到非pytorch目录，测试
cd ~ 
python3 -c "import torch;import torch_npu; a = torch.randn(3, 4).npu(); print(a + a);"
Warning: Device do not support double dtype now, dtype cast repalce with float.
tensor([[-2.5684, -1.0355,  2.1467, -1.9409],
        [-0.8765, -2.1909,  1.6815,  2.1013],
        [ 7.0935,  3.7160,  1.6914, -0.1912]], device='npu:0')