#!/bin/bash
#SBATCH -A nvr_elm_llm  
#SBATCH -p backfill_singlenode
#SBATCH -t 01:00:00
#SBATCH -N 1
#SBATCH -J dvlm_debug_1gpu_backfill
#SBATCH --gpus-per-node 1
#SBATCH -e workdir/%x_%j.err
#SBATCH -o workdir/%x_%j.out

set -eo pipefail

echo "[driver] start at $(date), host=$(hostname)"
echo "[driver] SLURM_JOB_ID=${SLURM_JOB_ID:-}"
echo "[driver] PWD=$(pwd)"

# job configuration
exp_id=dvlm_debug_simple_1gpu
dataset_path=/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/data/test_small
output_dir=output_models/${exp_id}
log_dir=log/${exp_id}
model_name_or_path=/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_DFM
ds_config=configs/ds_config_zero2_no_offload.json
trust_remote_code=1

# check dataset
if [ ! -d ${dataset_path} ]; then
  echo "[driver] ERROR: dataset path not found: ${dataset_path}" 1>&2
  exit 1
fi

# ensure output/log dirs
echo "[driver] creating output and log directories..."
mkdir -p ${output_dir} ${log_dir}

# check deepspeed config
if [ ! -f ${ds_config} ]; then
  echo "[driver] ERROR: deepspeed config not found: ${ds_config}" 1>&2
  exit 1
fi

# activate conda environment
echo "[driver] activating conda environment..."
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

echo "[driver] checking environment..."
which python3 && python3 -V
which deepspeed && deepspeed --version
nvidia-smi -L

# launch training (single node, no hostfile needed)
echo "[driver] launching training at $(date)"
set -x

deepspeed --num_gpus=1 examples/finetune.py \
  --model_name_or_path ${model_name_or_path} \
  --tokenizer_name Qwen/Qwen2.5-VL-3B-Instruct \
  --trust_remote_code ${trust_remote_code} \
  --dataset_path ${dataset_path} \
  --output_dir ${output_dir} \
  --max_steps 1 \
  --max_train_samples 100 \
  --learning_rate 2e-5 \
  --disable_group_texts 0 \
  --block_size 512 \
  --per_device_train_batch_size 2 \
  --gradient_accumulation_steps 1 \
  --deepspeed ${ds_config} \
  --bf16 \
  --conversation_template qwen2_5 \
  --run_name ${exp_id} \
  --validation_split_percentage 0 \
  --logging_steps 5 \
  --do_train \
  --ddp_timeout 72000 \
  --save_steps 500 \
  --dataloader_num_workers 4 \
  --preprocessing_num_workers 16 \
  --bd_size 128 \
  --mdm 1 \
  --mdm_loss 1 \
  > >(tee ${log_dir}/train.log) \
  2> >(tee ${log_dir}/train.err >&2)

set +x

echo "[driver] training completed at $(date)"

