#!/bin/bash
#SBATCH --job-name="fintune"
#SBATCH  --account=bckr-dtai-gh
#SBATCH --partition=ghx4
#SBATCH --nodes=1
#SBATCH --gpus-per-node=4
#SBATCH --tasks=1
#SBATCH --tasks-per-node=1
#SBATCH --cpus-per-task=20
#SBATCH --mem=512g
#SBATCH --time=23:59:00
#SBATCH --output="run.log"
#SBATCH --error="run.err"

set -e
export WANDB_API_KEY='1b2611814911cad498235f1ccb1a2e182638bd62'

# set up exp1 or exp3!!!!!
# launch this script after bilevel weighting and preparing data
# this script is for exp1 and exp3

# 1. finetune on bilevel and baseline

CUDA_VISIBLE=0,1

hf_ds=pxyyy/NuminaMath-CoT-smp10k
hf_val_ds=pxyyy/NuminaMath-CoT-smp10k
model_and_tok=Qwen/Qwen2.5-7B
# conv_template=llama3
conv_template=qwen2_5_math

hf_ds_str=$(echo ${hf_ds}|sed 's/\//-/g')
tmp_data_dir=./tmp_data/${hf_ds_str}/
val_data_dir=./tmp_data/${hf_ds_str}_val/
mkdir -p ${tmp_data_dir}
mkdir -p ${val_data_dir}
python3 hf2lmflow.py --ds_name ${hf_ds} --save ${tmp_data_dir}/data.json --split train
python3 hf2lmflow.py --ds_name ${hf_val_ds} --save ${val_data_dir}/data.json --split test

model_str=$(echo ${model_and_tok}|sed 's/\//-/g')

gradient_accumulation_steps=4
per_device_train_batch_size=8
epoch=1
project_dir=/u/xpan2/projects/mp-llm/MATH/finetune
for lr in 2e-5
do
    # Finetune
    exp_id=finetune-${model_str}-${hf_ds_str}-${epoch}-$lr
    # project_dir=$(cd "$(dirname $0)"; pwd)
    log_dir=${project_dir}/log/${exp_id}
    output_dir=${project_dir}/output_models/${exp_id}
    
    echo $exp_id
    
    mkdir -p ${output_dir} ${log_dir}

    export TRANSFORMERS_VERBOSITY=info

    deepspeed --master_port=7964 --include=localhost:${CUDA_VISIBLE} finetune.py \
        --model_name_or_path ${model_and_tok} \
        --trust_remote_code 1 \
        --dataset_path ${tmp_data_dir}/ \
        --eval_dataset_path ${val_data_dir}/ \
        --output_dir ${output_dir} --overwrite_output_dir \
        --conversation_template ${conv_template} \
        --num_train_epochs $epoch \
        --learning_rate $lr \
        --disable_group_texts 1 \
        --block_size 512 \
        --per_device_train_batch_size ${per_device_train_batch_size} \
        --per_device_eval_batch_size 1 \
        --bf16 \
        --deepspeed configs/ds_config_zero3_no_offload.json \
        --torch_dtype bfloat16 \
        --run_name ${exp_id} \
        --optim adamw_torch_fused \
        --logging_steps 1 \
        --do_train \
        --do_eval \
        --ddp_timeout 72000 \
        --save_total_limit 1 \
        --load_best_model_at_end False \
        --eval_steps 10 \
        --save_only_model \
        --evaluation_strategy "steps" \
        --dataloader_num_workers 1 \
        --lr_scheduler_type cosine \
        --warmup_ratio 0.03 \
        --gradient_checkpointing True \
        --use_flash_attention 1 \
        --gradient_accumulation_steps ${gradient_accumulation_steps} \
        | tee ${log_dir}/train.log \
        2> ${log_dir}/train.err
done