pretrain_multi_node.sh

#!/bin/bash
export

BASE_MODEL=$1 # qwen_7b
MP=$2 # mp layer
C_RATIO=$3 # compression ratio default=0.1

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

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

case "$BASE_MODEL" in
  qwen_7b)
    BASE_MODEL_DIR="./model/Qwen2.5-7B-Instruct/"
    ;;
  *)
    echo "❌ Unknown Model Name '$BASE_MODEL'！"
    echo "Available Options: qwen_7b"
    exit 1
    ;;
esac
###### setting for model ######

###### setting for dataset ######
echo "Pretrain on dataset: Academic"

OUTPUT_DIR="./output/pretrain/${BASE_MODEL}/pt_academic_mp_${MP}_ratio_${C_RATIO}/"

echo "Model Output dir: $OUTPUT_DIR"

TRAIN_DATA1="./data/pretrained_data/pretrained_Economics.json"
TRAIN_DATA2="./data/pretrained_data/pretrained_Mathematics.json"
TRAIN_DATA3="./data/pretrained_data/pretrained_Geology.json"
###### setting for dataset ######

EPOCH=1

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 \
  --model_name_or_path $BASE_MODEL_DIR \
  --train_data $TRAIN_DATA1 $TRAIN_DATA2 $TRAIN_DATA3 \
  --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 pretrain \
  --mp_layer_num $MP \
  --c_ratio $C_RATIO \
  --gradient_checkpointing \
  --save_strategy steps \
  --save_steps 50 \
  --save_total_limit 2 \
  --logging_steps 1 \
  --dataset_cache_dir ./cache \
  --model_cache_dir ./cache