近期知道了我校新建了高性能计算平台，于是找老师申请了个账号。上去一看发现居然三台 GPU 节点，每台都有两块 A100 40G，遂打算开始炼丹。结果发现学校为了国产化，服务器 CPU 居然是 arm64 指令集的鲲鹏 920，导致 PyTorch 无法安装。折腾许久，总结了本篇文章，用容器解决这个问题，希望能帮助到大家。

[2024/03/24 UPD]: 真尼玛难用，别折腾国产平台了，快跑！

1 问题表现

上服务器后发现无论是 pip 还是 conda，怎么都没法安装上 PyTorch CUDA 版本，即使安上了调用测试代码也报错无 CUDA 设备。这让我非常纳闷，明明复制的官网安装指令，nvidia-smi 也能回显两块显卡，怎么就跑不了代码？

经过研究发现，原来是 PyTorch 官方压根没有为 aarch64 平台编译带 CUDA 的二进制 whl，在 https://download.pytorch.org/whl/torch/ 可以发现，带 CUDA 的版本只有 x86_64 平台的：

这就麻烦了，如果想安装那就得从源码开始编译，遂开始 clone 源码，准备工具链。但是由于学校没有给用户 sudo 权限，编译环境很难准备。更令人窒息的是，有些 PyTorch 编译前置环境竟然也没有提供 aarch64 平台的二进制，为了编译 PyTorch 居然还得先编译前置环境！这谁受得了啊。

和运维拉扯了半天，最后决定用容器化解决该问题。

2 容器平台

众所周知，Docker 容器的管理均需要 root 权限。对于共享 HPC 平台这种应用场景，肯定不能把 root 权限给用户，因此运维选择安装 singularity 这个容器平台。

这个容器平台兼容 Docker 镜像，可以直接从 Docker Registry 拉取镜像。它的最大特点就是不需要 root 权限，普通用户就能启容器，并且 singularity 的不会像 Docker 完全隔离宿主机和容器，启动容器后甚至能直接访问到宿主家目录，省了 bind 绑定的事。

当然，如果你有 Docker 环境，那就不需要用这个了，直接参阅官方文档用 Docker 启容器就行了：https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/running.html

3 安装 PyTorch

之所以能用容器化解决这个问题，是因为 PyTorch 提供了 arm64 平台的 NGC 镜像，该镜像打包好了 CUDA, cuBLAS, cuDNN, TensorRT, Torch-TensorRT，环境已经是即开即用的级别了。

3.1 选择镜像

前往 NGC 储存库查看镜像列表：https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags

每个版本打包的内容版本也有表格：https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-02.html

需要注意的是要选择不带 igpu 结尾，另外镜像都有 amd64 和 arm64 两个版本，不需要手动指定，会自动拉取正确平台的镜像。

下面以撰文时最新版举例：nvcr.io/nvidia/pytorch:24.02-py3

3.2 拉取镜像

使用以下命令拉取 Docker 镜像：

singularity pull docker://nvcr.io/nvidia/pytorch:24.02-py3

以 docker:// 开头代表着该镜像是 Docker 镜像，后面可以接任意 Docker 镜像。拉取完成后，会获得一个 sif 文件，例如 pytorch_24.02-py3.sif.

3.3 运行镜像

接下来使用以下命令运行镜像：

singularity run --nv pytorch_24.02-py3.sif

需要注意的是，--nv 参数至关重要，加上它后 singularity 运行容器时才会初始化 NVIDIA 相关驱动和设备。

运行后关注回显，检查是否有 ERROR 字眼，如果有的话大概率是宿主机显卡相关环境有问题。我校平台宿主机并没有安装 cuda-tool-kits，只安装了 NVIDIA 驱动，但是实测下来是可以正常启动容器并使用的。

此时可以发现，目前的 Shell Prompt 是 Singularity>，不过实际上现在已经是 Bash 在做终端了。如果你不习惯这个 Prompt，可以输入 /bin/bash 切到原版 Bash 里。

3.4 测试环境

首先确保启动容器时没有 ERROR 字眼，因为有些时候下面两个测试回显正常，但是实际上是不能调用到显卡的。

1. 测试显卡驱动：nvidia-smi

2. 测试 CUDA 工具包：nvcc -V

3. 实际运行 PyTorch 进行测试

进入 Python 交互式界面，运行以下代码对环境进行测试，如果回显正确大概率就没问题了。

Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.version.cuda
'12.3'
>>> torch.cuda.is_available()
True
>>> torch.cuda.device_count()
2
>>> torch.cuda.current_device()
0
>>> torch.cuda.device(0)
<torch.cuda.device object at 0x40024a84cf70>
>>> torch.cuda.get_device_name(0)
'NVIDIA A100-PCIE-40GB'

4 一些感想

这几天为了解决这个问题，折腾了许久，最后还是靠让运维装了个容器平台来解决的，我想如果显卡也换成国产平台，大概率是真的没法解决了。在找资料的过程发现，国产平台的资料实在是少之又少，源头开发者基本上都没有关注国产平台的兼容性，甚至 PyTorch 官方发 aarch64 二进制是为了 macOS 平台的。

源头开发者不关注，那国内开发者是否关注呢？实际上也不算关注。因为目前国产平台还没有推广开，国内相关资料也少之又少。看到 CSDN 一篇博客和我遇到相同问题，但是也只表明了问题在这但解决不了。华为这个平台，基本上只有华为官网有少许资料，其他都是一些灌水吹嘘的新闻报道，毫无价值。

因此，我也切身体会到了国产化道阻且长，不仅是技术，而且是资料积累均还有一定差距。希望我这篇博客也能为国产平台的应用提供微小的助力吧。