#!/bin/bash
set -e

# Setup
export CUDA_VISIBLE_DEVICES="1"
export NUM_GPUS=$(python -c "import torch; print(torch.cuda.device_count())")

# Training
export MODELS_DIR="/nfsclient/os5800/vlm_project/data/pllum-vlm/models"
export BASE_MODEL="${MODELS_DIR}/pretrained/pllum-12b-nc-instruct-250715-llava-next-pretrained-2025-08-22_11-46-57"
export RUN_NAME="$(basename ${BASE_MODEL} | sed 's/-pretrained.*//')-finetuned-$(date +%Y-%m-%d_%H-%M-%S)"

export DATA_DIR="TODO: replace with your directory"
export DATA_PATH="${DATA_DIR}/processed/instruct"
export IMAGES_PATH="${DATA_DIR}/static/instruct"

export PER_DEVICE_TRAIN_BATCH_SIZE=1
export GLOBAL_BATCH_SIZE=128
export GRADIENT_ACCUMULATION_STEPS=$((GLOBAL_BATCH_SIZE / (NUM_GPUS * PER_DEVICE_TRAIN_BATCH_SIZE)))

export LAUNCHER="accelerate launch \
    --num_processes=${NUM_GPUS} \
    --num_machines=$SLURM_NNODES \
    --mixed_precision=bf16 \
    --dynamo_backend=no \
    "
export SCRIPT="-m src.modeling.llava.train"
export SCRIPT_ARGS="\
    --deepspeed=config/deepspeed/zero3.json \
    --model_name_or_path=${BASE_MODEL} \
    --data_path=${DATA_PATH} \
    --images_path=${IMAGES_PATH} \
    --trainer_output_dir=${MODELS_DIR}/checkpoints \
    --model_output_dir=${MODELS_DIR}/finetuned \
    --run_name=${RUN_NAME} \
    --seed=42 \
    --train_vision_tower=True \
    --train_multimodal_projector=True \
    --train_language_model=True \
    --max_seq_length=8192 \
    --epochs=1 \
    --lr=2e-5 \
    --vision_tower_lr=2e-6 \
    --lr_scheduler_type=cosine \
    --weight_decay=0. \
    --warmup_ratio=0.03 \
    --per_device_train_batch_size=${PER_DEVICE_TRAIN_BATCH_SIZE} \
    --per_device_eval_batch_size=4 \
    --gradient_accumulation_steps=${GRADIENT_ACCUMULATION_STEPS} \
    --gradient_checkpointing=True \
    --save_steps=100 \
    --logging_steps=10 \
    --report_to=wandb \
    --eval_strategy=steps \
    --eval_steps=100 \
    --use_peft=True \
    --lora_r=128 \
    --lora_alpha=256 \
    --lora_dropout=0.05 \
    --lora_target_modules=q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj \
    --lora_modules_to_save=multi_modal_projector,vision_tower \
    --lora_layers_to_transform=
    "

bash -c "${LAUNCHER} ${SCRIPT} ${SCRIPT_ARGS}"

python -m src.modeling.llava.merge \
    --base_model_name_or_path="${BASE_MODEL}" \
    --trainer_output_dir="${MODELS_DIR}/checkpoints/${RUN_NAME}" \
    --final_model_name_or_path="${MODELS_DIR}/finetuned/${RUN_NAME}"