A simple playground for testing Jax at Mila.
Install Python:
module load python/3.10
(add to .profile)cd ~
python -m venv python_environment
source $HOME/python_environment/bin/activate
(add to .profile)
Install cuDNN:
cd ~
tar --strip-components=1 -xf /network/scratch/b/bilaniuo/cuDNN/v8.6.0.163/cudnn-linux-x86_64-8.6.0.163_cuda11-archive.tar.xz -C ~/cudnn
export CPATH="$HOME/cudnn/include:$CPATH
(add to .profile)export LD_LIBRARY_PATH="$HOME/cudnn/lib:$LD_LIBRARY_PATH
(add to .profile)
Install poetry:
pip install poetry
alias poetry=$HOME/python_environment/bin/poetry
(add to .bashrc)
- Install this repository:
git clone https://github.com/NeilGirdhar/mila_jax.git
cd mila_jax
poetry install
Run the test:
cd mila_jax
poetry shell
sbatch job.sh