#!/bin/bash
#SBATCH --job-name=gpt-sft
#SBATCH --nodes=4
#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-[633]
##SBATCH --exclude=h100-ord01-03-[]

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

# 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_23_$USER

export COMPILE_MERT=0

export CUDA_LAUNCH_BLOCKING=1
# # NCCL: only shout when something's wrong; still fail fast
# export NCCL_DEBUG=WARN                 # INFO spams per-collective; WARN is quiet
# export NCCL_DEBUG_SUBSYS=INIT          # INIT only; skip COLL/P2P chatter
# export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
# export TORCH_NCCL_BLOCKING_WAIT=1
# export NCCL_TIMEOUT=600

# # PyTorch distributed: minimal breadcrumbs, full stack only on crash
# export TORCH_SHOW_CPP_STACKTRACES=1    # only prints on error

# # Allocator: no periodic dumps, but better chance to avoid fragmentation
# export PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6,max_split_size_mb:512


# 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

echo "Starting SFT train script..."
srun -K1 /home/victor/anaconda3/envs/suno_env/bin/python -u train.py \
    --debug_val_only=True \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    \
    --out_dir="/app2/suno/checkpoints" \
    --data_dir="/app2/suno/data/auk_v0" \
    --train_metas_filename="metas_v7_tr.jsonl" \
    --train_info_filename="ids_keep_sets_ext_v11_stems.json" \
    --val_metas_filename="metas_v7_val.jsonl" \
    --preload_strict=False \
    \
    --n_layer=8 \
    --n_head=32 \
    --d_head=128 \
    \
    --learning_rate=5e-5 \
    --max_iters=10_000 \
    --warmup_iters=1_000 \
    --eval_iters=1 \
    --batch_store_size=2 \
    --step_save_iters=5_000 \
    --allow_sample=True \
    --allow_vox=True \
    \
    --fsdp=True \
    --grad_checkpointing=True \
    --compile=False \
    \
    --wandb_log=False \
    --wandb_dir="/app2/suno/victor/wandb_logs" \
    --wandb_project="gpt-45-blocks" \
    --wandb_run_name="sft-v7" \
    \
    --local_cache_dir="/mnt/localdisk/tmp1/victor"