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

# Those variables need to modify
GPUS=""                 # how many gpus to use
MPI_PORT=""             # mpi port to use
HOST=""                 # hostname
PORT=""                 # port id
DATA_PATH=""            # path to my-mixtral_text_document
TOKENIZER_MODEL_PATH="" # path to tokenizer.model
CHECKPOINT_PATH=""      # path to ckpt

# Runs Mixtral 8x22B model
mpirun -np ${GPUS}  --hostfile hostfile_mixtral_8x22B \
                    --allow-run-as-root \
                    --bind-to none \
                    --mca plm_rsh_no_tree_spawn 1 \
                    --mca plm_rsh_args "-p ${MPI_PORT}" \
                    bash -c "
                    ./train_mixtral_8x22B_$((${GPUS} / 8))nodes.sh \
                    ${HOST} \
                    ${PORT} \
                    --data_path=$DATA_PATH \
                    --tokenizer_path=$TOKENIZER_MODEL_PATH \
                    --checkpoint_path=$CHECKPOINT_PATH \
                    --profiling=$profiling" > log-$((${GPUS} / 8))nodes-`date +%F-%H%M`.log 2>&1

wait