.. _jax:

####################
Using JAX on the HPC
####################

**********
Overview
**********


The current JAX environment on the Stats HPC is:

.. code-block:: text
   
    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:

.. code-block:: bash

   source /etc/profile.d/modules.sh

Then load JAX:

.. code-block:: bash

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

Check the selected Python interpreter:

.. code-block:: bash

   which python
   python --version

Expected output:

.. code-block:: text

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

Check the installed JAX version:

.. code-block:: bash

   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:

.. code-block:: bash

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

Inside a one-GPU allocation, output should resemble:

.. code-block:: text

   [CudaDevice(id=0)]

You can display additional device information with:

.. code-block:: bash

   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:

.. code-block:: text

   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:

.. code-block:: bash

   #!/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:

.. code-block:: bash

   sbatch jax_test.sbatch

A successful result should include:

.. code-block:: text

   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:

.. code-block:: text

   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:

.. code-block:: text

   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:

.. code-block:: bash

   nvidia-smi -L

or:

.. code-block:: bash

   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:

.. code-block:: bash

   #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:

.. code-block:: bash

   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:

.. code-block:: text

   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:

.. code-block:: bash

   export XLA_PYTHON_CLIENT_PREALLOCATE=false

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

.. code-block:: bash

   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:

.. code-block:: bash

   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:

.. code-block:: bash

   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:

.. code-block:: text

   /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 :ref:`modules` 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:

.. code-block:: bash

   module load jax/0.11.0
   python -m pip list

To display the principal JAX packages:

.. code-block:: bash

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

****************
Troubleshooting
****************

`module: command not found`
=============================

Initialise the module system:

.. code-block:: bash

   source /etc/profile.d/modules.sh

Then load the shared module path:

.. code-block:: bash

   module use /apps/modulefiles/Core

JAX does not detect a GPU
===========================

First confirm that the job requested a GPU:

.. code-block:: bash

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

Inside the job, check:

.. code-block:: bash

   echo "$CUDA_VISIBLE_DEVICES"
   nvidia-smi -L

Then check JAX:

.. code-block:: bash

   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:

.. code-block:: bash

   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:

.. code-block:: bash

   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:

.. code-block:: bash

   module list

Wrong Python interpreter
==========================

Check:

.. code-block:: bash

   which python

After loading the module, the result should be:

.. code-block:: text

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

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

.. code-block:: bash

   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:

.. code-block:: bash

   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:

.. code-block:: bash

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

Confirm GPU access with:

.. code-block:: bash

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