diff options
| author | YurenHao0426 <blackhao0426@gmail.com> | 2026-06-13 12:35:36 -0500 |
|---|---|---|
| committer | YurenHao0426 <blackhao0426@gmail.com> | 2026-06-13 12:35:36 -0500 |
| commit | 66e0d8b9fd4d0f7a2231d689c055e26fdf1cf04a (patch) | |
| tree | c29cba61124018755a19b02c9d33e3ad5f2e05cc /scripts/run_trm_sudoku.sh | |
Curated export for clone-and-run Maze training (2x A6000) + diagnostics.
trm/hrm pretrain.py carry trajectory-augmentation code (backward-compatible).
Heavy artifacts (checkpoints/wandb/npz) gitignored; see PROVENANCE.md.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Diffstat (limited to 'scripts/run_trm_sudoku.sh')
| -rwxr-xr-x | scripts/run_trm_sudoku.sh | 21 |
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 "$@" |