#!/bin/bash
export CUDA_LAUNCH_BLOCKING=0
export NCCL_DEBUG=WARN
export TORCH_DISTRIBUTED_DEBUG=INFO

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`

TRAIN_PATH=/home/minz/glockenspiel/sunoGPT
echo working from $TRAIN_PATH
cd $TRAIN_PATH

torchrun \
    /home/minz/glockenspiel/sunoGPT/train_textualize.py \
    \
    --out_dir="/app/suno/checkpoints" \
    --data_dir="/app/suno/data/chirp_v3_ft_v3_full" \
    \
    --learning_rate=5e-4 \
    --max_iters=100_000 \
    --weights_multiplier="mixed" \
    \
    --block_size=6016 \
    --t_text=3008 \
    --t_memmap=3008 \
    --t_audio=3008 \
    \
    --n_layer=6 \
    --n_head=12 \
    --d_head=64 \
    \
    --gradient_accumulation_steps=1 \
    --batch_size=16 \
    \
    --fsdp=False \
    --sharding_strategy="full_shard" \
    --debug_val_only=True \
    --eval_iters=50 \
    \
    --wandb_log=False \
    --wandb_project="gpt-textualize" \
    --wandb_run_name="tagging" \