for para in $*
do
    if [[ $para == --profiling* ]];then
        profiling=${para#*=}
    fi
done

# Runs GPT 567B model
source /opt/dtk/env.sh
HOST=""  # modify this variable
PORT=25900
DATA_PATH="path to redpajama_text_document"
TOKENIZER_MODEL_PATH="path to tokenizer.model"
CHECKPOINT_PATH="path to ckpt" 


mpirun -np 1024 --hostfile hostfile_gpt_567B \
                --allow-run-as-root \
                --bind-to none \
                --mca plm_rsh_no_tree_spawn 1 \
                train_gpt_567B_multinodes.sh \
                ${HOST} \
                ${PORT} \
                --data_path=$DATA_PATH \
                --tokenizer_path=$TOKENIZER_MODEL_PATH \
                --checkpoint_path=$CHECKPOINT_PATH \
                --profiling=$profiling > log-1024nodes-`date +%F-%H%M`.log 2>&1

wait