#!/bin/bash
#SBATCH -A nvr_elm_llm                                                                #account
#SBATCH -p batch_block1,batch_block3,batch_block4                                   #partition
#SBATCH -t 04:00:00                                                                   #wall time limit, hr:min:sec
#SBATCH -N 1                                                                           #number of nodes
#SBATCH -J throughput_test_qwen_7B                                                    #job name
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/throughput_error.log
#SBATCH -o workdir/throughput_out.log

# Set the environment variables first before running the command.
export HF_ALLOW_CODE_EVAL=1
export HF_DATASETS_TRUST_REMOTE_CODE=true

# 模型路径 - 7B模型（非AR）
model_path=output_models/finetune_qwen2.5-7B-bd-size-32-lr-1e-5/checkpoint-2500

# 任务配置 - 使用gsm8k进行throughput测试
task=gsm8k

# Throughput测试结果保存路径
throughput_log_path=eval_results/qwen2_5_7B_throughput_test.log
throughput_results_dir=eval_results/throughput_results_7B

# 创建必要的目录
mkdir -p workdir
mkdir -p workdir/slurm_log/throughput_test_7B_small_size_8
mkdir -p eval_results
mkdir -p ${throughput_results_dir}

echo "开始执行7B模型Throughput测试..."
echo "任务: ${task}"
echo "模型路径: ${model_path}"
echo "结果保存路径: ${throughput_results_dir}"

# 定义要测试的batch sizes
batch_sizes=(1 2 4 8 16 32 64)

# 计数器
experiment_count=0
total_experiments=${#batch_sizes[@]}

echo "总共需要运行 ${total_experiments} 个throughput测试"
echo "开始时间: $(date)" | tee -a ${throughput_log_path}

# 遍历所有batch size进行throughput测试
for batch_size in "${batch_sizes[@]}"; do
    experiment_count=$((experiment_count + 1))
    

    # 构建throughput测试命令
    eval_cmd="
    export HF_ALLOW_CODE_EVAL=1;
    export HF_DATASETS_TRUST_REMOTE_CODE=true;
    accelerate launch chengyue_script/eval_mdm_qwen_batch.py \
        --tasks ${task} \
        --num_fewshot 0 \
        --batch_size ${batch_size} \
        --confirm_run_unsafe_code \
        --model dfm_qwen \
        --fewshot_as_multiturn \
        --apply_chat_template \
        --model_args model_path=${model_path},show_speed=True,use_block_cache=True,small_block_size=8,dynamic_threshold=False,factor=0.9,threshold=0.9 
    "

    echo "执行throughput测试命令..."
    echo "命令: $eval_cmd"

    # 使用srun执行throughput测试命令
    srun -o workdir/slurm_log/throughput_test_7B_small_size_8/slurm_%A_%j_${task}_batch_${batch_size}.log \
         -e workdir/slurm_log/throughput_test_7B_small_size_8/error_%A_%j_${task}_batch_${batch_size}.log \
        bash -c "${eval_cmd}" 

done
