# dpo on one node 
# dpo on 16 hrs of data, takes ~ 400 steps to overfit
OMP_NUM_THREADS=8 torchrun \
    --nnodes=4 --nproc_per_node=8 \
    --rdzv_backend=c10d --rdzv_id=123 \
    --rdzv_endpoint="compute-hpc-node-973:1234" \
    dpo.py \
    \
    --out_dir="/app/suno/checkpoints" \
    --data_dir="/app/suno/data/dpo/v0" \
    \
    --train_filename="data_tr.bin" \
    --train_metas_filename="meta_tr.jsonl" \
    --train_info_filename="info_tr.json" \
    --val_filename="data_val.bin" \
    --val_metas_filename="meta_val.jsonl" \
    --val_info_filename="info_val.json" \
    \
    --learning_rate=1e-6 \
    --min_lr=1e-7 \
    --grad_clip=0.1 \
    --warmup_iters=200 \
    --max_iters=1_500 \
    --eval_interval=200 \
    --eval_iters=25 \
    --step_save_iters=500 \
    --step_save_infer=False \
    \
    --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=2 \
    \
    --fsdp=True \
    --sharding_strategy="full_shard" \
    --grad_checkpointing=True \
    \
    --preload_checkpoint="/app/suno/checkpoints/2024-01-20_22-01-25/last_ckpt.pt" \
    --preload_strict=False \
    \
    --wandb_log=True \
    --wandb_project="chirp-v3-dpo" \
    --wandb_run_name="v3_dpo_t0_exp_8_1k_beta2"

# calculate loss
# evaluate loss on A100
# 4 --> 28 GB; 8 --> 33.4 GB; 10 --> 35.7 GB;
# --val_filename="data_tr.bin" \
# --val_metas_filename="meta_tr.jsonl" \
# --val_info_filename="info_tr.json" \
# --val_filename="data_val.bin" \
# --val_metas_filename="meta_val.jsonl" \
# --val_info_filename="info_val.json" \
OMP_NUM_THREADS=4 torchrun \
    --nnodes=1 --nproc_per_node=8 \
    --rdzv_backend=c10d --rdzv_id=123 --rdzv_endpoint="compute-hpc-node-775:1234" \
    eval_loss.py \
    --out_dir="/app/suno/data/dpo/v0" \
    --data_dir="/app/suno/data/dpo/v0" \
    --preload_checkpoint="/app/suno/checkpoints/2024-01-20_22-01-25/last_ckpt_infer.pt" \
    --val_filename="data_tr.bin" \
    --val_metas_filename="meta_tr.jsonl" \
    --val_info_filename="info_tr.json" \
    \
    --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 \
    --batch_size=10