#!/bin/bash
#SBATCH -A nvr_lpr_llm
#SBATCH -p batch_block1
#SBATCH -t 04:00:00
#SBATCH -N 8
#SBATCH --ntasks-per-node=1
#SBATCH -J train_J63_fastdllmv2_8node_bsz256_anneal_bd32_block_causal_lr_1e-6_unfreeze_vision
#SBATCH --gpus-per-node 8
#SBATCH --array=1-30%1
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_J63_fastdllmv2_8node_bsz256_anneal_bd32_block_causal_lr_1e-6_unfreeze_vision_error.log
#SBATCH -o workdir/train_J63_fastdllmv2_8node_bsz256_anneal_bd32_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 "=========================================="
    
    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 $TRAIN_PID 2>/dev/null || true
    fi
    
    echo "Requesting SLURM to requeue job..."
    scontrol requeue $SLURM_JOB_ID
    exit 0
}

trap handle_timeout USR1

# Resolve project root
PROJECT_ROOT="${SLURM_SUBMIT_DIR:-$PWD}"
cd "${PROJECT_ROOT}"

# ============================================
# MULTI-NODE CONFIGURATION (8 NODES)
# ============================================
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 VLM Training (J63 FastDLLMV2)"
echo "=========================================="
echo "  NNODES: ${NNODES}  NGPUS_PER_NODE: ${NGPUS_PER_NODE}"
echo "  MASTER_ADDR: ${MASTER_ADDR}  MASTER_PORT: ${MASTER_PORT}"
echo "  Node: $(hostname)  Date: $(date)  Job: ${SLURM_JOB_ID:-N/A}"
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__)')"
nvidia-smi || true
echo ""

# ============================================
# Environment variables (aligned with LLM script)
# ============================================
export PYTHONPATH="${PROJECT_ROOT}/src:${PROJECT_ROOT}/third_party/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 TOKENIZERS_PARALLELISM=false

# ============================================
# Training parameters (aligned with LLM script where possible)
# ============================================
MODEL_PATH="${MODEL_PATH:-${PROJECT_ROOT}/output_models/j63_math_8node_bsz256_anneal_bd32_block_causal_lr_1e-6_unfreeze_vision/checkpoint-7400}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"
MAX_LENGTH="${MAX_LENGTH:-2048}"

# ShareGPT4V dataset configuration
DATASET_PATH="${DATASET_PATH:-j63_small_fastdllmv2}"
TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"

EPOCHS="${EPOCHS:-1}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-4}"                # 64 GPUs * 1 * 4 = 256 effective
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:-8}"                  # Was 1, aligned with LLM=8
BLOCK_SIZE="${BLOCK_SIZE:-2048}"

# MDM specific
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-32}"
LEARN_PADDING="${LEARN_PADDING:-1}"
COMPLEMENTARY_MASK="${COMPLEMENTARY_MASK:-1}"
ALWAYS_MASK_IM_END="${ALWAYS_MASK_IM_END:-1}"
USE_BLOCK_CAUSAL_MASK="${USE_BLOCK_CAUSAL_MASK:-1}"
ANNEAL_BLOCK_SIZE="${ANNEAL_BLOCK_SIZE:-0}"
FLEXIBLE_BD_SIZE="${FLEXIBLE_BD_SIZE:-0}"

# Activation checkpointing (saves memory)
ACTIVATION_CHECKPOINTING="${ACTIVATION_CHECKPOINTING:-1}"

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

mkdir -p "${OUTPUT_DIR}" "${LOG_DIR}"

# ============================================
# Check model path
# ============================================
if [ ! -d "${MODEL_PATH}" ]; then
    echo "ERROR: Model path does not exist: ${MODEL_PATH}"
    exit 1
fi

# ============================================
# 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)
fi

resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
    echo "Resuming from: ${latest_checkpoint}"
else
    resume_arg="--overwrite_output_dir"
    echo "Training from scratch"
fi

ac_arg=""
if [ "${ACTIVATION_CHECKPOINTING}" = "1" ]; then
    ac_arg="--gradient_checkpointing"
fi

echo ""
echo "=========================================="
echo "Config:"
echo "  Model: ${MODEL_PATH}"
echo "  Dataset: ${DATASET_PATH}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Epochs: ${EPOCHS}  Batch: ${BATCH}  GradAcc: ${GRAD_ACC}"
echo "  LR: ${LR}  MDM: ${MDM}  BD_SIZE: ${BD_SIZE}"
echo "  BlockCausal: ${USE_BLOCK_CAUSAL_MASK}  Complementary: ${COMPLEMENTARY_MASK}"
echo "  TotalGPUs: $((NNODES * NGPUS_PER_NODE))  EffBatch: $((NNODES * NGPUS_PER_NODE * BATCH * GRAD_ACC))"
echo "  Workers: ${WORKERS}"
echo "=========================================="
echo ""

# ============================================
# Launch training
# ============================================
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}"'/third_party/vila:'"${PROJECT_ROOT}"'/examples"
        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}"' \
            --complementary_mask '"${COMPLEMENTARY_MASK}"' \
            --bf16 \
            --fsdp "shard_grad_op" \
            --always_mask_im_end '"${ALWAYS_MASK_IM_END}"' \
            --flexible_bd_size '"${FLEXIBLE_BD_SIZE}"' \
            --use_block_causal_mask '"${USE_BLOCK_CAUSAL_MASK}"' \
            --anneal_block_size '"${ANNEAL_BLOCK_SIZE}"' \
            --validation_split_percentage 0 \
            '"${ac_arg}"' \
            '"${resume_arg}"'
    ' 2>&1 | tee "${LOG_DIR}/train_${SLURM_ARRAY_TASK_ID}-of-${SLURM_ARRAY_TASK_COUNT}.log" &
TRAIN_PID=$!

wait $TRAIN_PID
exit_code=$?

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "Training completed successfully!"
    
    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"
    fi
else
    echo "Training failed with exit code: ${exit_code}"
    echo "Check logs: ${LOG_DIR}/"
fi
echo "=========================================="

exit $exit_code
