#!/bin/bash
#SBATCH --job-name=gpt_cover_ft
#SBATCH --nodes=4
#SBATCH --gres=gpu:8
#SBATCH --ntasks-per-node=8  # match this to n_gpu if python
#SBATCH --cpus-per-task=4  # n_cpu*n_task has to be <= cpus per node (~64)
#SBATCH --output=/app/suno/slurm/logs/run_%j.txt
#SBATCH --error=/app/suno/slurm/logs/run_%j_err.txt

# other useful slurm commands
# --exclude=h100-ord01-03-[]
# --nodelist=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_$USER

# using bare python rather than torchrun prevents hanging on nightly
# -K1 means slurm job will crash if run crashes rather than just hang
# -u means python will stream stoud as it comes rather than buffer

#--artist_condition=False \
#--cover_condition=True \
#--overpaint_condition=False \
#--underpaint_condition=False \

echo "Starting train script..."
srun -K1 /home/georg/anaconda3/envs/gpt_n/bin/python -u train.py \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    --is_finetune=True \
    --out_dir="/app/suno/checkpoints" \
    --data_dir="/app/suno/sara/cover_filter_v0" \
    --train_info_filename="info_tr_0.4_0.8_w_def_all.json" \
    --train_metas_filename="metas_tr.jsonl" \
    --val_info_filename="info_val_0.4_0.8_w_def_all.json" \
    --val_metas_filename="metas_val.jsonl" \
    --step_save_iters=2_000 \
    --eval_interval=500 \
    --weights_multiplier=None \
    --n_layer=32 \
    --n_head=32 \
    --d_head=128 \
    --learning_rate=2e-5 \
    --max_iters=8_370 \
    --warmup_iters=840 \
    --fsdp=True \
    --grad_checkpointing=True \
    --compile=False \
    --wandb_log=True \
    --wandb_dir="/app/suno/sara/wandb_logs" \
    --wandb_project="gpt-cover-ft" \
    --wandb_run_name="cover-ft-4-to-9_sft_w_def_2e5_ep2" \
    --preload_checkpoint="/app/suno/checkpoints/2025-03-24_20-05-34/last_ckpt_infer.pt" \
    --preload_optimizer=False \
    --local_cache_dir="/mnt/localdisk/tmp/sma"