#!/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 8                                                                           #number of nodes
#SBATCH -J finetune_qwen_1.5B_nemotron_shizhe_pad_bd_size_128_8nodes_seq_length_8k       #job name
#SBATCH --array=1-100%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 NCCL_BLOCKING_WAIT=1
export NCCL_IB_DISABLE=1
export NCCL_SOCKET_IFNAME=eth0
export PDSH_RCMD_TYPE=ssh


export LOGLEVEL=INFO

# 训练参数设置
model_name_or_path=models/Qwen2.5-1.5B-Instruct
dataset_path=data/shizhe_stage1
output_dir=output_models/finetune_qwen_1.5B_nemotron_shizhe_pad_bd_size_128_8nodes_seq_length_8k
deepspeed_args="--master_port=23456 --num_nodes=$SLURM_NNODES --node_rank=$SLURM_PROCID"
conversation_template=qwen2_5
export CUDA_HOME=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/lmflow
trust_remote_code=1

job_name=finetune_qwen_1.5B_nemotron_shizhe_pad_bd_size_128_8nodes_seq_length_8k

# 添加节点信息打印
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=log/${job_name}
mkdir -p ${output_dir} ${log_dir}

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

nvidia-smi --query-gpu=timestamp,memory.used,memory.total,utilization.gpu,pstate \
           --format=csv,noheader,nounits -l 60 &

rm -rf ${log_dir}/deepspeed-hostfile
rm -rf ${log_dir}/temp-hostfile

cmd='
export PDSH_RCMD_TYPE=ssh;
export CUDA_HOME=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/lmflow;
# 获取当前机器的IP地址
echo "SLURMD_NODENAME:= " $SLURMD_NODENAME 1>&2;
echo "head_node:= " '"${head_node}"' 1>&2;
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo $SLURM_JOB_NODELIST > '"${log_dir}"'/temp-hostfile;
  read master_node slave_node <<< $(./expand_hostfile.sh '"${log_dir}"');
fi
./setup_ssh.sh '"${log_dir}"';
echo "MASTER_NODE:= " $master_node 1>&2;
# 检查当前机器是否为master节点
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo "Running on master node. Executing DeepSpeed..." "$current_ip" 1>&2;
  deepspeed --master_addr=${master_node} --hostfile '"${log_dir}"'/deepspeed-hostfile '"${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 8192 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 1 \
    --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 100 \
    --dataloader_num_workers 8 \
    --preprocessing_num_workers 32 \
    --learn_padding 1 \
    --bd_size 128 \
    --mdm 1 \
    --gradient_checkpointing 1 \
    --conplemenrary_mask 1 \
    --tokenizer_name models/Qwen2.5-0.5B-Instruct
else
    echo "Not running on master node. Skipping DeepSpeed execution." "$current_ip" 1>&2;
fi
'

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