How to Install JAX with ROCm Acceleration on Ubuntu 24.04
Introduction
JAX is an open-source library designed for high-performance numerical computing and machine learning research. It provides tools for automatic differentiation, GPU/TPU acceleration, and just-in-time compilation to optimize code execution. ROCm, AMD's platform for GPU computing, enables JAX to utilize the power of AMD GPUs for faster computations. A ROCm-enabled JAX container is a pre-configured, portable environment that includes JAX optimized for AMD GPUs. By using this container, developers and researchers can easily leverage AMD GPUs for their JAX-based machine learning tasks without worrying about setting up dependencies or hardware configurations, streamlining their workflow.
In this article, you are to download and run a ROCm supported JAX container, and install JAX using Pip for ROCm compute platform.
Use ROCm Supported JAX Containers
In this section, you are to download and run a ROCm supported JAX container. You are to also check the GPU availability from the container.
Before proceeding with the steps in this section, ensure you have downloaded Docker with ROCm GPU Support on Ubuntu 24.04.
Pull the ROCm supported container for Jax.
CONSOLERun a temporary Docker container.
CONSOLEThe above command runs a temporary container enabling access to GPU devices (
/dev/kfdand/dev/dri) for ROCm-supported JAX workloads. To learn more about ROCm devices.Verify GPU availability from the container.
CONSOLEThe output of the above command should list all devices along with their specifications.
Exit and destroy the temporary container.
CONSOLE
Install JAX on Host using Pip
In this section, you are to install jaxlib, JAX ROCm Plugin, and JAX on the host machine using Pip and check for GPU availability.
Fetch the Python version.
CONSOLEFetch the ROCm version.
CONSOLENavigate to the JAX ROCm GitHub fork releases page.
Find the installation commands by matching the Python version you retrieved in step 1 with the ROCm version you retrieved in step 2 in the latest release notes.
Copy and execute the
jaxliband JAX ROCm Plugin installation commands from the release notes into your terminal.jaxlib: It acts as a bridge between JAX and the hardware it runs on, such as CPUs, GPUs, or TPUs. Withoutjaxlib, JAX cannot execute computations efficiently on the target hardware.- JAX ROCm Plugin: It is a specialized component that enables JAX to utilize AMD GPUs via the ROCm platform. This plugin ensures compatibility and optimization for JAX workloads on AMD hardware, enabling tasks like automatic differentiation and parallelized matrix operations to run efficiently on ROCm-supported GPUs.
Install the JAX Python package.
CONSOLEVerify GPU availability.
CONSOLEThe output of the above command should display all devices along with their respective IDs.
Conclusion
In this article, you downloaded and ran a ROCm supported JAX container with access to ROCm supported devices. Furthermore, you also installed JAX using Pip for ROCm compute platform.
0 Comments
Be the first to comment and share your perspective with the community.