#!/bin/bash
#SBATCH -A nvr_elm_llm                                                                #account
#SBATCH -p batch_block1                                   #partition
#SBATCH -t 04:00:00                                                                 #wall time limit, hr:min:sec
#SBATCH -N 8                                                                        #number of nodes
#SBATCH -J train_sharegpt4v_vlm_ar_baseline_8node_bsz256_lr_1e-6                              #job name
#SBATCH --array=1-100%1
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/error_ar_baseline_8node_bsz256_lr_1e-6.log
#SBATCH -o workdir/out_ar_baseline_8node_bsz256_lr_1e-6.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

# Project root
PROJECT_ROOT="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow"

# 训练参数设置 (from ar_baseline)
MODEL_PATH="${MODEL_PATH:-/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2.5-VL-3B-Instruct}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"

DATASET_JSON="${DATASET_JSON:-/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/jason-filter-sharegpt4v_mix665k_cap23k_coco-ap9k_lcs3k_sam9k_div2k.json}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/data}"

EPOCHS="${EPOCHS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-4}"
LR="${LR:-1e-6}"
WARMUP_RATIO="${WARMUP_RATIO:-0.03}"
BF16="${BF16:-1}"
GRADIENT_CHECKPOINTING="${GRADIENT_CHECKPOINTING:-1}"
DISABLE_GROUP_TEXTS="${DISABLE_GROUP_TEXTS:-1}"
SAVE_STEPS="${SAVE_STEPS:-100}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-1}"
WORKERS="${WORKERS:-8}"
PREP_WORKERS="${PREP_WORKERS:-32}"
MDM="${MDM:-0}"
BD_SIZE="${BD_SIZE:-8}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-1}"
TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"

RUN_ID="${RUN_ID:-sharegpt4v_ar_baseline_8node_bsz256_lr_1e-6}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
DS_CONFIG="${DS_CONFIG:-configs/ds_config_zero2_no_offload.json}"

# Qwen2.5-VL processor dynamic resolution controls
# export MIN_PIXELS="${MIN_PIXELS:-200704}"
# export MAX_PIXELS="${MAX_PIXELS:-200704}"
# export RESIZED_HEIGHT="${RESIZED_HEIGHT:-}"
# export RESIZED_WIDTH="${RESIZED_WIDTH:-}"

# Environment variables
# export CUDA_HOME=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/lmflow

trust_remote_code=1
deepspeed_args="--master_port=23456 --num_nodes=$SLURM_NNODES --node_rank=$SLURM_PROCID"

job_name=train_sharegpt4v_vlm_ar_baseline_8node_bsz256_lr_1e-6

# 添加节点信息打印
echo "=========================================="
echo "MDM_Qwen2VL VLM Training (ShareGPT4V, bsz256, lr1e-6, 8node)"
echo "=========================================="
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"
echo "Date: $(date)"
echo ""

# 动态检测最新的checkpoint (AutoResume)
latest_checkpoint=""
if [ -d "${OUTPUT_DIR}" ]; then
    # 查找所有checkpoint目录，按数字排序，取最新的一个
    latest_checkpoint=$(find "${OUTPUT_DIR}" -maxdepth 1 -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=train_sharegpt4v_vlm_ar_baseline_8node_bsz256_lr_1e-6
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}"
    echo "📍 将从checkpoint恢复训练"
else
    resume_arg="--overwrite_output_dir"
    echo "🆕 从头开始训练"
fi

echo ""
echo "=========================================="
echo "Configuration:"
echo "  Model: ${MODEL_PATH}"
echo "  Tokenizer: ${TOKENIZER_NAME}"
echo "  Dataset JSON: ${DATASET_JSON}"
echo "  Image Folder: ${IMAGE_FOLDER}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Log Dir: ${log_dir}"
echo "  Epochs: ${EPOCHS}"
echo "  Block Size: ${BLOCK_SIZE}"
echo "  Batch: ${BATCH}"
echo "  Grad Acc: ${GRAD_ACC}"
echo "  LR: ${LR}"
echo "  Warmup Ratio: ${WARMUP_RATIO}"
echo "  BF16: ${BF16}"
echo "  MDM: ${MDM}  BD_SIZE: ${BD_SIZE}  LearnPadding: ${LEARN_PADDING}  ComplementaryMask: ${CONPLE_MASK}"
echo "  DS Config: ${DS_CONFIG}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  Template: ${TEMPLATE}"
echo "=========================================="
echo ""

# GPU monitoring
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='
# 获取当前机器的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 <<< $('"${PROJECT_ROOT}"'/expand_hostfile.sh '"${log_dir}"');
fi

wait 60
'"${PROJECT_ROOT}"'/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;
  
  # 激活conda环境
  source ~/.bashrc;
  conda activate shiyi-lmflow || true;
  
  cd '"${PROJECT_ROOT}"';
  # Build deepspeed command
  deepspeed --master_addr=${master_node} --hostfile '"${log_dir}"'/deepspeed-hostfile '"${deepspeed_args}"' \
    examples/finetune_qwenvl.py \
    --model_name_or_path '"${MODEL_PATH}"' \
    --tokenizer_name '"${TOKENIZER_NAME}"' \
    --conversation_template '"${TEMPLATE}"' \
    --trust_remote_code ${trust_remote_code} \
    --dataset_path '"${DATASET_JSON}"' \
    --image_folder '"${IMAGE_FOLDER}"' \
    --output_dir '"${OUTPUT_DIR}"' \
    '"${resume_arg}"' \
    --num_train_epochs '"${EPOCHS}"' \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio '"${WARMUP_RATIO}"' \
    --learning_rate '"${LR}"' \
    --disable_group_texts '"${DISABLE_GROUP_TEXTS}"' \
    --block_size '"${BLOCK_SIZE}"' \
    --per_device_train_batch_size '"${BATCH}"' \
    --gradient_accumulation_steps '"${GRAD_ACC}"' \
    --run_name '"${exp_id}"' \
    --logging_steps '"${LOGGING_STEPS}"' \
    --do_train \
    --validation_split_percentage 0 \
    --ddp_timeout 72000 \
    --save_steps '"${SAVE_STEPS}"' \
    --save_total_limit '"${SAVE_TOTAL_LIMIT}"' \
    --dataloader_num_workers '"${WORKERS}"' \
    --preprocessing_num_workers '"${PREP_WORKERS}"' \
    --bd_size '"${BD_SIZE}"' \
    --mdm '"${MDM}"' \
    --learn_padding '"${LEARN_PADDING}"' \
    --conplemenrary_mask '"${CONPLE_MASK}"' \
    --deepspeed '"${DS_CONFIG}"' \
    '"$([ "${BF16}" = "1" ] && echo "--bf16")"' \
    '"$([ "${GRADIENT_CHECKPOINTING}" = "1" ] && echo "--gradient_checkpointing")"'
    
  exit_code=$?;
  echo "" 1>&2;
  echo "==========================================" 1>&2;
  if [ $exit_code -eq 0 ]; then
    echo "✅ ShareGPT4V VLM Training Success!"
    # ============================================
    # AUTO EVAL: 训练成功后自动提交评测任务
    # ============================================
    echo ""
    echo "=========================================="
    echo "🚀 Auto-submitting evaluation jobs..."
    echo "=========================================="
    
    if [ -f "${PROJECT_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" ]; then
        bash "${PROJECT_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" \
            "${OUTPUT_DIR}" \
            --skip-existing 1 \
            2>&1 | tee "${LOG_DIR}/auto_eval_submit.log"
        echo "✅ Evaluation jobs submitted. Check: ${OUTPUT_DIR}/eval_vlmeval/"
    else
        echo "⚠️  submit_batch_eval.sh not found, skipping auto eval"
    fi
 1>&2;
  else
    echo "❌ ShareGPT4V VLM Training Failed (exit code: ${exit_code})" 1>&2;
  fi
  echo "==========================================" 1>&2;
  
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}"

