blob: 04d3a1ca6c6cf40a87957fd5932a47db38131259 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
|
#!/usr/bin/env bash
# 启动 HRM Sudoku 1k 训练 (HRM 官方推荐配置)
# 单 GPU 约 10h on RTX 4070; 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/hrm"
OMP_NUM_THREADS=${OMP_NUM_THREADS:-8} \
WANDB_MODE=${WANDB_MODE:-online} \
python pretrain.py \
data_path="$REPO_ROOT/data/sudoku-extreme-1k-aug-1000" \
epochs=20000 eval_interval=2000 global_batch_size=384 \
lr=7e-5 puzzle_emb_lr=7e-5 weight_decay=1.0 puzzle_emb_weight_decay=1.0 "$@"
|
