#!/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
export SLURM_NTASKS_PER_NODE=8

TRAIN_PATH=/home/tony/Work/neon/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
# hoot v5 pretrain: /home/tony/Data/checkpoints/hoot/2025-02-11_05-03-34/last_ckpt.pt
# hoot v5 ft: /home/tony/Data/checkpoints/hoot/2025-02-13_05-07-20/last_ckpt.pt
# https://wandb.ai/suno/chirp-v4_dev2/runs/335lzmnj/overview?nw=nwusertonytongsuno
# # pretrain:
# /home/tony/anaconda3/envs/suno_12/bin/python -u \
#     train_audios.py \
#     --out_dir=/home/tony/Data/checkpoints/hoot \
#     --max_duration_s=280 \
#     --max_iters=250_010 \
#     --warmup_iters=5_000 \
#     --log_interval=25 \
#     --eval_interval=5_000 \
#     --eval_iters=50 \
#     --learning_rate=1e-3 \
#     --min_lr=1e-6 \
#     --batch_size=4 \
#     --decoder_type=v2 \
#     --gradient_accumulation_steps=1 \
#     --preload_checkpoint=/app/suno/models/hoot_v3.pt \
#     --train_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_train.json \
#     --val_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_test.json \
#     --preload_decoder=False \
#     --lyrics_key=text \
#     --wandb_run_name=hoot_v6_v2_n16

# # fine tuning baby is the key
# /home/tony/anaconda3/envs/suno_12/bin/python -u \
#     train_audios.py \
#     --out_dir=/home/tony/Data/checkpoints/hoot \
#     --max_duration_s=280 \
#     --max_iters=50_010 \
#     --warmup_iters=200 \
#     --log_interval=25 \
#     --eval_interval=5_000 \
#     --eval_iters=50 \
#     --learning_rate=1e-6 \
#     --min_lr=1e-8 \
#     --batch_size=4 \
#     --decoder_type=v2 \
#     --gradient_accumulation_steps=1 \
#     --preload_checkpoint=/home/tony/Data/checkpoints/hoot/2025-03-19_04-47-27/last_ckpt.pt \
#     --train_input_path=/home/tony/Data/Hoot/v6_t1_cer_30_long_train.json \
#     --val_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_test.json \
#     --preload_decoder=True \
#     --custom_seed_offset=5678 \
#     --lyrics_key=text \
#     --wandb_run_name=hoot_v6_v2_n16_sft_2

# finish_pretrain:
# /home/tony/anaconda3/envs/suno_12/bin/python -u \
#     train_audios.py \
#     --out_dir=/home/tony/Data/checkpoints/hoot \
#     --max_duration_s=290 \
#     --max_iters=1_000 \
#     --warmup_iters=50 \
#     --log_interval=25 \
#     --eval_interval=1_000 \
#     --eval_iters=50 \
#     --learning_rate=1e-4 \
#     --min_lr=1e-5 \
#     --batch_size=4 \
#     --decoder_type=v2 \
#     --gradient_accumulation_steps=1 \
#     --preload_checkpoint=/home/tony/Data/checkpoints/hoot/2025-03-10_12-08-54/40k_ckpt.pt \
#     --train_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_test.json \
#     --val_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_test.json \
#     --preload_decoder=False \
#     --custom_seed_offset=5678 \
#     --preload_optimizer=False \
#     --lyrics_key=text \
#     --wandb_run_name=hoot_v6_v1_n2_test

# pretrain:
# /home/tony/anaconda3/envs/suno_12/bin/python -u \
#     train_audios.py \
#     --out_dir=/home/tony/Data/checkpoints/hoot \
#     --max_duration_s=280 \
#     --max_iters=250_010 \
#     --warmup_iters=20_000 \
#     --log_interval=25 \
#     --eval_interval=5_000 \
#     --eval_iters=50 \
#     --learning_rate=3e-4 \
#     --min_lr=3e-5 \
#     --batch_size=2 \
#     --decoder_type=v2 \
#     --gradient_accumulation_steps=1 \
#     --n_layers=24 \
#     --n_embd=1024 \
#     --preload_checkpoint=/home/tony/Data/Hoot/model/stt_en_fastconformer_ctc_xlarge.pt \
#     --train_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_train.json \
#     --val_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_test.json \
#     --preload_decoder=False \
#     --lyrics_key=text \
#     --wandb_run_name=hoot_v6_v2_n16_xlarge

# # run crashed at 230k, extend training
# /home/tony/anaconda3/envs/suno_12/bin/python -u \
#     train_audios.py \
#     --out_dir=/home/tony/Data/checkpoints/hoot \
#     --max_duration_s=280 \
#     --max_iters=100_010 \
#     --warmup_iters=1 \
#     --log_interval=25 \
#     --eval_interval=5_000 \
#     --eval_iters=50 \
#     --learning_rate=3e-5 \
#     --min_lr=1e-5 \
#     --batch_size=2 \
#     --decoder_type=v2 \
#     --gradient_accumulation_steps=1 \
#     --n_layers=24 \
#     --n_embd=1024 \
#     --preload_checkpoint=/home/tony/Data/checkpoints/hoot/2025-06-08_23-23-54/230k_ckpt.pt \
#     --train_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_train.json \
#     --val_input_path=/home/tony/Data/Hoot/v6_t1_cer_50_long_test.json \
#     --preload_decoder=True \
#     --custom_seed_offset=8888 \
#     --lyrics_key=text \
#     --wandb_run_name=hoot_v6_v2_n16_xlarge_continued

# run crashed at 230k, extend training
/home/tony/anaconda3/envs/suno_12/bin/python -u \
    train_audios.py \
    --out_dir=/home/tony/Data/checkpoints/hoot \
    --max_duration_s=280 \
    --max_iters=101 \
    --warmup_iters=10 \
    --log_interval=25 \
    --eval_interval=5_000 \
    --eval_iters=50 \
    --learning_rate=1e-6 \
    --min_lr=1e-6 \
    --batch_size=2 \
    --decoder_type=v2 \
    --gradient_accumulation_steps=1 \
    --n_augment_freq_masks=4 \
    --preload_checkpoint=/home/tony/Data/checkpoints/hoot/2025-04-23_02-24-11/last_ckpt.pt  \
    --preload_decoder=True \
    --train_suno_input_path=/home/tony/Data/Hoot/suno_hoot_20250617_subset100k.json \
    --custom_seed_offset=8889 \
    --lyrics_key=text \
    --wandb_run_name=hoot_v6_v2_sft_test

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