fix datapath issue of slurm run (#18)

Author: wenxie-amd

examples/megatron/run_pretrain.sh

@@ -159,14 +159,6 @@ if [ "$NODE_RANK" = "0" ]; then
     echo ""
 fi
 
-DISTRIBUTED_ARGS=(
-    --nproc_per_node "${GPUS_PER_NODE}"
-    --nnodes "${NNODES}"
-    --node_rank "${NODE_RANK}"
-    --master_addr "${MASTER_ADDR}"
-    --master_port "${MASTER_PORT}"
-)
-
 mkdir -p output
 TRAIN_LOG=output/log_torchrun_pretrain_${MODEL_CONFIG}.txt
 if [ "$NODE_RANK" = "0" ]; then
@@ -179,6 +171,14 @@ if [ "$RUN_ENV" = "torchrun" ]; then
     SITE_PACKAGES=$(python -c "import sysconfig; print(sysconfig.get_paths()['purelib'])")
     export PYTHONPATH=${SITE_PACKAGES}:${MEGATRON_PATH}:${PRIMUS_PATH}:${PYTHONPATH}
 
+    DISTRIBUTED_ARGS=(
+        --nproc_per_node "${GPUS_PER_NODE}"
+        --nnodes "${NNODES}"
+        --node_rank "${NODE_RANK}"
+        --master_addr "${MASTER_ADDR}"
+        --master_port "${MASTER_PORT}"
+    )
+
     # build helper_cpp of megatron
     pushd "${MEGATRON_PATH}/megatron/core/datasets" && make && popd || exit 1
 
@@ -228,6 +228,7 @@ elif [ "$RUN_ENV" = "slurm" ]; then
         --privileged --device=/dev/infiniband \
         -v $MEGATRON_PATH:$MEGATRON_PATH \
         -v $PRIMUS_PATH:$PRIMUS_PATH \
+        -v $DATA_PATH:$DATA_PATH \
         $DOCKER_IMAGE /bin/bash -c \
             "echo '[NODE-${NODE_RANK}]: begin, time=$(date +"%Y.%m.%d %H:%M:%S")' && \
             pip install -q loguru wandb && \
@@ -240,8 +241,8 @@ elif [ "$RUN_ENV" = "slurm" ]; then
                 --node_rank ${NODE_RANK} \
                 --master_addr ${MASTER_ADDR} \
                 --master_port ${MASTER_PORT} \
-                examples/deepseek/pretrain.py \
-                --exp examples/deepseek/exp_pretrain.yaml \
+                examples/megatron/pretrain.py \
+                --exp examples/megatron/exp_pretrain.yaml \
                 2>&1 | tee $TRAIN_LOG && \
             echo '[NODE-${NODE_RANK}]: end time=$(date +"%Y.%m.%d %H:%M:%S")'"
 else