ROCm on AMD Instinct MI355X: install and verify

#Check the GPUs and what is installed

lspci -d 1002:            # AMD GPUs on the PCI bus
lsmod | grep -w amdgpu    # is a kernel driver loaded?
dkms status               # amdgpu-dkms build state, if installed
amd-smi version           # present once ROCm is installed

The MI355X is a CDNA 4 GPU with the LLVM target gfx950. Support arrived in ROCm 7.0: older ROCm releases, and containers built on them, cannot run it.

AMD’s install guides use its own kernel driver package, amdgpu-dkms, which is versioned and tested together with ROCm. Install it rather than relying on the amdgpu module that ships with the Ubuntu kernel.

#Install the AMD GPU kernel driver

AMD now versions and documents the kernel driver separately from ROCm. Follow its Ubuntu native installation page; the repository line contains the driver version, so copy that version from the page. The steps look like this:

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> as listed on AMD's page; the codename is jammy (22.04) or 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

After the reboot, dkms status should show amdgpu as installed for the running kernel.

#Install ROCm

AMD’s ROCm installation page offers a package repository, Python wheels and tarballs, and still documents the older amdgpu-install script. ROCm ships often and package names carry the version number, so take the exact lines from that page. With the package repository on Ubuntu 24.04 (use ubuntu2204 on 22.04):

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 packages and versions on offer
sudo apt install amdrocm<version>             # all GPU targets; amdrocm<version>-gfx950 = MI355X only
sudo usermod -a -G render,video $LOGNAME      # GPU access for your user

Log out and back in so the new groups apply. If you plan to work only in containers, you can skip this step: containers bring their own ROCm, and the host needs only the kernel driver.

#Verify with rocminfo and amd-smi

groups                                          # should include render and video
rocminfo | grep -E 'Name:\s+gfx950' | wc -l     # number of MI355X GPUs ROCm can see
amd-smi list                                    # one entry per GPU, with PCI address and UUID
amd-smi static --asic                           # model and hardware details
amd-smi metric --usage                          # current utilization
amd-smi topology                                # links between the GPUs
amd-smi partition                               # compute and memory partition modes
watch -n 2 amd-smi monitor                      # power, temperature, utilization, memory

amd-smi replaces the older rocm-smi, which you may still find on older installs; the amd-smi reference lists every subcommand. MI355X also supports compute and memory partitioning, which makes one physical GPU appear as several devices. Keep the default mode unless you need it; NVIDIA’s MIG is the rough equivalent.

#PyTorch for ROCm

AMD recommends its prebuilt rocm/pytorch images as the most reliable route (see the next section). For a native install, use a PyTorch build for a ROCm release that supports gfx950: open pytorch.org, select Linux, Pip and ROCm, and copy the command. It has this form:

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 keeps the torch.cuda API on ROCm, so most code written for NVIDIA runs unchanged:

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

If gfx950 is missing from the target list, that build cannot run on MI355X: pick a newer ROCm build or use a container. AMD’s PyTorch on ROCm page lists validated images.

#ROCm containers

Containers carry their own ROCm user space; the host needs only the kernel driver. Pass the GPU device nodes to 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 is the compute interface shared by all GPUs; /dev/dri holds one render node per GPU. To give a container only some GPUs, pass single render nodes, such as --device=/dev/dri/renderD128, instead of the whole directory; ls -l /dev/dri/by-path/ shows which PCI address each node belongs to. Pick image tags built for ROCm 7.0 or later.

For LLM serving, vLLM publishes vllm/vllm-openai-rocm images (see Inference serving). Docker itself, Compose and AMD’s Container Toolkit are covered in Containers, and the details of device mapping in AMD’s Docker guide.

#Choose GPUs: HIP_VISIBLE_DEVICES and friends

HIP_VISIBLE_DEVICES=0,1 python train.py           # this process sees two GPUs, numbered 0 and 1
ROCR_VISIBLE_DEVICES=4,5,6,7 python serve.py      # applies to every ROCm library in the process

Set one of them, not several: they are applied one after another, and the numbering of the second becomes relative to the first. They are a convenience, not isolation; to separate users or jobs, use containers with specific render nodes. See AMD’s environment variables reference.

#Multi-GPU with RCCL

RCCL is AMD’s implementation of the NCCL API. PyTorch on ROCm uses it through the usual nccl backend, torchrun works unchanged, and RCCL reads the same NCCL_* variables (for example NCCL_DEBUG=INFO). The GPUs of a multi-GPU MI355X server are linked by Infinity Fabric; amd-smi topology shows the links.

git clone https://github.com/ROCm/rccl-tests.git
cd rccl-tests
make GPU_TARGETS=gfx950            # add HIP_HOME=... if ROCm is not in /opt/rocm
./build/all_reduce_perf -b 8 -e 2G -f 2 -g 8

Build rccl-tests inside a rocm/pytorch container if you did not install ROCm on the host. How to read the results, and how to launch training on several GPUs, is explained in Multi-GPU.