PyTorch and ROCm for Deep Learning

Deploying PyTorch on a local AMD GPU system

Getting Started with PyTorch on AMD ROCm

“Training models is the process of teaching a computer program to recognize patterns in data. … The ROCm software platform makes it easier to train models on AMD GPUs while maintaining compatibility with existing code and tools. The platform also provides features like multi-GPU support, allowing for scaling and parallelization of model training across multiple GPUs to enhance performance.” – Use ROCm for Training

Many PyTorch learning resources assume a working Nvidia GPU present in the local host. In this tutorial, we’ll explain how to deploy PyTorch to an Ubuntu Linux system with an AMD GPU attached.

ROCm stands for Radeon Open Compute, the open source stack for running high performance computing (HPC) and AI workloads on AMD GPUs. The library implements Nvidia’s CUDA API for interacting with the GPU; so some documents identify the Radeon software stack as cuda-ROCm. This document assumes the ROCm libraries and GPU drivers are preinstalled.

We can confirm the system drivers are working by opening a terminal and running this ROCm utility in the Linux terminal:

rocm-smi

If this prints a table displaying the AMD GPU, our Linux system is ready. If that fails, the official AMD ROCm documentation offfers a good starting point for troubleshooting the host Linux configuration.

PyTorch comes in different installation packages depending on the GPU and GPU driver software PyTorch was compiled for. A comprehensive listing can be found at the pytorch.org download page.

System ROCm Version

ROCm is the AMD software driver and library suite for Radeon GPUs. In a Linux Ubuntu system, run any of these commands in the terminal to inspect the system-wide ROCm that is installed.

• cat /opt/rocm/.info/version (Directly reads the primary ROCm version file)

• apt show rocm-core (Checks our package manager for the installed core version)

• /opt/rocm/bin/hipconfig --version (Displays the installed HIP compiler version)

PyTorch ROCm Version

In case PyTorch has already been deployed in our Linux system, we can verify exactly which ROCm version our active PyTorch environment was compiled against by printing the version attribute of HIP, the AMD GPU C++ API and kernel language:

python -c "import torch; print(torch.version.hip)"

PyTorch Installation

1. Uninstall Current PyTorch

We first need to clear out lingering versions of the library to prevent package conflicts.

pip uninstall torch torchvision torchaudio

2. Install PyTorch for ROCm

Next we pull the version of PyTorch explicitly compiled for AMD’s ROCm backend. We can use the command below to fetch the stable ROCm release (adjust the rocm7.2 suffix to target a specific ROCm version):

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2

3. Set the RDNA 3 Override Flag

Lastly, while ROCm officially supports the flagship RDNA 3 card (for example, the RX 7900 XTX, known as gfx1100), some Radeon cards requires a lightweight spoofing flag to compile kernels correctly. To accomplish this, we prefix our terminal executions with this environment variable:

HSA_OVERRIDE_GFX_VERSION=11.0.0 python our_script.py

To avoid typing this every time, we can add import os; os.environ["HSA_OVERRIDE_GFX_VERSION"] = "11.0.0" to the very top of our Python script, strictly before we import torch. See the example code in the section below.

Verify the PyTorch Deployment

We can use the following script to confirm thecuda-ROCm stack works.

import os; os.environ["HSA_OVERRIDE_GFX_VERSION"] = "11.0.0"
import torch

print(f"GPU Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"Device Name: {torch.cuda.get_device_name(0)}")

As a penultimate test, the following code trains a single-layer neural net and runs inference (taken from Moroney, adapted to AMD GPUs).

import os; os.environ["HSA_OVERRIDE_GFX_VERSION"] = "11.0.0"
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 1. Define the target device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Running torch on {device} ...\n")

# 2. Move the model to the GPU
model = nn.Sequential(nn.Linear(1, 1)).to(device)

# Loss and optimizer (unchanged)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 3. Move the data tensors to the GPU
xs = torch.tensor([[-1.0], [0.0], [1.0], [2.0], [3.0], [4.0]], dtype=torch.float32).to(device)
ys = torch.tensor([[-3.0], [-1.0], [1.0], [3.0], [5.0], [7.0]], dtype=torch.float32).to(device)

# Train (unchanged, operations now happen on the device where tensors live)
for _ in range(500):
    optimizer.zero_grad()
    outputs = model(xs)
    loss = criterion(outputs, ys)
    loss.backward()
    optimizer.step()

# 4. Move new prediction inputs to the GPU
with torch.no_grad():
    test_val = torch.tensor([[10.0]], dtype=torch.float32).to(device)
    print(model(test_val))