#!/bin/bash
#SBATCH -A nvr_elm_llm  
#SBATCH -p backfill_singlenode
#SBATCH -t 01:00:00
#SBATCH -N 1
#SBATCH -J dvlm_gpu_quick
#SBATCH --gpus-per-node 8
#SBATCH --ntasks-per-node 8
## no exclusive to avoid cancellation
#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_gpu_quick
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
trust_remote_code=1
ds_config=configs/ds_config_zero3_no_offload.json

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

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

which python3 && python3 -V || true
nvidia-smi -L || true

echo "[driver] launching torchrun 8-GPU training at $(date)"
set -x
DISABLE_TORCH_COMPILE=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True TORCHINDUCTOR_FREEZING=1 TORCHINDUCTOR_COMPILE_THREADS=1 torch_compile_backend=none torchrun --nproc_per_node=8 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} \
  --conversation_template qwen2_5 \
  --max_steps 1 \
  --max_train_samples 100 \
  --learning_rate 2e-5 \
  --disable_group_texts 1 \
  --block_size 128 \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 1 \
  --deepspeed ${ds_config} \
  --gradient_checkpointing \
  --bf16 \
  --run_name ${exp_id} \
  --overwrite_output_dir \
  --validation_split_percentage 0 \
  --logging_steps 1 \
  --do_train \
  --save_steps 1000 \
  --dataloader_num_workers 1 \
  --preprocessing_num_workers 1 \
  --bd_size 64 \
  --mdm 1 \
  --mdm_loss 1 \
  > >(tee ${log_dir}/train.log) \
  2> >(tee ${log_dir}/train.err >&2)
set +x

echo "[driver] end at $(date)"
