From addeb31f4adfb758781a2874008c3885ac78da54 Mon Sep 17 00:00:00 2001
From: inter <inter@qq.com>
Date: Sun, 21 Sep 2025 20:19:36 +0800
Subject: [PATCH] Add File

---
 tools/scripts/slurm_train_v2.sh | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 tools/scripts/slurm_train_v2.sh

diff --git a/tools/scripts/slurm_train_v2.sh b/tools/scripts/slurm_train_v2.sh
new file mode 100644
index 0000000..d5ec0ed
--- /dev/null
+++ b/tools/scripts/slurm_train_v2.sh
@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+
+set -x
+
+PARTITION=$1
+JOB_NAME=$2
+GPUS=$3
+PY_ARGS=${@:4}
+
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+CPUS_PER_TASK=${CPUS_PER_TASK:-40}
+SRUN_ARGS=${SRUN_ARGS:-""}
+
+while true
+do
+    PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
+    status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
+    if [ "${status}" != "0" ]; then
+        break;
+    fi
+done
+echo $PORT
+
+srun -p ${PARTITION} \
+    --job-name=${JOB_NAME} \
+    --gres=gpu:${GPUS_PER_NODE} \
+    --cpus-per-task=${CPUS_PER_TASK} \
+    --kill-on-bad-exit=1 \
+    ${SRUN_ARGS} \
+    python -m torch.distributed.launch --nproc_per_node=${GPUS_PER_NODE} train.py --launcher pytorch --tcp_port ${PORT} ${PY_ARGS}