merge_model.sh

#!/bin/bash
export

# --- 1. Get model and dataset name from command line ---
BASE_MODEL=$1 # qwen_7b
DATASET=$2 # cora
MP=$3 # 2
C_RATIO=0.1 # compression ratio

echo ">>>>>>>>>>>>> Important Parameters: >>>>>>>>>>>>>"
echo ">>>>>>>>>>>>> Base Model: $BASE_MODEL >>>>>>>>>>>>>"
echo ">>>>>>>>>>>>> Dataset: $DATASET >>>>>>>>>>>>>"
echo ">>>>>>>>>>>>> MP Layer Num: $MP >>>>>>>>>>>>>"
echo ">>>>>>>>>>>>> Compression RATIO: $C_RATIO >>>>>>>>>>>>>"

###### setting for model ######
echo "Merge model: $BASE_MODEL on $DATASET"

BASE_MODEL_DIR="./data/outputs/pretrain/${BASE_MODEL}/pt_academic_mp_${MP}_ratio_${C_RATIO}/"

###### setting for model ######


###### setting for dataset ######
echo "Saving model on dataset: $DATASET"

# --- 3. Auto find and select checkpoint ---
echo "====== Checkpoint Settings ======"

# Define parent directory for checkpoints
CHECKPOINT_PARENT_DIR="./data/outputs/finetune/${BASE_MODEL}/${DATASET}/"

echo "Searching checkpoints in: ${CHECKPOINT_PARENT_DIR}"

# Check if parent directory exists
if [ ! -d "$CHECKPOINT_PARENT_DIR" ]; then
    echo "Error: Checkpoint directory does not exist: $CHECKPOINT_PARENT_DIR"
    exit 1
fi

# Find all checkpoint-* directories and sort by version
# sort -V handles numbers like checkpoint-100 and checkpoint-900 correctly
# Use mapfile (or readarray) to read results into an array
mapfile -t sorted_checkpoints < <(find "${CHECKPOINT_PARENT_DIR}" -maxdepth 1 -type d -name "checkpoint-*" | sort -V)

# Check the number of checkpoints found
num_checkpoints=${#sorted_checkpoints[@]}
echo "Found ${num_checkpoints} checkpoint(s)."

# Check if enough checkpoints to select
if [ "$num_checkpoints" -lt 2 ]; then
  echo "Error: Found fewer than 2 checkpoints, cannot determine if model early-stopped properly"
  exit 1
fi

# Ascending order, early stop checkpoint index is 0
idx=0
LORA_PATH="${sorted_checkpoints[0]}"

echo "Success: Auto-selected early stop checkpoint (index 0): ${LORA_PATH}"
echo "==========================="
echo

# --- 4. Set output directory and execute merge script ---
OUTPUT_DIR="${CHECKPOINT_PARENT_DIR}/model/"

echo "Starting model merge..."
echo "LoRA path: ${LORA_PATH}"
echo "Final output directory: ${OUTPUT_DIR}"
echo

# Execute merge command
python -m merge \
  --model_name_or_path "${BASE_MODEL_DIR}" \
  --lora_path "${LORA_PATH}" \
  --final_output_dir "${OUTPUT_DIR}" \
  --max_length 15500 \
  --enable_beacon True \
  --local_window 3500 \
  --memory_stride 1024 \
  --mp_layer_num $MP \
  --c_ratio $C_RATIO \
  --beacon_param q k v o \

echo "Merge completed!"
###### setting for dataset ######