AMD显卡深度学习环境优化：ROCm下的Python包打包实践

1. 项目背景与核心需求

在深度学习领域，AMD显卡用户一直面临着一个尴尬的局面：主流深度学习框架对CUDA生态的高度依赖，使得ROCm环境下的开发体验远不如N卡流畅。作为一名长期使用AMD显卡进行模型训练的开发者，我深刻体会到每次搭建环境时重复下载依赖、编译组件的痛苦。

这个项目的核心目标很简单：将ROCm环境下的关键Python包（如PyTorch、TensorFlow等框架的ROCm版本）及其依赖项打包成独立的wheel文件，实现"一次构建，随处安装"的便捷体验。这不仅能节省大量重复配置时间，对于需要频繁切换开发环境或进行分布式部署的场景尤为实用。

2. 环境准备与基础工具链

2.1 硬件与系统要求

推荐使用以下配置作为打包环境：

• AMD显卡：至少Radeon VII或Instinct MI系列，支持PCIe 3.0以上
• 系统内存：32GB以上（编译大型包时非常消耗内存）
• 磁盘空间：至少100GB可用空间（源码编译会产生大量中间文件）
• 操作系统：Ubuntu 20.04/22.04 LTS（ROCm官方支持最完善的发行版）

注意：虽然ROCm 5.x声称支持更多显卡型号，但实际测试中发现RDNA架构显卡（如RX 6000系列）在编译环节常出现奇怪错误。如果条件允许，建议使用CDNA架构的专业卡。

2.2 基础软件安装

首先配置ROCm基础环境：

bash复制wget -q -O - https://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/5.4.2 ubuntu main' | sudo tee /etc/apt/sources.list.d/rocm.list
sudo apt update
sudo apt install rocm-hip-sdk rocm-opencl-sdk

验证安装：

bash复制/opt/rocm/bin/rocminfo | grep -i 'gpu arch'

2.3 Python环境隔离

强烈建议使用conda创建独立环境：

bash复制conda create -n rocm_build python=3.9 -y
conda activate rocm_build
pip install --upgrade pip wheel auditwheel

3. 核心打包流程详解

3.1 PyTorch的ROCm版本打包

从源码构建PyTorch的wheel文件：

bash复制git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
export ROCM_PATH=/opt/rocm
export PYTORCH_BUILD_VERSION=1.13.1
export PYTORCH_BUILD_NUMBER=1
python setup.py bdist_wheel

关键参数说明：

• MAX_JOBS=4：根据CPU核心数调整编译线程数
• USE_ROCM=1：强制启用ROCm支持
• BUILD_TEST=0：跳过测试以加快编译速度

生成的wheel文件位于dist/目录，使用auditwheel修复依赖：

bash复制auditwheel repair torch-1.13.1-cp39-cp39-linux_x86_64.whl

3.2 TensorFlow ROCm版打包

TensorFlow的构建更为复杂，需要先安装bazel：

bash复制sudo apt install bazel-5.3.0
wget https://github.com/tensorflow/tensorflow/archive/refs/tags/v2.10.1.tar.gz
tar xzf v2.10.1.tar.gz
cd tensorflow-2.10.1

配置编译参数：

bash复制python configure.py
# 交互式选择ROCm支持
bazel build --config=rocm --config=opt //tensorflow/tools/pip_package:build_pip_package
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

3.3 依赖项打包技巧

对于常见的科学计算包，可以使用pip wheel直接打包：

bash复制pip wheel --no-deps numpy scipy pandas

对于有C扩展的包，需要确保编译环境一致：

bash复制pip wheel --build-option="--with-hip" pycocotools

4. 高级技巧与问题排查

4.1 跨系统兼容性处理

不同Linux发行版的glibc版本差异会导致兼容性问题。解决方法：

1. 在较旧系统上构建（如CentOS 7）
2. 使用manylinux标签：

bash复制docker run -v `pwd`:/io quay.io/pypa/manylinux2014_x86_64 /io/build-wheels.sh

4.2 常见编译错误解决

问题1：HIP编译器找不到

code复制error: hipcc not found

解决方法：

bash复制export PATH=/opt/rocm/bin:$PATH
export HIP_PATH=/opt/rocm/hip

问题2：内存不足

code复制virtual memory exhausted: Cannot allocate memory

调整交换空间：

bash复制sudo fallocate -l 16G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile

4.3 性能优化技巧

1. 启用ccache加速重复编译：

bash复制sudo apt install ccache
export CMAKE_CXX_COMPILER_LAUNCHER=ccache

2. 并行编译设置：

bash复制export MAKEFLAGS="-j$(nproc)"

3. 选择性编译：

bash复制export BUILD_CAFFE2_OPS=0  # 禁用不常用的算子

5. 实际应用场景

5.1 离线环境部署

将打包好的wheel文件传输到目标机器后：

bash复制pip install --no-index --find-links=/path/to/wheels torch torchvision

5.2 集群批量部署

使用Ansible剧本批量安装：

yaml复制- name: Install ROCm wheels
  hosts: compute_nodes
  tasks:
    - name: Copy wheel files
      copy:
        src: "/local/wheels/"
        dest: "/opt/wheels/"
    - name: Install packages
      pip:
        requirements: /opt/wheels/requirements.txt
        extra_args: "--no-index --find-links=/opt/wheels"

5.3 持续集成方案

GitLab CI示例配置：

yaml复制build_wheels:
  stage: build
  script:
    - docker run --rm -v $PWD:/io quay.io/pypa/manylinux2014_x86_64 /io/build.sh
  artifacts:
    paths:
      - wheelhouse/*.whl
    expire_in: 1 week

6. 维护与更新策略

6.1 版本控制建议

建议采用以下目录结构管理wheel文件：

code复制/wheels
  /rocm5.4.2
    /python3.9
      torch-1.13.1+rocm5.4.2-cp39-cp39-linux_x86_64.whl
    /python3.8
  /rocm5.3.3

6.2 自动化构建脚本

示例build.sh：

bash复制#!/bin/bash
set -ex

# 环境检查
[[ -d "/opt/rocm" ]] || { echo "ROCm not found"; exit 1; }

# 清理旧构建
rm -rf build/ dist/

# 构建核心包
for PKG in torch tensorflow; do
  git clone https://github.com/${PKG}/${PKG}
  pushd ${PKG}
    git checkout tags/v${VER}
    python setup.py bdist_wheel
    auditwheel repair dist/*.whl
    mv wheelhouse/*.whl /output/
  popd
done

6.3 依赖关系管理

使用pip-tools生成精确依赖：

bash复制pip-compile requirements.in --output-file requirements.txt

典型requirements.in内容：

code复制torch @ file:///path/to/torch-1.13.1+rocm5.4.2-cp39-cp39-linux_x86_64.whl
torchvision @ https://example.com/torchvision-0.14.1+rocm5.4.2-cp39-cp39-linux_x86_64.whl

经过半年多的实践验证，这套方案使得我们的AMD显卡开发效率提升了3倍以上。特别是在需要频繁创建临时开发环境或为集群节点部署环境时，预先构建好的wheel文件节省了大量重复编译时间。一个实际案例是：在部署20节点的训练集群时，传统方法需要每台机器单独编译耗时约2小时，而使用预构建的wheel后，整个集群的环境部署仅需15分钟。