#! /bin/bash

MASTER_ADDR=localhost
MASTER_PORT=${2-2012}
NNODES=1
NODE_RANK=0
GPUS_PER_NODE=${3-16}

START_ALPHA_BETA=${4-0.9}  # Starting value of alpha + beta
END_ALPHA_BETA=${5-0.9}    # Ending value of alpha + beta
START_ALPHA=${6-0.1}       # Starting value of alpha
END_ALPHA=${7-0.1}         # Ending value of alpha

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE \
                  --nnodes $NNODES \
                  --node_rank $NODE_RANK \
                  --master_addr $MASTER_ADDR \
                  --master_port $MASTER_PORT"

# model
BASE_PATH=${1-"/home/MiniLLM"}
CKPT_NAME="Qwen/Qwen2.5-1.5B-Instruct"
#CKPT_NAME="results/qwen/train/seqkd/evolcode/1.5B_7B/1222"
CKPT="${BASE_PATH}/${CKPT_NAME}/"
TEACHER_CKPT_NAME="Qwen/Qwen2.5-Math-7B-Instruct"
TEACHER_CKPT="${BASE_PATH}/${TEACHER_CKPT_NAME}/"
# data
DATA_DIR="${BASE_PATH}/processed_data/metamath/pseudo/qwen/"
# hp
BATCH_SIZE=4
LR=0.00001
GRAD_ACC=4
EVAL_BATCH_SIZE=8
# length
MAX_LENGTH=2048
# seed
SEED=10


OPTS=""
# model
OPTS+=" --base-path ${BASE_PATH}"
OPTS+=" --model-path ${CKPT}"
OPTS+=" --teacher-model-path ${TEACHER_CKPT}"
OPTS+=" --ckpt-name ${CKPT_NAME}"
OPTS+=" --teacher-ckpt-name ${TEACHER_CKPT_NAME}"
OPTS+=" --teacher-model-fp16"
OPTS+=" --n-gpu ${GPUS_PER_NODE}"
OPTS+=" --model-type qwen"
OPTS+=" --gradient-checkpointing"
# data
OPTS+=" --data-dir ${DATA_DIR}"
# OPTS+=" --lm-data-dir ${LM_DATA_DIR}"
OPTS+=" --num-workers 4"
OPTS+=" --dev-num 1000"
# hp
OPTS+=" --lr ${LR}"
OPTS+=" --batch-size ${BATCH_SIZE}"
OPTS+=" --eval-batch-size ${EVAL_BATCH_SIZE}"
OPTS+=" --gradient-accumulation-steps ${GRAD_ACC}"
OPTS+=" --warmup-iters 0"
OPTS+=" --lr-decay-style cosine"
OPTS+=" --weight-decay 1e-2"
OPTS+=" --clip-grad 1.0"
OPTS+=" --epochs 4"
OPTS+=" --kd-ratio 1.0"
# length
OPTS+=" --max-length ${MAX_LENGTH}"
OPTS+=" --max-prompt-length 512"
# runtime
OPTS+=" --do-train"
OPTS+=" --do-valid"
OPTS+=" --eval-gen"
OPTS+=" --save-interval -1"
OPTS+=" --eval-interval -1"
OPTS+=" --log-interval 4"
OPTS+=" --mid-log-num -1"
# OPTS+=" --save ${SAVE_PATH}"
# lora
# OPTS+=" --peft lora"
OPTS+=" --do-train"
# OPTS+=" --peft-name ${PEFT_CKPT_NAME}"
# OPTS+=" --peft-path ${PEFT_CKPT}"
# OPTS+=" --teacher-peft-name ${TEACHER_PEFT_CKPT_NAME}"
# OPTS+=" --teacher-peft-path ${TEACHER_PEFT_CKPT}"
# seed
OPTS+=" --seed ${SEED}"
# deepspeed
OPTS+=" --deepspeed"
OPTS+=" --deepspeed_config ${BASE_PATH}/configs/deepspeed/ds_config_bf16.json"
# type
OPTS+=" --type ab"
# gen
OPTS+=" --do-sample"
OPTS+=" --top-k 0"
OPTS+=" --top-p 1.0"
OPTS+=" --temperature 1.0"
# GKD
OPTS+=" --student-gen"
OPTS+=" --init-threshold 0.2"
OPTS+=" --loss-eps 0.05"
OPTS+=" --mixed-alpha 0.35"

export NCCL_DEBUG=""
export WANDB_DISABLED=True
export TF_CPP_MIN_LOG_LEVEL=3
export PYTHONPATH=${BASE_PATH}

# for alpha_beta in 1.0; do
#    for alpha in 1.0; do
#        beta=$(echo "$alpha_beta - $alpha" | bc)
#         # 输出当前alpha和beta值
#        echo "alpha_beta=$alpha_beta, alpha=$alpha, beta=$beta"
#        # runtime
#        SAVE_PATH="${BASE_PATH}/results/qwen/train/evolcode/jsd_math/distill_1.5B_7B_final"
#       # SAVE_PATH="${BASE_PATH}/results/qwen/train/metamath/ab/distill_1.5B_7B_final/${alpha}_${beta}"
#        mkdir -p ${SAVE_PATH}

#        CURRENT_OPTS="${OPTS}"
#        CURRENT_OPTS+=" --save ${SAVE_PATH}"
#        CURRENT_OPTS+=" --ab_alpha ${alpha}"
#        CURRENT_OPTS+=" --ab_beta ${beta}"

#        CMD="torchrun ${DISTRIBUTED_ARGS} ${BASE_PATH}/finetune.py ${CURRENT_OPTS} $@"
#        echo ${CMD}
#        echo "PYTHONPATH=${PYTHONPATH}"
#        ${CMD}
#    done
# done


for alpha_beta in $(seq ${START_ALPHA_BETA} 0.1 ${END_ALPHA_BETA}); do
   for alpha in $(seq ${START_ALPHA} 0.1 ${END_ALPHA}); do
       beta=$(echo "$alpha_beta - $alpha" | bc)
       
       # runtime
    #    SAVE_PATH="${BASE_PATH}/results/qwen/train/evolcode_pseudo/ab/distill_1.5B_7B_final/${alpha}_${beta}"
       SAVE_PATH="${BASE_PATH}/results/qwen/train/metamath/ab/distill_1.5B_7B_final/${alpha}_${beta}"
       mkdir -p ${SAVE_PATH}

       CURRENT_OPTS="${OPTS}"
       CURRENT_OPTS+=" --save ${SAVE_PATH}"
       CURRENT_OPTS+=" --ab_alpha ${alpha}"
       CURRENT_OPTS+=" --ab_beta ${beta}"

       CMD="torchrun ${DISTRIBUTED_ARGS} ${BASE_PATH}/finetune.py ${CURRENT_OPTS} $@"
       echo ${CMD}
       echo "PYTHONPATH=${PYTHONPATH}"
       ${CMD}
   done
done