飞腾S5000C ARM服务器PyTorch+CUDA环境搭建指南

1. 飞腾 S5000C 服务器环境搭建概述

飞腾 S5000C 是一款基于 ARM 架构的高性能服务器处理器，在国产化替代和特定计算场景中具有重要地位。搭建 PyTorch + CUDA + RTX 4090D 的开发环境，需要特别注意 ARM 架构与传统 x86 环境的差异。本文将详细介绍从零开始的环境搭建过程，包括驱动安装、CUDA 配置、PyTorch 适配等关键环节。

在 ARM 架构上部署深度学习环境面临几个独特挑战：首先是软件生态的兼容性问题，许多预编译的二进制包默认只支持 x86 架构；其次是驱动和工具链的版本匹配，特别是 NVIDIA 显卡驱动对 ARM 平台的支持情况；最后是性能调优，需要针对 ARM 架构特点进行特定优化。

2. 基础系统环境准备

2.1 操作系统选择与安装

推荐使用 Ubuntu 22.04 LTS for ARM64 作为基础操作系统，这是目前对飞腾平台支持最完善的发行版之一。安装时需注意：

1. 下载官方 ARM64 版本镜像（非 x86 版本）
2. 安装时选择最小化安装，减少不必要的软件包
3. 确保系统内核版本不低于 5.15（可通过 uname -r 检查）

安装完成后，首先更新系统：

bash复制sudo apt update && sudo apt upgrade -y

2.2 基础依赖安装

深度学习环境需要以下基础依赖：

bash复制sudo apt install -y build-essential cmake git wget \
    libopenblas-dev liblapack-dev libatlas-base-dev \
    python3-dev python3-pip python3-venv

对于 ARM 平台，特别需要安装：

bash复制sudo apt install -y gfortran-arm-linux-gnueabihf

3. NVIDIA 驱动与 CUDA 安装

3.1 驱动安装准备

RTX 4090D 需要 NVIDIA 驱动版本 >= 525.60.13。在安装前需：

1. 禁用 Nouveau 驱动：

   bash复制echo "blacklist nouveau" | sudo tee /etc/modprobe.d/blacklist-nouveau.conf
   echo "options nouveau modeset=0" | sudo tee -a /etc/modprobe.d/blacklist-nouveau.conf
   sudo update-initramfs -u
   

2. 重启后验证：

   bash复制lsmod | grep -i nouveau  # 应无输出
   

3.2 驱动安装

从 NVIDIA 官网下载 ARM64 版本的驱动安装包（.run 格式），然后：

bash复制chmod +x NVIDIA-Linux-aarch64-*.run
sudo ./NVIDIA-Linux-aarch64-*.run

安装完成后验证：

bash复制nvidia-smi  # 应显示显卡信息

3.3 CUDA Toolkit 安装

由于官方 CUDA 对 ARM 支持有限，推荐使用以下方法：

1. 下载 CUDA 11.7 网络安装包（ARM64 版本）：

   bash复制wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda-repo-ubuntu2204-11-7-local_11.7.0-515.43.04-1_arm64.deb
   

2. 安装：

   bash复制sudo dpkg -i cuda-repo-ubuntu2204-11-7-local_11.7.0-515.43.04-1_arm64.deb
   sudo cp /var/cuda-repo-ubuntu2204-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
   sudo apt-get update
   sudo apt-get -y install cuda
   

3. 环境变量配置：

   bash复制echo 'export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}' >> ~/.bashrc
   echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
   source ~/.bashrc
   

验证安装：

bash复制nvcc --version  # 应显示 CUDA 11.7

4. PyTorch 安装与配置

4.1 创建 Python 虚拟环境

推荐使用 conda 管理环境：

bash复制wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
bash Miniconda3-latest-Linux-aarch64.sh
# 按照提示安装后
conda create -n pytorch_env python=3.9
conda activate pytorch_env

4.2 PyTorch 安装

由于官方 PyTorch 预编译包不支持 ARM，需要从源码编译：

1. 安装编译依赖：

   bash复制conda install numpy pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests dataclasses
   

2. 获取源码：

   bash复制git clone --recursive https://github.com/pytorch/pytorch
   cd pytorch
   git checkout v1.12.1  # 选择稳定版本
   git submodule sync
   git submodule update --init --recursive
   

3. 编译配置：

   bash复制export USE_CUDA=1
   export USE_CUDNN=1
   export USE_MKLDNN=1
   export USE_OPENMP=1
   export BUILD_TEST=0
   

4. 开始编译：

   bash复制python setup.py install
   

编译过程可能需要数小时，建议在 screen 或 tmux 会话中进行。

4.3 验证安装

编译完成后验证：

python复制import torch
print(torch.__version__)  # 应显示版本号
print(torch.cuda.is_available())  # 应返回 True

5. 性能优化与问题排查

5.1 ARM 架构特定优化

1. 启用 OpenBLAS：

   bash复制sudo update-alternatives --config libblas.so.3
   # 选择 OpenBLAS 选项
   

2. 调整 CPU 频率策略：

   bash复制sudo apt install cpufrequtils
   sudo cpufreq-set -g performance
   

5.2 常见问题解决

1. CUDA 版本不匹配：

   python复制# 如果遇到类似错误：
   # "CUDA error: no kernel image is available for execution"
   # 需要确保 PyTorch 编译时的 CUDA 版本与系统安装的一致
   

2. 驱动兼容性问题：

   bash复制# 如果 nvidia-smi 显示的 CUDA 版本与 nvcc 不一致
   # 需要重新安装匹配版本的驱动和 CUDA Toolkit
   

3. PyTorch 运行缓慢：

   python复制# 在代码开头添加：
   torch.backends.cudnn.benchmark = True
   

6. 实际测试与基准对比

6.1 基础性能测试

使用以下代码测试矩阵乘法性能：

python复制import torch
import time

device = torch.device("cuda")
size = 10240
a = torch.randn(size, size, device=device)
b = torch.randn(size, size, device=device)

start = time.time()
c = torch.matmul(a, b)
torch.cuda.synchronize()
print(f"Time: {time.time()-start:.4f}s")

6.2 与 x86 平台对比

在相同 RTX 4090D 显卡下，ARM 与 x86 平台的典型性能差异：

差异主要来自 ARM 平台的内存访问模式和指令集优化空间。

6.3 实际模型训练测试

以 MNIST 分类为例：

python复制import torch
import torchvision
import torch.nn as nn
import torch.optim as optim

# 数据准备
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.1307,), (0.3081,))
])
train_set = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(
    train_set, batch_size=64, shuffle=True)

# 模型定义
model = nn.Sequential(
    nn.Conv2d(1, 32, 3, 1),
    nn.ReLU(),
    nn.Conv2d(32, 64, 3, 1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(9216, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
).cuda()

# 训练循环
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

for epoch in range(5):
    for data, target in train_loader:
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

7. 环境维护与升级建议

7.1 日常维护

1. 定期检查驱动更新：

   bash复制sudo apt list --upgradable | grep nvidia
   

2. 监控 GPU 使用情况：

   bash复制watch -n 1 nvidia-smi
   

7.2 升级策略

1. CUDA 升级：

   • 先卸载旧版本：sudo apt remove --purge '^cuda.*'
   • 然后安装新版本（确保有 ARM 支持）

2. PyTorch 升级：

   • 建议在新的虚拟环境中测试新版本
   • 确认新版本对 ARM 的支持情况

3. 驱动升级：

   • 下载新版驱动包后，先卸载旧版：

     bash复制sudo nvidia-uninstall
     

   • 然后安装新版

7.3 备份与恢复

1. 环境备份：

   bash复制conda env export > pytorch_env.yaml
   pip freeze > requirements.txt
   

2. 恢复环境：

   bash复制conda env create -f pytorch_env.yaml
   pip install -r requirements.txt
   

在飞腾 S5000C 上成功搭建 PyTorch + CUDA 环境后，可以充分发挥 RTX 4090D 的计算能力，为国产化平台的深度学习应用提供强大支持。实际使用中，建议针对 ARM 架构特点进行算法和实现上的优化，以获得最佳性能。