#!/bin/bash
#SBATCH -A llmservice_fm_vision                                                                #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 finetune_qwen_0.5B_shizhe_stage1       #job name
#SBATCH --array=1-5%1
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/error.log
#SBATCH -o workdir/out.log
#SBATCH --exclusive
# SBATCH --exclude

idx=$SLURM_ARRAY_TASK_ID
total=$SLURM_ARRAY_TASK_COUNT

nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
slurm_name=${idx}-of-${total}

export LOGLEVEL=INFO

# 训练参数设置
model_name_or_path=models/Qwen2.5-0.5B-Instruct
dataset_path=split_by_instances
output_dir=output_models/finetune_qwen_0.5B_shizhe_stage1_constant_lr_warmup_test
deepspeed_args="--master_port=11000"
conversation_template=qwen2_5
export CUDA_HOME=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/lmflow
trust_remote_code=1

job_name=finetune_qwen_0.5B_shizhe_stage1

# 添加节点信息打印
echo "Nodes allocated: $SLURM_JOB_NODELIST"
echo "Number of nodes: $SLURM_NNODES"
echo "head_node: $head_node"
echo "head_node_ip: $head_node_ip"
echo "slurm_name: $slurm_name"

# 动态检测最新的checkpoint
latest_checkpoint=""
if [ -d "${output_dir}" ]; then
    # 查找所有checkpoint目录，按数字排序，取最新的一个
    latest_checkpoint=$(find "${output_dir}" -name "checkpoint-*" -type d | sort -V | tail -1)
    if [ -n "${latest_checkpoint}" ]; then
        echo "找到最新的checkpoint: ${latest_checkpoint}"
    else
        echo "在 ${output_dir} 中没有找到checkpoint"
        latest_checkpoint=""
    fi
else
    echo "输出目录 ${output_dir} 不存在，将从头开始训练"
    latest_checkpoint=""
fi

# 创建必要的目录
exp_id=finetune
project_dir=$(cd "$(dirname $0)"/..; pwd)
log_dir=${project_dir}/log/${exp_id}
mkdir -p ${output_dir} ${log_dir}

# 构建resume_from_checkpoint参数
resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
fi

read -r -d '' cmd <<EOF
deepspeed ${deepspeed_args} \
  examples/finetune.py \
    --model_name_or_path ${model_name_or_path} \
    --trust_remote_code ${trust_remote_code} \
    --dataset_path ${dataset_path} \
    --output_dir ${output_dir} \
    ${resume_arg} \
    --conversation_template ${conversation_template} \
    --num_train_epochs 1 \
    --learning_rate 2e-5 \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio 0.03 \
    --disable_group_texts 0 \
    --block_size 2048 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 2 \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    --bf16 \
    --run_name finetune \
    --validation_split_percentage 0 \
    --logging_steps 20 \
    --do_train \
    --ddp_timeout 72000 \
    --save_steps 1000 \
    --dataloader_num_workers 8 \
    --preprocessing_num_workers 8
EOF

echo $cmd
srun -o ${output_dir}/${job_name}/slurm_log/train/slurm_%A_${slurm_name}.log -e ${output_dir}/${job_name}/slurm_log/train/error_%A_${slurm_name}.log \
    bash -c "${cmd}" 