#!/bin/bash
#SBATCH -A nvr_elm_llm
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 04:00:00
#SBATCH -N 8
#SBATCH --ntasks-per-node=1
#SBATCH -J train_sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision
#SBATCH --gpus-per-node 8
#SBATCH --array=1-5%1
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision_error.log
#SBATCH -o workdir/train_sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision_out.log
#SBATCH --open-mode=append
#SBATCH --requeue
#SBATCH --signal=B:USR1@120

# ============================================
# SIGNAL HANDLER FOR SLURM TIMEOUT/PREEMPTION
# ============================================
TRAIN_PID=""

handle_timeout() {
    echo ""
    echo "=========================================="
    echo "Received USR1 signal (SLURM timeout warning)"
    echo "Gracefully stopping training for checkpoint save..."
    echo "=========================================="
    
    # Send SIGTERM to training process to trigger graceful shutdown
    if [ -n "$TRAIN_PID" ] && kill -0 $TRAIN_PID 2>/dev/null; then
        echo "Sending SIGTERM to training process (PID: $TRAIN_PID)..."
        kill -TERM $TRAIN_PID 2>/dev/null || true
        # Wait for process to finish saving checkpoint
        wait $TRAIN_PID 2>/dev/null || true
    fi
    
    echo "Requesting SLURM to requeue job..."
    scontrol requeue $SLURM_JOB_ID
    exit 0
}

# Set up signal handler
trap handle_timeout USR1

# Resolve project root and change to it immediately
PROJECT_ROOT="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow"
cd "${PROJECT_ROOT}"

# ============================================
# MULTI-NODE CONFIGURATION (8 NODES)
# ============================================
# Get master node address from SLURM
MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
MASTER_PORT=${MASTER_PORT:-29500}
NNODES=${SLURM_NNODES:-8}
NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}

export MASTER_ADDR
export MASTER_PORT

echo "=========================================="
echo "MDM_Qwen2.5-VL Training (ShareGPT4V)"
echo "=========================================="
echo "Multi-Node Configuration (8 Nodes):"
echo "  NNODES: ${NNODES}"
echo "  NGPUS_PER_NODE: ${NGPUS_PER_NODE}"
echo "  MASTER_ADDR: ${MASTER_ADDR}"
echo "  MASTER_PORT: ${MASTER_PORT}"
echo "  SLURM_JOB_NODELIST: ${SLURM_JOB_NODELIST}"
echo ""
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "SLURM Job ID: ${SLURM_JOB_ID:-N/A}"
echo "Working Directory: $(pwd)"
echo ""

# Activate environment
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

echo "Python: $(python3 --version)"
echo "PyTorch: $(python3 -c 'import torch; print(torch.__version__)')"
echo "CUDA available: $(python3 -c 'import torch; print(torch.cuda.is_available())')"
echo "nvidia-smi:"
nvidia-smi || true
echo ""

# VILA dataset registry (cs-oci-ord has the correct paths for this cluster)
# VILA_DATASETS — loaded from cluster_env.sh
# DATASET_BACKEND — loaded from cluster_env.sh
# Set environment variables
export PYTHONPATH="${PROJECT_ROOT}/src:${PROJECT_ROOT}/VILA:${PROJECT_ROOT}/examples:${PYTHONPATH:-}"
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-16}"
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:64,garbage_collection_threshold:0.9"
export TORCHDYNAMO_DISABLE=1
export TORCH_COMPILE=0
export MDM_VL_KEEP_VISION=0
export TOKENIZERS_PARALLELISM=false

# Qwen2.5-VL processor dynamic resolution controls
# Default to ~448x448 (200704 pixels) to stay aligned with 14x14 patching.
# export MIN_PIXELS="${MIN_PIXELS:-200704}"
# export MAX_PIXELS="${MAX_PIXELS:-200704}"

# Configurable parameters (override via env before submitting)
# NOTE: Must use MDM version of the model!
MODEL_PATH="${MODEL_PATH:-${PROJECT_ROOT}/src/models/Qwen2.5-VL-3B-Instruct-mdm}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"
MAX_LENGTH="${MAX_LENGTH:-2048}"

# ShareGPT4V dataset configuration
DATASET_PATH="${DATASET_PATH:-sharegpt4v_sft_text}"

EPOCHS="${EPOCHS:-1}"
BATCH="${BATCH:-1}"
# With 8 nodes (64 GPUs): 64 GPUs * 1 batch * 4 grad_acc = 256 effective batch
GRAD_ACC="${GRAD_ACC:-4}"
LR="${LR:-1e-6}"
WARMUP_RATIO="${WARMUP_RATIO:-0.03}"
SAVE_STEPS="${SAVE_STEPS:-50}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-10}"
WORKERS="${WORKERS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"

# MDM specific parameters
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-8}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-1}"
ALWAYS_MASK_IM_END="${ALWAYS_MASK_IM_END:-1}"

# Enable Activation Checkpointing (saves memory by recomputing activations)
ACTIVATION_CHECKPOINTING="${ACTIVATION_CHECKPOINTING:-1}"

TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"
# Use fixed RUN_ID for auto-resume across sbatch submissions
RUN_ID="${RUN_ID:-sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision}"
FLEXIBLE_BD_SIZE="${FLEXIBLE_BD_SIZE:-0}"
USE_BLOCK_CAUSAL_MASK="${USE_BLOCK_CAUSAL_MASK:-1}"

OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/log/${RUN_ID}}"

# Skip dataset adaptation by default (set to 0 to enable)
SKIP_ADAPT="${SKIP_ADAPT:-0}"

echo "Debug: Current PWD=$(pwd)"
echo "Debug: OUTPUT_DIR=${OUTPUT_DIR}"
echo "Debug: LOG_DIR=${LOG_DIR}"
mkdir -p "${OUTPUT_DIR}" "${LOG_DIR}"

# Sanity-check pixel settings stay compatible with the model patch size (14x14).
if [[ "${MIN_PIXELS}" =~ ^[0-9]+$ ]] && (( MIN_PIXELS % 196 != 0 )); then
    echo "Warning: MIN_PIXELS=${MIN_PIXELS} is not divisible by 14*14; expect padding overhead."
fi
if [[ "${MAX_PIXELS}" =~ ^[0-9]+$ ]] && (( MAX_PIXELS % 196 != 0 )); then
    echo "Warning: MAX_PIXELS=${MAX_PIXELS} is not divisible by 14*14; expect padding overhead."
fi

# ============================================
# CHECK DATASET PATHS
# ============================================
echo "Checking paths..."
if [ ! -d "${MODEL_PATH}" ]; then
    echo "❌ Model path does not exist: ${MODEL_PATH}"
    echo "⚠️  Make sure you're using the MDM version of the model!"
    exit 1
fi
echo "✅ Model path: ${MODEL_PATH}"

# ============================================
# AUTO-RESUME: Detect latest checkpoint
# ============================================
latest_checkpoint=""
if [ -d "${OUTPUT_DIR}" ]; then
    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

# Build resume_from_checkpoint argument
resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
    echo "📍 将从checkpoint恢复训练"
else
    resume_arg="--overwrite_output_dir"
    echo "🆕 从头开始训练"
fi

# Activation Checkpointing configuration
ac_arg=""
fsdp_wrap_arg=""
if [ "${ACTIVATION_CHECKPOINTING}" = "1" ]; then
    ac_arg="--gradient_checkpointing"
    fsdp_wrap_arg="" 
    echo "✅ Activation Checkpointing enabled (relying on model._no_split_modules for FSDP wrapping)"
fi

echo ""
echo "=========================================="
echo "Configuration (MDM Mode - 8 Nodes - ShareGPT4V - Bsz 256 - Bd 8 - Block Causal - LR 1e-6 - Unfreeze Vision):"
echo "  Model: ${MODEL_PATH}"
echo "  Tokenizer: ${TOKENIZER_NAME}"
echo "  Dataset Path: ${DATASET_PATH}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Log Dir: ${LOG_DIR}"
echo "  Epochs: ${EPOCHS}"
echo "  Batch: ${BATCH}"
echo "  Grad Acc: ${GRAD_ACC}"
echo "  LR: ${LR}"
echo "  Warmup Ratio: ${WARMUP_RATIO}"
echo "  MDM: ${MDM}  BD_SIZE: ${BD_SIZE}  LearnPadding: ${LEARN_PADDING}  ComplementaryMask: ${CONPLE_MASK}"
echo "  Always Mask IM End: ${ALWAYS_MASK_IM_END}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  FSDP: shard_grad_op (ZeRO-2)"
echo "  Num Nodes: ${NNODES}"
echo "  GPUs per Node: ${NGPUS_PER_NODE}"
echo "  Total GPUs: $((NNODES * NGPUS_PER_NODE))"
echo "  Effective Batch Size: $((NNODES * NGPUS_PER_NODE * BATCH * GRAD_ACC))"
echo "  Resume Arg: ${resume_arg}"
echo "  SLURM Requeue: enabled"
echo "=========================================="
echo ""

echo "=========================================="
echo "Starting MDM training with FSDP (8 nodes, 64 GPUs - ShareGPT4V - Bsz 256 - Bd 8 - Block Causal - LR 1e-6 - Unfreeze Vision)..."
echo "=========================================="

# Run training using srun to launch on all nodes
# Each node runs torchrun which spawns 8 GPU processes
srun --nodes=${NNODES} --ntasks-per-node=1 \
    --export=ALL,MASTER_ADDR=${MASTER_ADDR},MASTER_PORT=${MASTER_PORT} \
    bash -c '
        cd '"${PROJECT_ROOT}"'
        source '"${PROJECT_ROOT}"'/scripts/cluster_env.sh
        lmflow_activate_conda
        
        export PYTHONPATH="'"${PROJECT_ROOT}"'/src:'"${PROJECT_ROOT}"'/VILA:'"${PROJECT_ROOT}"'/examples"
        # VILA_DATASETS — loaded from cluster_env.sh
        # DATASET_BACKEND — loaded from cluster_env.sh
        export MIN_PIXELS="'"${MIN_PIXELS}"'"
        export MAX_PIXELS="'"${MAX_PIXELS}"'"
        export MDM_VL_KEEP_VISION=0
        export OMP_NUM_THREADS=16
        export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:64,garbage_collection_threshold:0.9"
        export TORCHDYNAMO_DISABLE=1
        export TORCH_COMPILE=0
        export TOKENIZERS_PARALLELISM=false
        
        torchrun \
            --nnodes='"${NNODES}"' \
            --nproc_per_node='"${NGPUS_PER_NODE}"' \
            --rdzv_id='"${SLURM_JOB_ID}"' \
            --rdzv_backend=c10d \
            --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
            examples/finetune_qwenvl.py \
            --model_name_or_path "'"${MODEL_PATH}"'" \
            --tokenizer_name "'"${TOKENIZER_NAME}"'" \
            --conversation_template "'"${TEMPLATE}"'" \
            --trust_remote_code 1 \
            --dataset_path "'"${DATASET_PATH}"'" \
            --output_dir "'"${OUTPUT_DIR}"'" \
            --num_train_epochs '"${EPOCHS}"' \
            --lr_scheduler_type constant_with_warmup \
            --warmup_ratio '"${WARMUP_RATIO}"' \
            --learning_rate '"${LR}"' \
            --per_device_train_batch_size '"${BATCH}"' \
            --gradient_accumulation_steps '"${GRAD_ACC}"' \
            --logging_steps '"${LOGGING_STEPS}"' \
            --save_steps '"${SAVE_STEPS}"' \
            --save_total_limit '"${SAVE_TOTAL_LIMIT}"' \
            --do_train \
            --dataloader_num_workers '"${WORKERS}"' \
            --ddp_timeout 7200 \
            --mdm '"${MDM}"' \
            --bd_size '"${BD_SIZE}"' \
            --block_size '"${BLOCK_SIZE}"' \
            --learn_padding '"${LEARN_PADDING}"' \
            --conplemenrary_mask '"${CONPLE_MASK}"' \
            --bf16 \
            --fsdp "shard_grad_op" \
            --max_length '"${MAX_LENGTH}"' \
            --always_mask_im_end '"${ALWAYS_MASK_IM_END}"' \
            --flexible_bd_size '"${FLEXIBLE_BD_SIZE}"' \
            --use_block_causal_mask '"${USE_BLOCK_CAUSAL_MASK}"' \
            '"${ac_arg}"' \
            '"${fsdp_wrap_arg}"' \
            '"${resume_arg}"'
    ' 2>&1 | tee "${LOG_DIR}/train_${SLURM_ARRAY_TASK_ID}-of-${SLURM_ARRAY_TASK_COUNT}.log" &
TRAIN_PID=$!

# Wait for training to complete (or be interrupted by signal)
wait $TRAIN_PID
exit_code=$?

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ ShareGPT4V Text MDM Training (8-node - Bsz 256 - Bd 8 - Block Causal - LR 1e-6 - Unfreeze Vision) 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
else
    echo "❌ ShareGPT4V Text MDM Training (8-node - Bsz 256 - Bd 8 - Block Causal - LR 1e-6 - Unfreeze Vision)
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log\n  Error: workdir/train_sharegpt4v_text_8node_bsz256_bd8_block_causal_lr_1e-6_unfreeze_vision_error.log"
fi
echo "=========================================="

exit $exit_code

