`sbatch trlx-nemo.sh`. Note that you have to edit `num_nodes` in the model yml to match the number set in the this file

trlx-nemo.sh

#!/bin/bash

#SBATCH --job-name=nemo-trlx
#SBATCH --partition=a100-cu117
#SBATCH --nodes=4
#SBATCH --gres=gpu:8
#SBATCH --ntasks-per-node=8
#SBATCH --output="%x.out"
#SBATCH --signal=B:INT@60 

# optional if cluster has it
source /opt/hpcx/hpcx-init.sh
hpcx_load

# Log env vars
export $(grep -v '^#' nemo.env | xargs)

export WANDB_API_KEY=<your wandb key>

# Python where you installed NeMo
PY=/mnt/nvme/home/uwu/conda/nemo-113/bin/python

# for debugging
# export HYDRA_FULL_ERROR=1
# export CUDA_LAUNCH_BLOCKING=1

srun echo $PY

# if using stability wandb
export WANDB_BASE_URL=https://stability.wandb.io/

# helps not OOM when close
export PYTORCH_CUDA_ALLOC_CONF='max_split_size_mb:512'

srun $PY -u /mnt/nvme/home/uwu/trlx/examples/ilql_sentiments.py