#!/bin/bash -l

# SLURM SUBMIT SCRIPT
#SBATCH --job-name=rewrite-1b-25-2m
#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


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

# # activate conda env
# conda activate stable-audio-tools-env

# debugging flags (optional)
# export NCCL_DEBUG=INFO
export PYTHONFAULTHANDLER=1
export CUDA_LAUNCH_BLOCKING=0
# export NCCL_DEBUG=WARN
# export TORCH_DISTRIBUTED_DEBUG=INFO

export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12802
export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
# 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 \

# have to lower batch size to 3 for 1b model and qk_norm

# run script from above
cd /home/victor/neon/diffusion-infer
echo "Starting training..."
/home/victor/anaconda3/envs/stable-audio-tools-env/bin/torchrun \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    --nnodes=$COUNT_NODE \
    --nproc_per_node=$SLURM_GPUS_ON_NODE \
    --rdzv_id $SLURM_JOB_ID \
    --rdzv_backend c10d \
    --rdzv_endpoint "$MASTER_ADDR:$MASTER_PORT" \
    train.py