Installing PyTorch with ROCm Acceleration on Ubuntu 24.04

PyTorch is an open-source framework for building and training machine learning models, especially deep learning networks. ROCm is AMD's platform for GPU computing, enabling high-performance computing and machine learning on AMD GPUs. A ROCm-enabled PyTorch container is a pre-built, portable environment with PyTorch configured to run efficiently on AMD GPUs, so you can skip manually setting up or configuring dependencies. This guide walks through pulling and running a ROCm-supported PyTorch container on a GPU-enabled server, then installing PyTorch directly on the host with Pip for the ROCm compute platform. By the end, you'll have PyTorch running with GPU acceleration on AMD hardware, verified from both a container and a native host install.

Use ROCm-Supported PyTorch Containers

Before you begin, make sure Docker is installed on your server with ROCm GPU support enabled.

1. Pull the ROCm-supported container for PyTorch:

$ docker pull rocm/pytorch:latest

2. Run a temporary Docker container:

$ docker run --rm -it --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --shm-size 8G rocm/pytorch:latest

This command runs a temporary container with access to the GPU devices (/dev/kfd and /dev/dri) needed for ROCm-supported PyTorch workloads.

3. Verify GPU availability from the container:

$ rocm-smi
$ python3 -c 'import torch; print(torch.cuda.is_available())'

The output should list all available devices along with their specifications, and print true if the GPU is available.

4. Exit and destroy the temporary container:

$ exit

Install PyTorch on the Host Using Pip

1. Select your install options: Visit the PyTorch Start Locally page and choose PyTorch Build: Stable, Your OS: Linux, Package: Pip, Language: Python, Compute Platform: ROCm. Copy the command the page generates and run it in your terminal.

Note: If you encounter a timeout error while downloading the PyTorch package, add the --timeout flag with an appropriate value to the installation command. Alternatively, use the --no-cache-dir flag to restart the installation from scratch, ignoring any incomplete or corrupted cached downloads.

2. Verify GPU availability:

$ python3 -c 'import torch; print(torch.cuda.is_available())'

This should print true if the GPU is available.

Next Steps

• Benchmark training throughput on AMD GPUs against your existing CPU workflow.
• Install JAX or TensorFlow with ROCm acceleration alongside PyTorch to compare frameworks on the same hardware.
• Set up a Jupyter environment to prototype PyTorch models interactively on GPU.
• Containerize your training scripts on top of the rocm/pytorch image for repeatable runs.

For the full guide with additional tips, visit the original article on Vultr Docs.