#!/bin/bash
# Finetune MDM_Qwen2.5-VL on j63_math_mmfinereasoning (single node, DeepSpeed ZeRO-2)

# Parses arguments
model_name_or_path=${MODEL_PATH:-output_models/j63_small_fastdllmv2_8node_bd32_block_causal_anneal_lr_5e-6_ds}
tokenizer_name=${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}
dataset_path=${DATASET_PATH:-j63_recipe}
output_dir=${OUTPUT_DIR:-output_models/post_train_j63_mdm_small_fastdllmv2_8node_bd32_block_causal_anneal_lr_5e-6_ds}
deepspeed_args="--master_port=$(python3 -c 'import random; print(random.randint(20000,65000))')"
conversation_template=qwen2_5_no_reasoning
trust_remote_code=1

export CUDA_HOME=${CUDA_HOME:-/cm/shared/apps/cuda12.4/toolkit/12.4.1}
export MAX_PIXELS="${MAX_PIXELS:-1048576}"
# export MDM_VL_KEEP_VISION=1
export TOKENIZERS_PARALLELISM=false
# export OMP_NUM_THREADS="${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

# Activate environment
project_dir=$(cd "$(dirname $0)"/../../..; pwd)
source "${project_dir}/../../scripts/cluster_env.sh" 2>/dev/null || true
type lmflow_activate_conda &>/dev/null && lmflow_activate_conda

export PYTHONPATH="${project_dir}/../../src:${project_dir}/../../third_party/vila:${project_dir}/../../examples:${PYTHONPATH:-}"

# Auto-detect latest checkpoint for resume
latest_checkpoint=""
if [ -d "${output_dir}" ]; then
    latest_checkpoint=$(find "${output_dir}" -name "checkpoint-*" -type d | sort -V | tail -1)
    if [ -n "${latest_checkpoint}" ]; then
        echo "Found checkpoint: ${latest_checkpoint}"
    fi
fi

# Finetune
exp_id=finetune
log_dir=${project_dir}/log/${exp_id}
mkdir -p ${output_dir} ${log_dir}

resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
fi

cmd="deepspeed ${deepspeed_args} \
  examples/finetune_qwenvl.py \
    --model_name_or_path ${model_name_or_path} \
    --tokenizer_name ${tokenizer_name} \
    --trust_remote_code ${trust_remote_code} \
    --dataset_path ${dataset_path} \
    --output_dir ${output_dir} \
    ${resume_arg} \
    --conversation_template ${conversation_template} \
    --num_train_epochs 1 \
    --learning_rate 2e-5 \
    --block_size 8192 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    --bf16 \
    --logging_steps 1 \
    --do_train \
    --ddp_timeout 7200 \
    --save_steps 500 \
    --dataloader_num_workers 12 \
    --learn_padding 1 \
    --bd_size 32 \
    --use_block_causal_mask 1 \
    --mdm 1 \
    --always_mask_im_end 1 \
    --complementary_mask 1 \
    --save_total_limit 20 \
    --max_length 4096 \
    --gradient_checkpointing 1 \
    --block_causal_no_dynamic 1 \
    --report_to none
    > >(tee ${log_dir}/train.log) \
    2> >(tee ${log_dir}/train.err >&2)"

echo $cmd
eval $cmd
