ROCm на AMD Instinct MI355X: установка и проверка

Проверьте GPU и то, что уже установлено

lspci -d 1002:            # GPU AMD на шине PCI
lsmod | grep -w amdgpu    # загружен ли драйвер ядра?
dkms status               # состояние сборки amdgpu-dkms, если он установлен
amd-smi version           # появляется после установки ROCm

MI355X — это GPU архитектуры CDNA 4 с LLVM-таргетом gfx950. Поддержка появилась в ROCm 7.0: более старые версии ROCm и контейнеры, собранные на их основе, не могут его запустить.

В руководствах AMD по установке используется собственный пакет драйвера ядра, amdgpu-dkms, который версионируется и тестируется вместе с ROCm. Устанавливайте именно его, а не полагайтесь на модуль amdgpu, поставляемый с ядром Ubuntu.

Установите драйвер ядра для GPU AMD

Теперь 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; кодовое имя дистрибутива — jammy (22.04) или noble (24.04)
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

Страница установки ROCm от AMD предлагает пакетный репозиторий, готовые пакеты для Python и архивы для сборки, а также по-прежнему описывает старый скрипт 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     # сколько GPU MI355X видит ROCm
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 выглядит как несколько устройств. Оставляйте режим по умолчанию, если он вам не нужен; примерным аналогом служит MIG от NVIDIA.

PyTorch для ROCm

AMD рекомендует свои готовые образы rocm/pytorch как самый надёжный вариант (см. следующий раздел). Для установки непосредственно на систему используйте сборку PyTorch для версии ROCm, поддерживающей gfx950: откройте 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 сохраняет тот же API torch.cuda и на ROCm, поэтому большая часть кода, написанного для 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 PyTorch на ROCm содержит список проверенных образов.

Контейнеры ROCm

Контейнеры несут собственное пользовательское окружение ROCm; хосту нужен только драйвер ядра. Передайте Docker узлы устройств GPU:

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 содержит по одному render-узлу на каждый GPU. Чтобы выделить контейнеру только некоторые GPU, передавайте отдельные render-узлы, например --device=/dev/dri/renderD128, вместо всего каталога; команда ls -l /dev/dri/by-path/ показывает, какому PCI-адресу принадлежит каждый узел. Выбирайте теги образов, собранные для ROCm 7.0 или новее.

Для инференса LLM vLLM публикует образы vllm/vllm-openai-rocm (см. Инференс-сервер). Docker, Compose и Container Toolkit от AMD описаны в Контейнеры, а подробности о маппинге устройств — в руководстве по Docker от AMD.

Выбор GPU: HIP_VISIBLE_DEVICES и другие переменные

HIP_VISIBLE_DEVICES=0,1 python train.py           # этот процесс видит два GPU, с номерами 0 и 1
ROCR_VISIBLE_DEVICES=4,5,6,7 python serve.py      # применяется ко всем библиотекам ROCm в процессе

Задавайте только одну из них, не несколько: они применяются последовательно, и нумерация во второй переменной становится относительной к первой. Это удобство, а не изоляция; чтобы разделить пользователей или задачи, используйте контейнеры с отдельными render-узлами. См. справочник AMD по переменным окружения.

Multi-GPU с RCCL

RCCL — это реализация API NCCL от AMD. PyTorch на ROCm использует её через обычный бэкенд nccl, torchrun работает без изменений, а RCCL читает те же переменные NCCL_* (например, NCCL_DEBUG=INFO). GPU в составе multi-GPU-сервера MI355X связаны через Infinity Fabric; amd-smi topology показывает эти связи.

git clone https://github.com/ROCm/rccl-tests.git
cd rccl-tests
make GPU_TARGETS=gfx950            # добавьте HIP_HOME=..., если ROCm не в /opt/rocm
./build/all_reduce_perf -b 8 -e 2G -f 2 -g 8

Соберите rccl-tests внутри контейнера rocm/pytorch, если вы не устанавливали ROCm на хост. Как читать результаты и как запускать обучение на нескольких GPU, объясняется в разделе Multi-GPU.