#!/bin/bash
# Finetune J63 MDM Ministral on j63_recipe (single node, DeepSpeed ZeRO-2)

# Parses arguments
model_name_or_path=${MODEL_PATH:-src/models/Nemotron-Diffusion-Final-merged}
tokenizer_name=${TOKENIZER_NAME:-src/models/Nemotron-Diffusion-Final-merged}
dataset_path=${DATASET_PATH:-j63_recipe}
output_dir=${OUTPUT_DIR:-output_models/post_train_j63_mdm_ministral_final_template_unfreeze_vision_len2048_new}
deepspeed_args="--master_port=$(python3 -c 'import random; print(random.randint(20000,65000))')"
conversation_template=ministral
trust_remote_code=1

# export CUDA_HOME=${CUDA_HOME:-/cm/shared/apps/cuda12.4/toolkit/12.4.1}
export CUDA_HOME=/home/chengyuew/chengyuew/miniconda3/envs/dvlm
export MDM_VL_KEEP_VISION=1
export TOKENIZERS_PARALLELISM=false
# export VILA_DATASETS=cw-dfw-cs
export DATASET_BACKEND=webdataset_multi_modal
export WANDB_DISABLED=true
export MINISTRAL_OVERFLOW_POLICY=${MINISTRAL_OVERFLOW_POLICY:-vision_dummy}
export MINISTRAL_FORCE_VISION=${MINISTRAL_FORCE_VISION:-1}

# 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_unfreeze_vision_len1024
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_ministral.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 50 \
    --dataloader_num_workers 12 \
    --learn_padding 1 \
    --bd_size 32 \
    --use_block_causal_mask 1 \
    --mdm 1 \
    --complementary_mask 1 \
    --save_total_limit 20 \
    --max_length 2048 \
    --gradient_checkpointing 1 \
    --freeze_vision_encoder 0 \
    --block_causal_no_dynamic 1 \
    --report_to none \
    > >(tee ${log_dir}/train.log) \
    2> >(tee ${log_dir}/train.err >&2)"

echo $cmd
eval $cmd
