Update sbatch_8b7_178b_25b_jz_tmp.sh

sbatch_8b7_178b_25b_jz_tmp.sh

@@ -28,7 +28,10 @@ source $six_ALL_CCFRWORK/start-tr13f-6B3-ml-t0
 GPUS_PER_NODE=8
 NNODES=$SLURM_NNODES
 
-DATA_PATH="XXXX"
+TRAIN_DATA_PATH=train55boscar.txt
+# "train: 1.0 0:1 /scratch/project_462000119/data/c4_subsampled/gpt2tok_c4_en_55B_text_document"
+VALID_DATA_PATH=val.txt
+# "validation: 1.0 0:1 /scratch/project_462000119/data/c4_validation/gpt2tok_c4validation_rerun_text_document"
 
 
 
@@ -171,7 +174,24 @@ echo $CMD
 
 echo "START $SLURM_JOBID: $(date)"
 
+
+### JZ ###
+export LAUNCHER="python -u -m torch.distributed.run \
+    --nproc_per_node $GPUS_PER_NODE \
+    --nnodes $NNODES \
+    --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT \
+    --rdzv_backend c10d \
+    --max_restarts 0 \
+    --tee 3 \
+    "
+SRUN_ARGS=" \
+    --wait=60 \
+    --kill-on-bad-exit=1 \
+    "
+srun $SRUN_ARGS --jobid $SLURM_JOBID bash -c "$LAUNCHER --node_rank \$SLURM_PROCID $CMD" 2>&1 | tee -a $LOGS_PATH/main_log.txt
+
+### LUMI ###
 # bash launch_srun.sh $CMD
-srun --label launch.sh $CMD
+# srun --label launch.sh $CMD
 
 echo "END $SLURM_JOBID: $(date)"