#!/bin/bash
#SBATCH -A nvr_elm_llm  
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 02:00:00
#SBATCH -N 1
#SBATCH -J preprocess_dvlm_dataset
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/preprocess_%j.err
#SBATCH -o workdir/preprocess_%j.out
#SBATCH --exclusive

set -eo pipefail

echo "[preprocess] Start at $(date)"
echo "[preprocess] Single node preprocessing with 64 workers to generate dataset cache"

# activate environment (temporarily disable -u for conda)
set +u
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda
set -u

# parameters (must match training script for cache compatibility)
model_name_or_path=/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_DFM
dataset_path=/home/chengyuew/chengyuew/workspace/LMFlow/data/nemotron_no_thinking_subset
conversation_template=qwen2_5_no_reasoning
block_size=2048
trust_remote_code=1

# output (temporary, just for preprocessing)
output_dir=output_models/preprocess_temp
log_dir=log/preprocess_temp
mkdir -p ${output_dir} ${log_dir}

echo "[preprocess] Dataset: ${dataset_path}"
echo "[preprocess] Using 64 preprocessing workers (8 GPUs x 8 workers each)"
echo "[preprocess] Cache will be stored in ~/.cache/huggingface/datasets"

# Run preprocessing with max_steps=1 (just to trigger dataset processing)
deepspeed --num_gpus 8 \
  examples/finetune.py \
  --model_name_or_path ${model_name_or_path} \
  --tokenizer_name Qwen/Qwen2.5-VL-3B-Instruct \
  --conversation_template ${conversation_template} \
  --trust_remote_code ${trust_remote_code} \
  --dataset_path ${dataset_path} \
  --output_dir ${output_dir} \
  --max_steps 1 \
  --learning_rate 1e-5 \
  --lr_scheduler_type constant_with_warmup \
  --warmup_ratio 0.03 \
  --disable_group_texts 0 \
  --block_size ${block_size} \
  --per_device_train_batch_size 8 \
  --gradient_accumulation_steps 1 \
  --deepspeed configs/ds_config_zero3_no_offload.json \
  --bf16 \
  --run_name preprocess_test \
  --validation_split_percentage 0 \
  --logging_steps 1 \
  --do_train \
  --save_steps 9999999 \
  --dataloader_num_workers 8 \
  --preprocessing_num_workers 1 \
  --bd_size 128 \
  --gradient_checkpointing 1 \
  --mdm 1 \
  --learn_padding 1 \
  --conplemenrary_mask 1 \
  2>&1 | tee ${log_dir}/preprocess.log

echo "[preprocess] Completed at $(date)"
echo "[preprocess] Check cache: ls ~/.cache/huggingface/datasets/"
ls -lh ~/.cache/huggingface/datasets/ | tail -n 20

# Clean up temporary output
echo "[preprocess] Cleaning up temp output..."
rm -rf ${output_dir}

echo "[preprocess] Done! Dataset cache generated for multi-node training."

