#!/bin/bash -l

# SLURM SUBMIT SCRIPT
#SBATCH --job-name=1b-100-30s
#SBATCH --nodes=1            # This needs to match Trainer(num_nodes=...)
#SBATCH --gres=gpu:8
#SBATCH --ntasks-per-node=8   # This needs to match Trainer(devices=...)
#SBATCH --cpus-per-task=8    # Number of cores per tasks


# --nodelist=compute-h100-ord-node-[]

# ACTIVATE ANACONDA
eval "$(conda shell.bash hook)"

# activate conda env
conda activate suno_env

# debugging flags (optional)
export NCCL_DEBUG=INFO
export PYTHONFAULTHANDLER=1

# on your cluster you might need these:
# set the network interface
# export NCCL_SOCKET_IFNAME=^docker0,lo

# might need the latest CUDA
#module load NCCL/2.4.7-1-cuda.10.0

# --local-data-shard-dir /mnt/localdisk/cjs_shards \

# sm 200m model batch size 28
# lg 1b model batch size 4
# lg 1b model qk_norm batch size 3
# 4b model qk_norm batch size 4 with checkpointing layers 

#

# run script from above
srun python3 ./train.py \
--dataset-config stable_audio_tools/configs/dataset_configs/memmap_semantic+text_vae_100hz_30s.json \
--model-config stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_infill.json \
--name harmonai_train \
--checkpoint-every 20000 \
--pretrained-ckpt-path /home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_1b_full_scale=2.5_epoch=15-step=2360000.ckpt \
--num-gpus 8 \
--num-nodes 1 \
--batch-size 6