#!/bin/bash
export CUDA_LAUNCH_BLOCKING=0
export NCCL_DEBUG=VERSION
export TORCH_DISTRIBUTED_DEBUG=OFF

export NCCL_IB_SL=0
export NCCL_IB_TC=41
export NCCL_IB_QPS_PER_CONNECTION=4
export NCCL_IB_GID_INDEX=3
export NCCL_IB_HCA="=mlx5_5,mlx5_6,mlx5_7,mlx5_8,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_14,mlx5_15,mlx5_16,mlx5_17,mlx5_9,mlx5_10,mlx5_11,mlx5_12"

export OMP_NUM_THREADS=4
export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12802
export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

TRAIN_PATH=/home/tony/Work/glockenspiel/sunoGPT
echo working from $TRAIN_PATH
cd $TRAIN_PATH
# 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
/home/tony/anaconda3/envs/suno_env/bin/torchrun \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    --nnodes=$SLURM_JOB_NUM_NODES \
    --nproc_per_node=$SLURM_GPUS_ON_NODE \
    --rdzv_id $SLURM_JOB_ID \
    --rdzv_backend c10d \
    --rdzv_endpoint "$MASTER_ADDR:$MASTER_PORT" \
    train.py \
    \
    --out_dir="/app/suno/checkpoints" \
    --data_dir="/app/suno/data/chirp_v3_ft_v3" \
\
    --learning_rate=6e-5 \
    --max_iters=50_000 \
    \
    --block_size=4288 \
    --t_text=1152 \
    --t_memmap=3008 \
    --t_audio=3136 \
    --coarse_vocab_size=2112 \
    --coarse_codebook_size=2048 \
    --coarse_pad_token=2048 \
    --coarse_infer_token=2049 \
    \
    --n_layer=32 \
    --n_head=32 \
    --d_head=128 \
    --n_kv_head=4 \
    \
    --gradient_accumulation_steps=1 \
    --batch_size=4 \
    \
    --fsdp=True \
    --sharding_strategy="full_shard" \
    --grad_checkpointing=True \
    \
    --wandb_log=True \
    --wandb_project="chirp-v3-ft-tony" \
    --wandb_run_name="7b_ft" \
    \
    --preload_checkpoint="/app/suno/checkpoints/2024-01-20_22-01-25/last_ckpt.pt" \
    --preload_strict=False \
    --custom_seed_offset=2