scripts/train_mllmseg_gres.sh
set -x

GPUS=4
BATCH_SIZE=32 # 16 for internvl2
PER_DEVICE_BATCH_SIZE=4
GRADIENT_ACC=$((BATCH_SIZE / PER_DEVICE_BATCH_SIZE / GPUS))
LORA_RANK=64
BASE_MODEL="OpenGVLab/InternVL2_5-8B"
PRETRAINED_MODEL="checkpoints/MLLMSeg_InternVL25_RES"
arr=(${BASE_MODEL//\// })
MODEL_NAME=${arr[-1]}
TOTAL_STEPS=20000

export PYTHONPATH="${PYTHONPATH}:$(pwd)"
export MASTER_PORT=$(shuf -i 20000-65000 -n 1)
export TF_CPP_MIN_LOG_LEVEL=3
export LAUNCHER=pytorch

OUTPUT_DIR=work_dirs/mllmseg_${MODEL_NAME}_lora${LORA_RANK}_b${BATCH_SIZE}_g${GPUS}_gres

if [ ! -d "$OUTPUT_DIR" ]; then
  mkdir -p "$OUTPUT_DIR"
fi
cp "$0" "$OUTPUT_DIR/"

torchrun --nnodes=1 --node_rank=0 --master_addr=127.0.0.1 --nproc_per_node=${GPUS} --master_port=${MASTER_PORT} \
  mllmseg/finetune_mllmseg.py \
  --model_name_or_path ${BASE_MODEL} \
  --pretrained_model ${PRETRAINED_MODEL} \
  --train_dataset "refer_seg" \
  --refer_seg_data "grefcoco" \
  --output_dir ${OUTPUT_DIR} \
  --overwrite_output_dir True \
  --freeze_llm True \
  --freeze_mlp True \
  --freeze_backbone True \
  --use_llm_lora ${LORA_RANK} \
  --dataloader_num_workers 4 \
  --bf16 True \
  --num_train_epochs 1 \
  --total_steps ${TOTAL_STEPS} \
  --ce_loss_weight 1.0 \
  --focal_loss_weight 1.0 \
  --dice_loss_weight 1.0 \
  --per_device_train_batch_size ${PER_DEVICE_BATCH_SIZE} \
  --gradient_accumulation_steps ${GRADIENT_ACC} \
  --evaluation_strategy "no" \
  --save_strategy "steps" \
  --save_steps 1000 \
  --save_total_limit 2 \
  --learning_rate 4e-5 \
  --weight_decay 0.05 \
  --warmup_ratio 0.03 \
  --lr_scheduler_type "cosine" \
  --logging_steps 1000 \
  --do_train True \
  --ps_version 'v2' \
  --deepspeed "zero_stage1_config.json" \
  --report_to "tensorboard" \
  2>&1 | tee -a "${OUTPUT_DIR}/training_log.txt"
