在 AMD Instinct MI355X 上安装并验证 ROCm

检查 GPU 与已安装的组件

lspci -d 1002:            # PCI 总线上的 AMD GPU
lsmod | grep -w amdgpu    # 是否已加载内核驱动？
dkms status               # amdgpu-dkms 的构建状态（如果已安装）
amd-smi version           # 安装 ROCm 后才会出现

MI355X 是一款 CDNA 4 GPU，LLVM 目标架构为 gfx950。相关支持是从 ROCm 7.0 开始加入的：更早的 ROCm 版本，以及基于它们构建的容器，都无法运行在这款 GPU 上。

AMD 的安装指南使用它自己的内核驱动包 amdgpu-dkms，该驱动包会与 ROCm 一起进行版本管理和测试。请安装这个驱动包，而不要依赖 Ubuntu 内核自带的 amdgpu 模块。

安装 AMD GPU 内核驱动

AMD 现在把内核驱动和 ROCm 分开进行版本管理和文档编写。请参照它的 Ubuntu 原生安装页面操作；仓库那一行包含驱动版本号，请从页面上复制该版本号。步骤大致如下：

sudo apt update
sudo apt install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
sudo mkdir --parents --mode=0755 /etc/apt/keyrings
wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null

# <driver-version> 请使用 AMD 页面上列出的版本；代号在 22.04 上是 jammy，在 24.04 上是 noble
echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/amdgpu/<driver-version>/ubuntu $(. /etc/os-release; echo $VERSION_CODENAME) main" \
  | sudo tee /etc/apt/sources.list.d/amdgpu.list
sudo apt update
sudo apt install amdgpu-dkms
sudo reboot

重启之后，dkms status 应该会显示当前运行的内核已安装 amdgpu。

安装 ROCm

AMD 的 ROCm 安装页面提供了软件包仓库、Python wheel 包和 tarball，也仍然记录着较旧的 amdgpu-install 脚本。ROCm 发布很频繁，包名中带有版本号，因此请直接从该页面复制准确的命令行。以下以 Ubuntu 24.04 上的软件包仓库方式为例（22.04 上请使用 ubuntu2204）：

wget https://stable.repo.amd.com/rocm/gpg/packages.gpg -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/amdrocm.gpg > /dev/null
sudo tee /etc/apt/sources.list.d/amdrocm-stable.sources > /dev/null <<'EOF'
Types: deb
URIs: https://stable.repo.amd.com/rocm/core/packages/ubuntu2404/
Suites: stable
Components: main
Architectures: amd64
Signed-By: /etc/apt/keyrings/amdrocm.gpg
EOF
sudo apt update
apt-cache search --names-only '^amdrocm'      # 可选的 ROCm 软件包及版本
sudo apt install amdrocm<version>             # 全部 GPU 目标架构；amdrocm<version>-gfx950 = 仅限 MI355X
sudo usermod -a -G render,video $LOGNAME      # 让您的用户获得 GPU 访问权限

注销后重新登录，让新的用户组生效。如果您只打算在容器中工作，可以跳过这一步：容器会自带 ROCm，主机只需要内核驱动。

用 rocminfo 和 amd-smi 验证

groups                                          # 应该包含 render 和 video
rocminfo | grep -E 'Name:\s+gfx950' | wc -l     # ROCm 能看到的 MI355X GPU 数量
amd-smi list                                    # 每块 GPU 一条记录，包含 PCI 地址和 UUID
amd-smi static --asic                           # 型号和硬件详情
amd-smi metric --usage                          # 当前利用率
amd-smi topology                                # GPU 之间的连接
amd-smi partition                               # 计算与显存分区模式
watch -n 2 amd-smi monitor                      # 功耗、温度、利用率、显存

amd-smi 取代了较旧的 rocm-smi，在一些较旧的安装中您可能仍会见到它；amd-smi 参考文档列出了它的每个子命令。MI355X 还支持计算与显存分区，可以让一块物理 GPU 表现为多个设备。除非您确实需要，否则请保持默认模式；NVIDIA 的 MIG 大致相当于这一功能。

面向 ROCm 的 PyTorch

AMD 建议使用它预先构建好的 rocm/pytorch 镜像，这是最可靠的方式（见下一节）。如果要原生安装，请选择支持 gfx950 的 ROCm 版本所对应的 PyTorch 构建：打开 pytorch.org，选择 Linux、Pip 和 ROCm，然后复制命令。命令的形式大致如下：

sudo apt install -y python3-venv
python3 -m venv ~/venvs/torch
source ~/venvs/torch/bin/activate
pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm<X.Y>

PyTorch 在 ROCm 上保留了 torch.cuda 接口，因此大多数为 NVIDIA 编写的代码无需修改即可运行：

python3 - <<'EOF'
import torch
print("HIP runtime:", torch.version.hip)
print("GPUs visible:", torch.cuda.device_count())
print("GPU 0:", torch.cuda.get_device_name(0))
print("gfx targets in this build:", torch.cuda.get_arch_list())
x = torch.randn(4096, 4096, device="cuda")
print("matmul OK:", (x @ x).shape)
EOF

如果目标列表中没有 gfx950，说明这个版本无法在 MI355X 上运行：请换一个更新的 ROCm 版本，或使用容器。AMD 的面向 ROCm 的 PyTorch 页面列出了已验证可用的镜像。

ROCm 容器

容器会自带一套完整的 ROCm 用户空间；主机只需要内核驱动。把 GPU 设备节点传给 Docker：

docker run -it --rm \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video --ipc=host \
  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
  -v /data:/data \
  rocm/pytorch:latest

/dev/kfd 是所有 GPU 共享的计算接口；/dev/dri 中每块 GPU 对应一个渲染节点。如果只想把部分 GPU 分配给某个容器，可以只传递单个渲染节点，例如 --device=/dev/dri/renderD128，而不是整个目录；ls -l /dev/dri/by-path/ 可以显示每个节点对应的 PCI 地址。请选择基于 ROCm 7.0 或更高版本构建的镜像标签。

如果要做 LLM 推理服务，vLLM 提供了 vllm/vllm-openai-rocm 镜像（参见推理服务）。Docker 本身、Compose 以及 AMD 的 Container Toolkit 在容器中有介绍，设备映射的细节则在 AMD 的 Docker 指南中说明。

选择 GPU：HIP_VISIBLE_DEVICES 及相关变量

HIP_VISIBLE_DEVICES=0,1 python train.py           # 该进程只能看到编号为 0 和 1 的两块 GPU
ROCR_VISIBLE_DEVICES=4,5,6,7 python serve.py      # 作用于该进程中的每一个 ROCm 库

请只设置其中一个，不要同时设置多个：它们会依次生效，后一个变量的编号会相对于前一个变量重新计算。它们只是一种便利手段，并不能起到隔离作用；如果要区分不同用户或任务，请使用带有特定渲染节点的容器。参见 AMD 的环境变量参考文档。

使用 RCCL 实现多 GPU

RCCL 是 AMD 对 NCCL API 的实现。ROCm 上的 PyTorch 通过常规的 nccl 后端使用它，torchrun 的用法不变，RCCL 读取的也是同样的 NCCL_* 变量（例如 NCCL_DEBUG=INFO）。多 GPU MI355X 服务器上的各块 GPU 由 Infinity Fabric 连接；amd-smi topology 可以显示这些连接。

git clone https://github.com/ROCm/rccl-tests.git
cd rccl-tests
make GPU_TARGETS=gfx950            # 如果 ROCm 不在 /opt/rocm，请添加 HIP_HOME=...
./build/all_reduce_perf -b 8 -e 2G -f 2 -g 8

如果您没有在主机上安装 ROCm，可以在 rocm/pytorch 容器内构建 rccl-tests。如何解读结果，以及如何在多块 GPU 上启动训练，请参见多 GPU。