1 files changed, 21 insertions, 0 deletions
diff --git a/scripts/run_trm_sudoku.sh b/scripts/run_trm_sudoku.sh
new file mode 100755
index 0000000..ae6db21
--- /dev/null
+++ b/scripts/run_trm_sudoku.sh
@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+# 启动 TRM Sudoku 1k 训练 (TRM 官方 pretrain_mlp_t_sudoku 配置)
+# 单 GPU L40S 48GB 约 18h; A6000 应该接近
+set -euo pipefail
+REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+source "$(conda info --base)/etc/profile.d/conda.sh"
+conda activate rrm
+
+cd "$REPO_ROOT/trm"
+OMP_NUM_THREADS=${OMP_NUM_THREADS:-8} \
+WANDB_MODE=${WANDB_MODE:-online} \
+python pretrain.py \
+  arch=trm \
+  data_paths="[$REPO_ROOT/data/sudoku-extreme-1k-aug-1000]" \
+  evaluators="[]" \
+  epochs=50000 eval_interval=5000 \
+  lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
+  arch.mlp_t=True arch.pos_encodings=none \
+  arch.L_layers=2 \
+  arch.H_cycles=3 arch.L_cycles=6 \
+  +run_name=pretrain_mlp_t_sudoku ema=True "$@"