[For Brown Community] Setting up HSSM on Oscar

[digicosmos86]

This is a tutorial on how to get HSSM to work on Oscar, Brown University's HPC cluster. This may or may not work with other HPC environments, depending on the setup.

You can choose to work with HSSM either in a venv or a conda environment. The former option, if linked correctly with the Intel MKL library on Oscar, can lead to a significant performance boost on Oscar.

Option 1: Using HSSM in venv (Faster)

In general, we recommend that we work with environment created with venv on Oscar and install packages with pip. However, the set up for PyMC is slightly more involved because we need to tell PyMC where it can find the BLAS library to accelerate its computation. Fortunately, we only need to do this once, unless the path to the library changes. When that happens, we just need to run this process again.

Create a .pytensorrc file with path to a BLAS library

The BLAS library that seems to work well with pytensor is the Intel one MKL library. We need to load the intel-oneapi-mkl module and tell pymc how to link to that library:

module purge
unset LD_LIBRARY_PATH
module load intel-oneapi-mkl
echo $LD_LIBRARY_PATH

You will see a few paths printed out separated with colons. Copy the path that ends in intel64 after running the above command. Then, create a .pytensorrc file under your home directory ~, and add the following content:

[blas]
ldflags=-Lpath_just_copied -Wl,--no-as-needed -lmkl_gf_lp64 -Wl,--start-group -lmkl_gnu_thread -lmkl_core  -Wl,--end-group -fopenmp  -ldl -lpthread -lm

Replace path_just_copied with what you have just copied. Be sure to keep the -L before the path. This will ensure that pymc finds the lapack library. You only need to do this once unless the LAPACK module gets changed. The file will look like this:

[blas]
ldflags=-L/oscar/rt/9.2/software/0.20-generic/0.20.1/opt/spack/linux-rhel9-x86_64_v3/gcc-11.3.1/intel-oneapi-mkl-2023.1.0-xbcd2g3myicmoznxp4pajepy7y2xzsvh/mkl/2023.1.0/lib/intel64 -Wl,--no-as-needed -lmkl_gf_lp64 -Wl,--start-group -lmkl_gnu_thread -lmkl_core  -Wl,--end-group -fopenmp  -ldl -lpthread -lm

Load modules and create environment

We need the python module to create the virtual environment and activate it. You can do this in the login node.

module load python intel-oneapi-mkl
python -m venv .venv
source .venv/bin/activate

Install HSSM

Once the environment is loaded, we can install hssm with pip:

pip install hssm

Or you can install the dev version directly from github:

pip install git+https://github.com/lnccbrown/HSSM.git

The installation is complete. Please ensure that there is no warning like this: WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions. when hssm is imported.

Optional: Install JAX with GPU support

If you need JAX with GPU support, please perform the installation in an interactive GPU node:

interact -q gpu -g 1 -f ampere -m 20g -n 4

Once the node is up, run the following commands to set up the virtual environment. Note that you also need to load the cuda and cudnn modules before proceeding:

module load python netlib-lapack cuda cudnn
python -m venv .venv
source .venv/bin/activate

Then, install the GPU version of JAX:

pip install jax[cuda12]

Finally, you can choose to install the release or dev versions of HSSM.

Option 2: Using HSSM in a conda environment:

The use of conda is discouraged on Oscar. It might be slow in resolving dependencies. If you have to use conda, we recommend using mamba in place of conda for considerable speed-up.

With the conda environment, you do NOT need to set up the .pytensorrc file. In fact, it's not a good idea to have that file in your home directory.

Load modules and create environment

We need the miniforge module to create the virtual environment and activate it. It comes with mamba installed. You can also use the anaconda module, but you won't be able to use mamba with it.

module load miniforge
mamba create -n my_env python=3.11 #3.10 also works
mamba activate my_env

Install HSSM

HSSM is directly available on conda-forge now, so you can directly install it:

mamba install -c conda-forge hssm

If you need the dev version of HSSM, then you need to install pymc from conda first, before installing the dev version from github with pip:

mamba install -c conda-forge pymc
pip install git+https://github.com/lnccbrown/HSSM.git

Optional: Install JAX with GPU support

If you need JAX with GPU support, please perform the installation in an interactive GPU node:

interact -q gpu -g 1 -f ampere -m 20g -n 4

Once the node is up, run the following commands to set up the conda environment. Note that you also need to load the cuda and cudnn modules before proceeding:

module load miniforge cuda cudnn
mamba create -n my_env python=3.11 #3.10 also works
mamba activate my_env

Then, install the GPU version of JAX:

pip install jax[cuda12]

Finally, follow the instruction in the previous section to install either the release or the dev version of HSSM.

[igrahek]

Getting this error on Oscar:

[igrahek@login009 ~]$ module load intel-one-mkl
Lmod has detected the following error:  The following module(s) are unknown: "intel-one-mkl"
Please check the spelling or version number. Also try "module spider ..."
It is also possible your cache file is out-of-date; it may help to try:
  $ module --ignore_cache load "intel-one-mkl"
Also make sure that all modulefiles written in TCL start with the string #%Module

[digicosmos86]

Sorry the module should be intel-oneapi-mkl. I've updated the document