Installing JAX
This page describes how to install JAX with Python virtual environments
In this example, we will install Jax.
Step 1: Request an interactive session on a GPU node with Ampere architecture GPUs
interact -q gpu -g 1 -f ampere -m 20g -n 4
Here, -f = feature. We only need to build on Ampere once.
Step 2: Once your session has started on a compute node, run nvidia-smi to verify the GPU and then load the appropriate modules
module purge
unset LD_LIBRARY_PATH
module load cuda cudnn
Step 3: Create and activate the virtual environment
python -m venv jax.venv
source jax.venv/bin/activate
Step 4: Install the required packages
pip install --upgrade pip
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Step 5: Test that JAX is able to detect GPUs
python
>>> from jax.lib import xla_bridge
>>> print(xla_bridge.get_backend().platform)
gpu
If the above function returns gpu, then it's working correctly. You are all set, now you can install other necessary packages.
Modify batch file: See below the example batch file with the created environment
#SBATCH -J RBC
#SBATCH -N 1
#SBATCH --ntasks=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=3:30:00
#SBATCH --mem=64GB
#SBATCH --partition=gpu
#SBATCH --gres=gpu:1
#SBATCH -o RBC_job_%j.o
#SBATCH -e RBC_job_%j.e
echo $LD_LIBRARY_PATH
unset LD_LIBRARY_PATH
echo $LD_LIBRARY_PATH
source /oscar/data/gk/psaluja/jax_env.venv/bin/activate
python3 -u kernel.py