#!/bin/bash
#SBATCH --job-name=dac-large
#SBATCH --nodes=1
#SBATCH --gres=gpu:8
#SBATCH --ntasks-per-node=8  # match this to n_gpu if python
#SBATCH --cpus-per-task=14  # n_cpu*n_task has to be <= cpus per node (~64)
#SBATCH --mem=0


# other useful slurm commands
# SBATCH --nodelist=h100-ord01-03-[]
##SBATCH --exclude=h100-ord01-03-[]

# this is fairly arbity
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12885

# cluster specific defaults
export OMP_NUM_THREADS=1
export NCCL_CROSS_NIC=2

# explicit cache dirs to not get user conflicts
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_1_$USER

export CUDA_LAUNCH_BLOCKING=0

echo "Starting train script..."
srun -K1 /home/minz/anaconda3/envs/nenv/bin/python -u train.py \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    \
    --out_dir="/app2/suno/checkpoints" \
    --data_dir="/app2/suno/data/diffusion/v1" \
    --train_metas_filename="metas_diff_v4_val.jsonl" \
    --val_metas_filename="metas_diff_v4_val.jsonl" \
    --preload_strict=False \
    \
    --batch_size=19 \
    \
    --encoder_dim=128 \
    --vae_dim=128 \
    --decoder_dim=1536 \
    \
    --weight_mel_loss=15.0 \
    --weight_kl_loss=0.0001 \
    --weight_feat_loss=2.0 \
    --weight_adv_loss=1.0 \
    --weight_disc_loss=1.0 \
    \
    --lr_scheduler_type="exponential" \
    --learning_rate_codec=1e-4 \
    --learning_rate_disc=1e-4 \
    --max_iters=500_000 \
    --eval_iters=20 \
    \
    --fsdp=False \
    --grad_checkpointing=True \
    --compile=False \
    \
    --wandb_log=True \
    --wandb_dir="/app/suno/minz/wandb_logs" \
    --wandb_project="codec-2025q4" \
    --wandb_run_name="dac-large" \
    \
    --local_cache_dir="/mnt/localdisk/tmp/minz"