#!/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_llava80k_vlm
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_llava80k_vlm_error.log
#SBATCH -o workdir/train_llava80k_vlm_out.log

set -e

echo "=========================================="
echo "MDM_Qwen2VL VLM Training (LLaVA 80K)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "GPUs: 8"
echo "CPUs: 16"
echo "Memory: 1024GB"
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 - <<'PY' 2>/dev/null || true
import torch
print(torch.cuda.is_available())
PY
)"
echo "nvidia-smi:"
nvidia-smi || true
echo ""

# Set paths
export PYTHONPATH="$PWD:$PWD/chengyue_script:$PWD/src:$PYTHONPATH"
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

# 配置
exp_id=train_llava80k_vlm_${SLURM_JOB_ID:-$(date +%Y%m%d_%H%M%S)}
model_name_or_path=/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_MDM
tokenizer_name=Qwen/Qwen2.5-VL-3B-Instruct

# Multi-modal dataset (json + images)
dataset_json="/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/data/llava_instruct_80k.json"
image_folder="/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/data/coco2017/train2017"

output_dir=output_models/${exp_id}
log_dir=log/${exp_id}

echo "=========================================="
echo "Configuration:"
echo "  Model: ${model_name_or_path}"
echo "  Tokenizer: ${tokenizer_name}"
echo "  Dataset JSON: ${dataset_json}"
echo "  Image Folder: ${image_folder}"
echo "  Output: ${output_dir}"
echo "  Num Train Epochs: 1"
echo "  Sequence Length: 2048"
echo "  per_device_train_batch_size: 4"
echo "  gradient_accumulation_steps: 2"
echo "  Effective batch size: 4 × 2 × 8 GPUs = 64"
echo "  Precision: bf16"
echo "  bd_size: 128"
echo "  learn_padding: 1"
echo "  conplemenrary_mask: 1"
echo "=========================================="
echo ""

# Create directories
mkdir -p ${output_dir} ${log_dir}

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

if [ ! -f "${dataset_json}" ]; then
    echo "❌ Dataset json not found: ${dataset_json}"
    exit 1
fi
echo "✅ Dataset json: ${dataset_json}"

if [ ! -d "${image_folder}" ]; then
    echo "❌ Image folder not found: ${image_folder}"
    exit 1
fi
echo "✅ Image folder: ${image_folder}"
echo ""

echo "Starting VLM training (LLaVA 80K, with images)..."
echo "=========================================="
echo ""

# Keep vision modules when loading base VLM (MDM naming)
export MDM_VL_KEEP_VISION=1

# Run training (single node, 8 GPUs)
deepspeed --num_gpus 8 examples/finetune_qwenvl.py \
    --model_name_or_path ${model_name_or_path} \
    --tokenizer_name ${tokenizer_name} \
    --conversation_template qwen2_5_no_reasoning \
    --trust_remote_code 1 \
    --dataset_path ${dataset_json} \
    --image_folder ${image_folder} \
    --output_dir ${output_dir} \
    --num_train_epochs 1 \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio 0.03 \
    --learning_rate 2e-5 \
    --disable_group_texts 1 \
    --block_size 2048 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 2 \
    --gradient_checkpointing \
    --run_name ${exp_id} \
    --overwrite_output_dir \
    --logging_steps 50 \
    --do_train \
    --validation_split_percentage 0 \
    --ddp_timeout 72000 \
    --save_steps 1000 \
    --save_total_limit 20 \
    --dataloader_num_workers 8 \
    --preprocessing_num_workers 32 \
    --bd_size 128 \
    --mdm 1 \
    --learn_padding 1 \
    --conplemenrary_mask 1 \
    --bf16 \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    2>&1 | tee ${log_dir}/train.log

exit_code=${PIPESTATUS[0]}

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ VLM Training (LLaVA 80K) Success!"
else
    echo "❌ VLM Training (LLaVA 80K) Failed (exit code: ${exit_code})"
    echo -e "Please check logs:\n  Output: ${log_dir}/train.log\n  Error: workdir/train_llava80k_vlm_error.log"
fi
echo "=========================================="

exit $exit_code


