add slurm template

Signed-off-by: Ye Yu <yeyu@nvidia.com>

Author: yeyu-nvidia

examples/speculative_decoding/launch_train.sh

@@ -106,7 +106,7 @@ while [ $# -gt 0 ]; do
       if [[ "$1" != *=* ]]; then shift; fi
       NUM_NODES="${1#*=}"
       ;;
-    --head_node_ip*
+    --head_node_ip*)
       if [[ "$1" != *=* ]]; then shift; fi
       HEAD_NODE_IP="${1#*=}"
       ;;
@@ -199,7 +199,7 @@ if [[ "$NUM_NODES" != 1 ]]; then
                    --num_machines $NUM_NODES \
                    --machine_rank $SLURM_PROCID \
                    --rdzv_backend c10d \
-                   --main_process_ip $HEAD_NODE_IP \ 
+                   --main_process_ip $HEAD_NODE_IP \
                    --main_process_port 29500"
 else
   MULTI_NODE_ARGS=""

examples/speculative_decoding/slurm.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+#SBATCH -A {account}
+#SBATCH --job-name={job_name}
+#SBATCH --nodes={num_nodes} --ntasks-per-node=1 --gpus-per-node={num_gpus_per_node}
+#SBATCH -p {partition}
+#SBATCH -t {time_limit}
+
+CONTAINER_IMAGE={container_image}
+WORK_DIR={path_to_modelopt}
+
+CONTAINER_MOUNT="${WORK_DIR}:/modelopt"
+
+OUTPUT_DIR={path_to_output_dir}
+MODEL={path_to_model_dir}
+DATA={path_to_data_dir}
+OFFLINE_DATA={path_to_offline_data_dir}
+
+CMD="./launch_train.sh --model $MODEL \
+            --output_dir $OUTPUT_DIR \
+            --data $DATA \
+            --num_epochs 1 \
+            --train_bs 1 \
+            --lr 1e-4 \
+            --eagle_config eagle_config.json \
+            --training_seq_len 4096 \
+            --save_steps 1000 \
+            --estimate_ar True \
+            --disable_tqdm True \
+            --offline-data $OFFLINE_DATA \
+            --num_nodes $SLURM_NNODES \
+            --head_node_ip $head_node_ip \
+"
+
+srun -l \
+    --mpi=pmix \
+    --output=%x_%j_$DATETIME.log \
+    --container-workdir "/modelopt/examples/speculative_decoding" \
+    --container-image ${CONTAINER_IMAGE} --container-mounts ${CONTAINER_MOUNT} \
+    bash -lc "$CMD"
+
+set +x