1 files changed, 54 insertions, 0 deletions
diff --git a/run_maze_a6000.sh b/run_maze_a6000.sh
new file mode 100755
index 0000000..787647a
--- /dev/null
+++ b/run_maze_a6000.sh
@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+# Clone-and-run Maze training, defaulted for 2x A6000 (48G). Repo-relative; no hardcoded paths.
+# Usage:
+#   bash run_maze_a6000.sh              # 2x A6000, gbs 384 (default)
+#   bash run_maze_a6000.sh 2 384        # explicit
+#   bash run_maze_a6000.sh 1 192        # single A6000
+#   SMOKE=1 bash run_maze_a6000.sh      # 200-step smoke test first (minutes)
+set -eo pipefail
+
+REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+NGPU="${1:-2}"
+GBS="${2:-384}"
+DATA="$REPO/data/maze-30x30-hard-1k"
+TRM="$REPO/trm"
+RUN_NAME="pretrain_att_maze30x30_${NGPU}gpu_gbs${GBS}"
+
+# conda env (default 'rrm'); override with CONDA_ENV / CONDA_SH
+CONDA_SH="${CONDA_SH:-}"
+[[ -z "$CONDA_SH" ]] && for p in "$HOME/miniconda3/etc/profile.d/conda.sh" \
+  "$HOME/anaconda3/etc/profile.d/conda.sh" "/opt/conda/etc/profile.d/conda.sh"; do
+  [[ -f "$p" ]] && CONDA_SH="$p" && break; done
+[[ -f "$CONDA_SH" ]] && source "$CONDA_SH" && conda activate "${CONDA_ENV:-rrm}"
+
+[[ -d "$DATA" ]] || { echo "FATAL: dataset missing at $DATA"; exit 1; }
+[[ -f "$TRM/pretrain.py" ]] || { echo "FATAL: trm/pretrain.py missing"; exit 1; }
+cd "$TRM"
+export WANDB_MODE=offline
+
+EPOCHS=50000; EVAL=5000
+if [[ "${SMOKE:-0}" == "1" ]]; then EPOCHS=200; EVAL=200; RUN_NAME="smoke_${RUN_NAME}"; fi
+
+ARGS=(
+  arch=trm
+  "data_paths=[$DATA]"
+  "evaluators=[]"
+  epochs=$EPOCHS eval_interval=$EVAL
+  lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0
+  global_batch_size="$GBS"
+  arch.L_layers=2 arch.H_cycles=3 arch.L_cycles=4
+  +run_name="$RUN_NAME" ema=True
+  +checkpoint_every_eval=true
+)
+LOG="$REPO/maze_${RUN_NAME}.log"
+
+echo "launching $RUN_NAME on $NGPU GPU(s), gbs=$GBS  (epochs=$EPOCHS)"
+if [[ "$NGPU" -gt 1 ]]; then
+  nohup torchrun --nproc-per-node "$NGPU" --rdzv_backend=c10d --rdzv_endpoint=localhost:0 \
+    --nnodes=1 pretrain.py "${ARGS[@]}" > "$LOG" 2>&1 &
+else
+  nohup python pretrain.py "${ARGS[@]}" > "$LOG" 2>&1 &
+fi
+echo "pid $!   log: $LOG"
+echo "watch:  tail -f $LOG | grep -E 'exact|accuracy|Traceback|Error'"
+echo "ckpts:  $TRM/checkpoints/maze-30x30-hard-1k.../$RUN_NAME/  (1 per ${EVAL} epochs)"