↓Skip to main content

jax-based samplers on AMD's ROCm - a recipe

·3 mins

This is a small note on how to use a GPU-based sampler on linux with an AMD GPU.

When using PyMC, it is possible to use external samplers, that increase the sampling speed. For CPU-based samplers, the best-known choice is the Rust-based nutpie.

If the underlying model has a high number of levels and/or any type of significant hierarchies, one may want to consider a GPU-based sampler. For that, the common solutions are jax-based numpyro and blackjax samplers. To properly perform them, we need a dedicated GPU. As of today, the most solid ecosystem, even on Linux, is, of course, the Nvidia’s CUDA. So the most frictionless way to make the common libraries is using the CUDA infrastructure. I am, currently, an AMD user, so making common libraries work to use GPU’s capability, is often not trivial.

My goal here is to use my newly installed AMD GPU to use efficient JAX-based samplers.
After considerable effort of googling, reading AMD’s guides, github issues and using LLM’s, I am sharing the discovered recipe to achieve what I needed!

System’s info #

====================================== Product Info ======================================
GPU[0]          : Card Series:          AMD Radeon RX 7900 XT
GPU[0]          : Card Model:           0x744c
GPU[0]          : Card Vendor:          Advanced Micro Devices, Inc. [AMD/ATI]
GPU[0]          : Card SKU:             D70401XT
GPU[0]          : Subsystem ID:         0x471e
GPU[0]          : Device Rev:           0xcc
GPU[0]          : Node ID:              1
GPU[0]          : GUID:                 37920
GPU[0]          : GFX Version:          gfx1100

============================== Version of System Component ===============================
Driver version: 6.19.6-arch1-1

and the jax version

JAX version = 0.8.2

Running the docker container #

Create the Dockerfile and build it. The Dockerfile in question, that includes the basic python packages needed for the work. Note again the usage of marimo, my default notebook, which I use for my experiments.

FROM rocm/jax:rocm7.2.4-jax0.8.2-py3.12
WORKDIR /workspace
RUN python3 -m pip install --upgrade --ignore-installed pip && \
    python3 -m pip install \
        "marimo[recommended]" \
        "pymc>=6" \
        "jax==0.8.2" \
        "jaxlib==0.8.2+rocm7.2.4" \
        numpyro \
        blackjax \
        arviz \
        pandas \
        polars \
        matplotlib \
        scipy \
        scikit-learn \
	bokeh \
	holoviews \
	hvplot \
        ipykernel \
	xarray \
	netCDF4

EXPOSE 2718
CMD ["bash", "-lc", "exec marimo edit --headless --host 0.0.0.0 --port 2718 /workspace/notebook.py"]

A script, created and adjusted with a help of an LLM, to facilitate the daily run of the container and facilitate the usage:

#!/usr/bin/env bash
set -e
IMAGE_NAME=marimo-amd-jax
CONTAINER_NAME=marimo-amd-jax
PORT=2718
docker build -t "$IMAGE_NAME" .
docker run --rm -it \
  --name "$CONTAINER_NAME" \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  -p $PORT:$PORT \
  -v "$PWD":/workspace \
  -w /workspace \
  "$IMAGE_NAME" \
  marimo edit --headless --host 0.0.0.0 --port $PORT notebook.py

Inside the container #

Within the container, we run python scripts/notebooks as usual. For my case, a necessary variables must be defined for a functional sampling.

Before jax is imported, some variables must be exported:

import os
os.environ["XLA_FLAGS"] = "--xla_gpu_enable_command_buffer=''"
os.environ["HIP_VISIBLE_DEVICES"] = "0"
os.environ["ROCR_VISIBLE_DEVICES"] = "0"
#os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "true"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.95"
os.environ["JAX_TRACEBACK_FILTERING"] = "off"
os.environ["JAX_ENABLE_X64"] = "False"
os.environ["PYTENSOR_FLAGS"] = "floatX=float32"
###
import jax
import pytensor
pytensor.config.floatX = "float32"

In order to sample using the jax sampler, we implicitly import the jax sampler’s functionality. The first option is to use the numpyro’s sampler:

from pymc.sampling.jax import sample_jax_nuts
###
idata = sample_jax_nuts(
    model=pymc_model,
    draws=1000,
    tune=1000,
    chains=4,
    nuts_sampler='numpyro',
    chain_method="vectorized",
    progressbar=True,
    random_seed=42,
)

A second option is the blackjax sampler, which still experiences issues with IO, parallelism. The working configuration with a progress tracker is the following:

idata = sample_jax_nuts(
    model=pymc_model,
    draws=1000,
    tune=1000,
    chains=1, ## note here
    nuts_sampler='blackjax',
    chain_method="vectorized",
    #chain_method="parallel",
    progressbar=True,
    random_seed=42,
)

it is possible, however, to increase the number of chains, which is solved by setting chain_method="vectorized", and higeher number of chains.

Author
Leo