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

CURRENT_DIR=$( cd "$( dirname "$0" )" && pwd )
MEGATRON_PATH=$( dirname $( dirname ${CURRENT_DIR}))

# Those variables need to modify
DTK_ENV="/opt/dtk/env.sh"                                                               # where env.sh of dtk
DATA_PATH="/workspace/dcu_megatron-core_v0.15.0/data/oscar-1GB_head-qwen_text_document"                           # path to mmap_qwen3_datasets_text_document
TOKENIZER_MODEL_PATH="/workspace/dcu_megatron-core_v0.15.0/qwen3_dataset"                                                  # path to qwen3_dataset
CHECKPOINT_PATH="./ckpt"                                                       # path to ckpt
NCCL_ENV=${MEGATRON_PATH}/requirements/env.sh                            # Please adjust the variables based on the actual NET being used
LAUNCH_WITH_BINDING=${MEGATRON_PATH}/requirements/launch_with_binding.sh # Please adjust the variables based on the actual NET being used

# Those variables no need to modify
HOSTFILE="hostfile_$(basename "$0" | sed -E 's/^run_(.+)\.sh$/\1/')"
GPUS=$(($(cat ${HOSTFILE}|sort|uniq |wc -l)*8))
HOST="$(cat ${HOSTFILE} |sed -n "1p"|awk -F ' ' '{print $1}')"
PORT="25900"
train_type="pretrain"
if [ ${train_type} = "pretrain" ]; then
    shell=${CURRENT_DIR}/train_qwen3_30B_A3B_pretrain.sh
elif [ ${train_type} = "sft" ]; then
    shell=train_qwen3_30B_A3B_sft.sh
fi

# Runs Qwen3 30B A3B model
source ${NCCL_ENV}
mpirun -np ${GPUS}  --hostfile ${HOSTFILE} \
                    --allow-run-as-root \
                    --bind-to none \
                    --mca plm_rsh_no_tree_spawn 1 \
                    --mca plm_rsh_args "-p 0211" \
                    bash -c "
                    source ${DTK_ENV} && \
                    source ${NCCL_ENV} && \
                    ${shell} \
                    ${HOST} \
                    ${PORT} \
                    --data_path=$DATA_PATH \
                    --tokenizer_path=$TOKENIZER_MODEL_PATH \
                    --checkpoint_path=$CHECKPOINT_PATH \
                    --launch_with_binding=${LAUNCH_WITH_BINDING} \
                    --profiling=$profiling" 
                    #> log-$((${GPUS} / 8))nodes-`date +%F-%H%M`.log 2>&1

wait