Using JAX on the HPC

Overview

The current JAX environment on the Stats HPC is:

jax/0.11.0

It includes:

  • JAX 0.11.0

  • JAXlib 0.11.0

  • Flax 0.12.7

  • Optax 0.2.8

  • Orbax Checkpoint 0.12.1

  • Python 3.12.13

  • CUDA 12 libraries for NVIDIA GPU support

JAX provides NumPy-style numerical computing, automatic differentiation, just-in-time compilation, and GPU acceleration.

Flax provides neural-network tools built on JAX, while Optax provides optimisers and gradient transformations. Orbax Checkpoint provides model and application checkpointing.

Loading the Module

Before loading the module, initialise the module system:

source /etc/profile.d/modules.sh

Then load JAX:

module purge
module use /apps/modulefiles/Core
module load jax/0.11.0

Check the selected Python interpreter:

which python
python --version

Expected output:

/apps/conda/envs/jax-0.11.0/bin/python
Python 3.12.13

Check the installed JAX version:

python -c "import jax; print(jax.**version**)"

Note

The JAX module includes its own CUDA 12 user-space libraries. Do not load a separate CUDA module alongside jax/0.11.0.

Checking GPU Detection

GPU access should be tested within a Slurm GPU allocation.

The following command displays the devices available to JAX:

python -c "import jax; print(jax.devices())"

Inside a one-GPU allocation, output should resemble:

[CudaDevice(id=0)]

You can display additional device information with:

python - <<'PY'
import jax

device = jax.devices()[0]

print("Device:", device)
print("Platform:", device.platform)
print("Device kind:", device.device_kind)
print("Backend:", jax.default_backend())
PY

Example output:

Device: cuda:0
Platform: gpu
Device kind: NVIDIA RTX 6000 Ada Generation
Backend: gpu

Using JAX in a Slurm Job

GPU computations must be submitted to a GPU compute node through Slurm.

Example Slurm sbatch script:

#!/bin/bash

#SBATCH --job-name=jax_test
#SBATCH --clusters=srf_gpu_01
#SBATCH --partition=standard-gpu
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=1
#SBATCH --mem=4G
#SBATCH --time=00:05:00
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err

set -euo pipefail

source /etc/profile.d/modules.sh
module purge
module use /apps/modulefiles/Core
module load jax/0.11.0

echo "Node: $(hostname)"
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unset}"
echo "Python: $(command -v python)"

python -u - <<'PY'
import os

import jax
import jax.numpy as jnp

print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("JAX devices:", jax.devices())

@jax.jit
def matrix_test(x):
return x @ x

x = jnp.ones((4096, 4096), dtype=jnp.float32)
result = matrix_test(x)
result.block_until_ready()

device = jax.devices()[0]

print("Backend:", jax.default_backend())
print("Result:", float(result[0, 0]))
print("Device:", device)
print("Platform:", device.platform)
print("Device kind:", device.device_kind)
PY

Submit the job with:

sbatch jax_test.sbatch

A successful result should include:

CUDA_VISIBLE_DEVICES: 0
JAX devices: [CudaDevice(id=0)]
Backend: gpu
Result: 4096.0

Understanding CUDA_VISIBLE_DEVICES

Slurm restricts each job so that it can access only the GPU devices allocated to it.

For a one-GPU job, the environment may contain:

CUDA_VISIBLE_DEVICES=0

This means that one GPU is visible to the job. It does not mean that no GPU was allocated.

Within the job, the allocated GPU is normally presented to JAX as device number zero:

CudaDevice(id=0)

The physical GPU may have a different index on the compute node, but Slurm isolates and renumbers the devices visible to the job.

You can inspect the allocated GPU with:

nvidia-smi -L

or:

nvidia-smi \
--query-gpu=index,name,uuid,memory.total \
--format=csv

CPU and GPU Memory

The Slurm –mem option requests ordinary system RAM on the compute node. It does not request GPU video memory.

For example:

#SBATCH --mem=4G

requests 4 GiB of host memory.

Even a small JAX calculation may require more than 1 GiB of host memory while Python, JAX, CUDA libraries and the XLA compiler are initialised. The example job therefore requests 4 GiB.

Larger simulations may require substantially more memory. You should select memory according to the requirements of your application rather than copying the example value unchanged.

After a job completes, inspect its peak host-memory use with:

sacct -M srf_gpu_01 -j <job-id> \
--units=M \
--format=JobID,State,ExitCode,ReqMem,MaxRSS,MaxRSSNode

If a job exceeds its requested system memory, Slurm may report:

Detected 1 oom_kill event
OUT_OF_MEMORY

In that case, increase the –mem request and submit the job again.

GPU Memory Allocation

By default, JAX may reserve a large proportion of the visible GPU memory when the first GPU operation is run. This reduces allocation overhead and memory fragmentation.

For most jobs using an exclusively allocated GPU, the default behaviour should be left unchanged.

When troubleshooting GPU-memory errors, preallocation can be disabled:

export XLA_PYTHON_CLIENT_PREALLOCATE=false

Alternatively, the fraction of GPU memory reserved by JAX can be limited:

export XLA_PYTHON_CLIENT_MEM_FRACTION=0.50

These settings must be exported before starting Python.

Disabling preallocation can increase the risk of GPU-memory fragmentation and should not be used unless required.

Running Project Code

Copy the project source code and dependency files to the HPC, but do not copy a virtual environment which was created on a desktop or laptop.

Python virtual environments can contain:

  • Absolute paths to the original Python installation

  • Operating-system-specific compiled extensions

  • CPU-only or incompatible GPU packages

  • CUDA libraries built for a different driver or GPU generation

For a project compatible with the centrally managed environment:

module purge
module use /apps/modulefiles/Core
module load jax/0.11.0

python my_simulation.py

In a Slurm script, always load the module explicitly before running the application:

source /etc/profile.d/modules.sh
module purge
module use /apps/modulefiles/Core
module load jax/0.11.0

python -u my_simulation.py

Using python -u disables standard-output buffering, allowing progress and diagnostic messages to appear promptly in the Slurm output file.

Additional Packages

The shared JAX environment is centrally managed and should be treated as read-only.

Do not attempt to install packages directly into:

/apps/conda/envs/jax-0.11.0

Additional packages that are broadly useful may be added to the managed environment by the IT team on request.

For experimental or project-specific dependencies, use a personal environment as described in the HPC Software Modules (Python, PyTorch, Julia) guide.

Do not activate another Conda environment after loading jax/0.11.0. Doing so may replace the Python interpreter and libraries selected by the module.

To see all packages installed in the shared environment:

module load jax/0.11.0
python -m pip list

To display the principal JAX packages:

python -m pip list | grep -Ei '^(jax|jaxlib|jax-cuda|flax|optax|orbax)'

Troubleshooting

module: command not found

Initialise the module system:

source /etc/profile.d/modules.sh

Then load the shared module path:

module use /apps/modulefiles/Core

JAX does not detect a GPU

First confirm that the job requested a GPU:

#SBATCH --clusters=srf_gpu_01
#SBATCH --partition=standard-gpu
#SBATCH --gres=gpu:1

Inside the job, check:

echo "$CUDA_VISIBLE_DEVICES"
nvidia-smi -L

Then check JAX:

python -c "import jax; print(jax.devices()); print(jax.default_backend())"

A GPU job should report a CUDA device and the gpu backend.

Job killed with OUT_OF_MEMORY

This normally indicates that the job exceeded the system-memory request given by –mem.

Inspect the job with:

sacct -M srf_gpu_01 -j <job-id> \
--units=M \
--format=JobID,State,ExitCode,ReqMem,MaxRSS,MaxRSSNode

Increase the Slurm memory request as appropriate.

CUDA library errors

Start from a clean module environment:

module purge
module use /apps/modulefiles/Core
module load jax/0.11.0

Do not load a separate CUDA module. Custom CUDA settings in shell startup files, including modifications to LD_LIBRARY_PATH, may interfere with the libraries supplied by the JAX environment.

Check the loaded modules with:

module list

Wrong Python interpreter

Check:

which python

After loading the module, the result should be:

/apps/conda/envs/jax-0.11.0/bin/python

If another interpreter is shown, purge and reload the modules:

module purge
module use /apps/modulefiles/Core
module load jax/0.11.0

Best Practices

  • Always use module purge before loading the JAX environment.

  • Load jax/0.11.0 explicitly in every Slurm script.

  • Do not copy local Python virtual environments to the cluster.

  • Do not install packages into the shared /apps environment.

  • Do not load a separate CUDA module with JAX.

  • Request only the number of GPUs required by the application.

  • Treat –mem as host RAM, separate from GPU memory.

  • Use jax.devices() to confirm which devices are available.

  • Use sacct to review memory use after representative jobs.

  • Record the JAX, Flax and Optax versions used for reproducibility.

Summary

To use the centrally managed JAX environment:

source /etc/profile.d/modules.sh
module purge
module use /apps/modulefiles/Core
module load jax/0.11.0

For GPU work, include the same commands in a Slurm job that requests at least one GPU:

#SBATCH --clusters=srf_gpu_01
#SBATCH --partition=standard-gpu
#SBATCH --gres=gpu:1

Confirm GPU access with:

python -c "import jax; print(jax.devices()); print(jax.default_backend())"