blob: 787647ac007f1b7af28624cf5c79304b5657eac5 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
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)"
|
