finetune_multi_node.sh

#!/bin/bash
export

# --- 1. Get model and dataset name from command line ---
BASE_MODEL=$1 # qwen_3b
DATASET=$2 # citeseer
MP=$3 # mp layer
C_RATIO=0.1 # compression ratio
LORA_RANK=64
LORA_ALPHA=64
LORA_DROPOUT=0.1

EPOCH=3

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

###### setting for model ######
echo "Finetuning on model: $BASE_MODEL"

case "$BASE_MODEL" in
  qwen_7b)
    BASE_MODEL_DIR="./output/pretrain/${BASE_MODEL}/pt_academic_mp_${MP}_ratio_${C_RATIO}/"
    ;;
  # Add more model mappings here...
  # *) is the default case for handling unknown model names
  *)
    echo "Error: Unknown model name '$BASE_MODEL'!"
    echo "Available options: qwen_7b"
    exit 1
    ;;
esac
###### setting for model ######

echo ">>>>>>>>>>>>> Set Base Model: $BASE_MODEL_DIR >>>>>>>>>>>>>"

###### setting for dataset ######
echo "Finetuing on dataset: $DATASET"

OUTPUT_DIR="./output/finetune/${BASE_MODEL}/${DATASET}/"
FINAL_OUTPUT_DIR="./final_output/finetune/${BASE_MODEL}/${DATASET}/model/"

# ###### scan existing ckpt ######
# if ls -d "${OUTPUT_DIR}/checkpoint"*/ >/dev/null 2>&1; then
#     echo "=========================================================="
#     echo "Warning: Found existing checkpoint directory in ${OUTPUT_DIR}"
#     echo "Script terminated automatically to avoid overwriting existing training progress"
#     echo "=========================================================="
#     exit 1
# fi
# ###### scan existing ckpt ######

case "$DATASET" in
  cora)
    TRAIN_DATA="./data_preprocess/cora/processed/finetuned_${DATASET}_v1.json"
    EVAL_DATA="./data_preprocess/cora/processed/finetuned_${DATASET}_val_v1.json"
    ;;
  *)
    echo "Error: Unknown dataset '$DATASET'!"
    ;;
esac
###### setting for dataset ######

WORLD_SIZE=${WORLD_SIZE:-1}
NODE_RANK=${RANK:-0}
MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
MASTER_PORT=${MASTER_PORT:-12346}

gpus=$(nvidia-smi -L | wc -l)
torchrun --nproc_per_node=$gpus --nnode=$WORLD_SIZE --node_rank=$NODE_RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT -m main.train \
  --output_dir $OUTPUT_DIR \
  --final_output_dir $FINAL_OUTPUT_DIR \
  --model_name_or_path $BASE_MODEL_DIR \
  --train_data $TRAIN_DATA \
  --eval_data $EVAL_DATA \
  --load_best_model_at_end True \
  --eval_steps 20 \
  --max_length 8192 \
  --min_length 3072 \
  --max_train_num_per_data 200000 \
  --num_train_epochs $EPOCH \
  --enable_beacon True \
  --beacon_param q k v o \
  --training_stage finetune \
  --mp_layer_num $MP \
  --c_ratio $C_RATIO \
  --lora_rank $LORA_RANK \
  --lora_alpha $LORA_ALPHA \
  --lora_dropout $LORA_DROPOUT \
  --group_by_length \
  --gradient_checkpointing \
  --save_strategy steps \
  --save_steps 20 \
  --save_total_limit 2 \
  --logging_steps 1 \
  --dataset_cache_dir ./cache \
  --model_cache_dir ./cache