#!/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_llada_tulu-3       #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

output_dir=output
job_name=finetune_llada_tulu-3_block-diffusion

# 添加节点信息打印
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"

read -r -d '' cmd <<EOF
fabric run --node-rank=\$SLURM_PROCID --accelerator=cuda --devices=8 --num-nodes=$SLURM_NNODES \
    --main-address=$head_node_ip --main-port=26495 sft/finetune_llada.py --bs 256 --epoch 3 --nodes_num 8 --devices_num 8 --model GSAI-ML/LLaDA-8B-Base --dataset_name tulu-3 --load_local --save_step_interval 500 --micro_batch_size 4 --learning_rate 2.5e-5 --hp_name GSAI-ML-LLaDA-8B-Base-tulu-3-block-diffusion-pad-more-eos --use_block_diffusion  --max_length 2048 --pad_to_max_length
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}" 