#!/bin/bash
#SBATCH -A nvr_elm_llm
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 4:00:00
#SBATCH -N 1
#SBATCH -J train_sharegpt4v_vlm_use_block_causal
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_sharegpt4v_vlm_mdm_fix_eos_use_block_causal_error.log
#SBATCH -o workdir/train_sharegpt4v_vlm_mdm_fix_eos_use_block_causal_out.log
#SBATCH --requeue
#SBATCH --signal=B:USR1@120

set -euxo pipefail

# Signal handler for requeue
on_signal() {
    echo "Caught signal; requeue current job $SLURM_JOB_ID"
    REQUEUED=1
    scontrol requeue $SLURM_JOB_ID || true
}
trap on_signal SIGTERM SIGUSR1

# 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}"

echo "=========================================="
echo "MDM_Qwen2VL VLM Training (ShareGPT4V, single-node)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "Working Directory: $(pwd)"
echo ""

# Activate environment
# source ~/.bashrc
# conda activate shiyi-lmflow || true

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
DATASET_PATH="${DATASET_PATH:-sharegpt4v_recipe}"

# Set environment variables (merge from interactive_gpu.slurm and run_deepspeed_sharegpt4v.sh)
# export PYTHONPATH="${PROJECT_ROOT}:${PROJECT_ROOT}/chengyue_script:${PROJECT_ROOT}/src:${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=1
export TOKENIZERS_PARALLELISM=false

# Qwen2.5-VL processor dynamic resolution controls (read by examples/finetune_qwenvl.py)
# Reduced to ~448x448 fixed to limit vision tokens (avoid OOM and extreme prefill lengths)
export MIN_PIXELS="${MIN_PIXELS:-200704}"
export MAX_PIXELS="${MAX_PIXELS:-200704}"
# Optional fixed resize override (must be multiples of 28 if set)
export RESIZED_HEIGHT="${RESIZED_HEIGHT:-}"
export RESIZED_WIDTH="${RESIZED_WIDTH:-}"

# Configurable parameters (override via env before submitting)
MODEL_PATH="${MODEL_PATH:-output_models/finetune_qwen2.5-VL-3B-bd-size-32-2k-no-dynamic-block-causal/checkpoint-2000}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"

EPOCHS="${EPOCHS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-1}"
LR="${LR:-2e-5}"
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:-50}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-1}"
WORKERS="${WORKERS:-8}"
PREP_WORKERS="${PREP_WORKERS:-32}"
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-32}"
LEARN_PADDING="${LEARN_PADDING:-1}"
COMPLEMENTARY_MASK="${COMPLEMENTARY_MASK:-1}"
ENTROPY_LOSS="${ENTROPY_LOSS:-1}"
ENTROPY_LOSS_WEIGHT="${ENTROPY_LOSS_WEIGHT:-1.0}"
TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"
RUN_ID="${RUN_ID:-sharegpt4v_vlm_mdm_fix_eos_use_block_causal}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/log/${RUN_ID}}"
DS_CONFIG="${DS_CONFIG:-configs/ds_config_zero2_no_offload.json}"
NGPUS="${NGPUS:-8}"
SKIP_ADAPT="${SKIP_ADAPT:-0}"
FLEXIBLE_BD_SIZE="${FLEXIBLE_BD_SIZE:-0}"
USE_BLOCK_CAUSAL_MASK="${USE_BLOCK_CAUSAL_MASK:-1}"
BLOCK_CAUSAL_NO_DYNAMIC="${BLOCK_CAUSAL_NO_DYNAMIC:-1}"

# AutoResume: check for existing checkpoints
AUTO_RESUME="${AUTO_RESUME:-1}"
RESUME_FROM_CHECKPOINT="${RESUME_FROM_CHECKPOINT:-}"

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

echo "=========================================="
echo "Configuration:"
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 "  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: ${COMPLEMENTARY_MASK}"
echo "  DS Config: ${DS_CONFIG}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  Processor Fixed WxH (opt): ${RESIZED_WIDTH:-auto}x${RESIZED_HEIGHT:-auto}"
echo "  Num GPUs: ${NGPUS}"
echo "  Auto Resume: ${AUTO_RESUME}"
echo "  Flexible BD Size: ${FLEXIBLE_BD_SIZE}"
echo "  Use Block Causal Mask: ${USE_BLOCK_CAUSAL_MASK}"
echo "  Entropy Loss: ${ENTROPY_LOSS} Entropy Loss Weight: ${ENTROPY_LOSS_WEIGHT}"
echo "  Block Causal No Dynamic: ${BLOCK_CAUSAL_NO_DYNAMIC}"
echo "=========================================="
echo ""

# Check paths
echo "Checking paths..."
if [ ! -d "${MODEL_PATH}" ]; then
    echo "❌ Model path does not exist: ${MODEL_PATH}"
    exit 1
fi
echo "✅ Model path: ${MODEL_PATH}"

# AutoResume: detect latest checkpoint
if [ "${AUTO_RESUME}" = "1" ] && [ -z "${RESUME_FROM_CHECKPOINT}" ]; then
    if [ -d "${OUTPUT_DIR}" ]; then
        # Find the latest checkpoint directory
        LATEST_CHECKPOINT=$(find "${OUTPUT_DIR}" -maxdepth 1 -type d -name "checkpoint-*" | sort -V | tail -n 1)
        if [ -n "${LATEST_CHECKPOINT}" ]; then
            RESUME_FROM_CHECKPOINT="${LATEST_CHECKPOINT}"
            echo "🔄 AutoResume: Found existing checkpoint: ${RESUME_FROM_CHECKPOINT}"
        else
            echo "ℹ️ AutoResume: No checkpoint found, starting from scratch"
        fi
    else
        echo "ℹ️ AutoResume: Output directory does not exist, starting from scratch"
    fi
elif [ -n "${RESUME_FROM_CHECKPOINT}" ]; then
    echo "📍 AutoResume: Using explicitly specified checkpoint: ${RESUME_FROM_CHECKPOINT}"
fi

# Build deepspeed command
CMD=(
    deepspeed --num_gpus "${NGPUS}" 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}"
    --disable_group_texts "${DISABLE_GROUP_TEXTS}"
    --block_size "${BLOCK_SIZE}"
    --per_device_train_batch_size "${BATCH}"
    --gradient_accumulation_steps "${GRAD_ACC}"
    --run_name "${RUN_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}"
    --complementary_mask "${COMPLEMENTARY_MASK}"
    --flexible_bd_size "${FLEXIBLE_BD_SIZE}"
    --use_block_causal_mask "${USE_BLOCK_CAUSAL_MASK}"
    --entropy_loss "${ENTROPY_LOSS}"
    --entropy_loss_weight "${ENTROPY_LOSS_WEIGHT}"
    --block_causal_no_dynamic "${BLOCK_CAUSAL_NO_DYNAMIC}"
    --deepspeed "${DS_CONFIG}"
)

# Add resume_from_checkpoint if available
if [ -n "${RESUME_FROM_CHECKPOINT}" ]; then
    CMD+=( --resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}" )
    echo "📍 Resuming from checkpoint: ${RESUME_FROM_CHECKPOINT}"
else
    # Only add overwrite_output_dir if NOT resuming
    CMD+=( --overwrite_output_dir )
    echo "🆕 Starting fresh training (will overwrite output dir)"
fi

if [ "${BF16}" = "1" ]; then
    CMD+=( --bf16 )
fi
if [ "${GRADIENT_CHECKPOINTING}" = "1" ]; then
    CMD+=( --gradient_checkpointing )
fi

echo ""
echo "=========================================="
echo "Starting Deepspeed training..."
echo "=========================================="

# Run training with proper error handling
set +e
"${CMD[@]}" 2>&1 | tee "${LOG_DIR}/train.log"
EXIT_CODE=${PIPESTATUS[0]}
set -e

echo ""
echo "=========================================="
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

else
    echo "❌ ShareGPT4V VLM Training Failed (exit code: ${EXIT_CODE})"
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log\n  Error: workdir/train_sharegpt4v_vlm_error.log"
fi
echo "=========================================="

# Handle requeue logic (optional: add auto-resubmit on failure if desired)
SCRIPT_PATH="${PROJECT_ROOT}/chengyue_script/dvlm/train_sharegpt4v_vlm_autoresume.sh"
if [ "${REQUEUED:-0}" = "1" ]; then
    echo "Job was requeued via signal. Exit code: ${EXIT_CODE}"
elif [ "$EXIT_CODE" -ne 0 ]; then
    echo "Training failed with exit code: ${EXIT_CODE}"
    # Optionally add auto-resubmit logic here if needed
    # echo "Resubmitting due to non-zero exit..."
    # sbatch --dependency=afterany:$SLURM_JOB_ID --export=ALL ${SCRIPT_PATH}
fi

exit $EXIT_CODE


