DEV Community

Cover image for Installing PyTorch with ROCm Acceleration on Ubuntu 24.04
Vultr profile image Sanskriti Harmukh
Sanskriti Harmukh for Vultr

Posted on with Aashish Chaurasiya Originally published at docs.vultr.com

Installing PyTorch with ROCm Acceleration on Ubuntu 24.04

#pytorch #rocm #machinelearning #ubuntu

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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())'
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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())'
Enter fullscreen mode Exit fullscreen mode

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.

Top comments (0)