DEV Community

Cover image for Installing JAX 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 JAX with ROCm Acceleration on Ubuntu 24.04

#jax #rocm #machinelearning #ubuntu

JAX is an open-source library for high-performance numerical computing and machine learning research, offering tools for automatic differentiation, GPU/TPU acceleration, and just-in-time compilation. ROCm is AMD's platform for GPU computing, and a ROCm-enabled JAX container is a pre-configured, portable environment with JAX optimized for AMD GPUs, so you can skip manual dependency and hardware setup. This guide walks through pulling and running a ROCm-supported JAX container on a GPU-enabled server, then installing JAX directly on the host with Pip for the ROCm compute platform. By the end, you'll have JAX running with GPU acceleration on AMD hardware, verified from both a container and a native host install.

Use ROCm-Supported JAX Containers

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

1. Pull the ROCm-supported container for JAX:

$ docker pull rocm/jax: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/jax: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 JAX workloads.

3. Verify GPU availability from the container:

$ rocm-smi
$ python3 -c 'import jax; print(jax.devices())'
Enter fullscreen mode Exit fullscreen mode

The output should list all available devices along with their specifications.

4. Exit and destroy the temporary container:

$ exit
Enter fullscreen mode Exit fullscreen mode

Install JAX on the Host Using Pip

1. Check your Python version:

$ python3 -V
Enter fullscreen mode Exit fullscreen mode

2. Check your ROCm version:

$ amd-smi version
Enter fullscreen mode Exit fullscreen mode

3. Find the matching install commands: Visit the JAX ROCm GitHub fork releases page and, in the latest release notes, find the installation commands that match the Python version and ROCm version you just retrieved.

4. Install jaxlib and the JAX ROCm Plugin: Copy and run the jaxlib and JAX ROCm Plugin installation commands from the release notes in your terminal.

  • jaxlib bridges JAX to the hardware it runs on — CPU, GPU, or TPU. Without it, JAX can't execute computations efficiently on the target hardware.
  • The JAX ROCm Plugin lets JAX use AMD GPUs through the ROCm platform, enabling tasks like automatic differentiation and parallelized matrix operations to run efficiently on ROCm-supported GPUs.

5. Install the JAX Python package:

$ python3 -m pip install jax
Enter fullscreen mode Exit fullscreen mode

6. Verify GPU availability:

$ python3 -c 'import jax; print(jax.devices())'
Enter fullscreen mode Exit fullscreen mode

The output should display all devices along with their respective IDs.

Next Steps

  • Explore JAX's jit, grad, and vmap transformations to accelerate and differentiate your numerical code.
  • Benchmark training throughput on AMD GPUs against your existing CPU workflow.
  • Install PyTorch or TensorFlow with ROCm acceleration alongside JAX to compare frameworks on the same hardware.
  • Set up a Jupyter environment to prototype JAX models interactively on GPU.

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

Top comments (0)