#!/bin/bash
export CUDA_LAUNCH_BLOCKING=0
export NCCL_DEBUG=WARN
export TORCH_DISTRIBUTED_DEBUG=OFF
export TORCH_CPP_LOG_LEVEL=WARNING

export OMP_NUM_THREADS=8
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12835
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER

TRAIN_PATH=/home/tony/Work/hoot/scripts
# echo working from $TRAIN_PATH
# echo MASTER_ADDR: $MASTER_ADDR
# echo MASTER_PORT: $MASTER_PORT
cd $TRAIN_PATH
pkill -f 'spawn_main'
rm -rf /mnt/localdisk/tmp_tony
mkdir -p /mnt/localdisk/tmp_tony
chmod -R 777 /mnt/localdisk/tmp_tony
# super old prod: /app/suno/checkpoints/2023-12-21_10-59-38/step_450k_ckpt.pt
# ft prod: /app/suno/checkpoints/2024-01-20_22-01-25
# dpo prod: /app/suno/checkpoints/2024-02-02_15-26-29
# v3 prod:  /app/suno/checkpoints/2024-03-19_03-11-05 
# 13b test: /app/suno/data/dpo/models/model_13b_full.pt
# 13b dpo: /app/suno/checkpoints/2024-05-16_04-33-58/last_ckpt_infer.pt
# 13b dpo launch: /app/suno/checkpoints/2024-05-27_02-51-46/last_ckpt_infer.pt
# 13b ft: /app/suno/data/dpo/models/13_ft_4head_20k.pt
# s-8: /app/suno/checkpoints/2024-06-17_01-25-46/last_ckpt_infer.pt
# https://wandb.ai/suno/chirp-v4_dev2/runs/335lzmnj/overview?nw=nwusertonytongsuno
/home/tony/anaconda3/envs/suno_env_dev/bin/python -u \
    train_audios.py \
    --out_dir=/home/tony/Data/checkpoints/hoot \
    --max_duration_s=400 \
    --max_iters=10000 \
    --warmup_iters=500 \
    --log_interval=20 \
    --eval_interval=5000 \
    --eval_iters=30 \
    --learning_rate=5e-5 \
    --min_lr=1e-6 \
    --batch_size=2 \
    --gradient_accumulation_steps=1 \
    --preload_checkpoint=/app/suno/models/hoot_v3.pt \
    --train_input_path=/home/tony/Data/Hoot/v4_t1_cer_50_long_train.json \
    --val_input_path=/home/tony/Data/Hoot/v4_t1_cer_50_long_test.json \
    --preload_decoder=False \
    --wandb_run_name=hoot_v4_test_scan_lr5e5

rm -rf /mnt/localdisk/tmp_tonyc
echo working from $TRAIN_PATH
echo DONE