Views
No views yet

1# Install CUDA 12.6
2CUDA_URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run
3
4wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer.run $CUDA_URL
5sudo sh cuda_installer.run --silent --toolkit --override
6
7export CUDA_HOME=/usr/local/cuda-12.6
8
9# Install PyTorch with CUDA 12.6
10PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu126
11
12pip3 install torch torchvision torchaudio --index-url $PYTORCH_INDEX_URL
13
14# Additional packages for building extensions
15pip3 install packaging ninja wheel setuptools setuptools-scm1git clone git@github.com:Dao-AILab/flash-attention.git
2cd flash-attention/hopper
3python setup.py installpip3 install flash-attnpip install -r requirements.txtwandb login1# Download and build Sudoku dataset
2python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000 --subsample-size 1000 --num-aug 1000
3
4# Start training (single GPU, smaller batch size)
5OMP_NUM_THREADS=8 python pretrain.py data_path=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.01# Initialize submodules
2git submodule update --init --recursive
3
4# ARC-1
5python dataset/build_arc_dataset.py # ARC offical + ConceptARC, 960 examples
6# ARC-2
7python dataset/build_arc_dataset.py --dataset-dirs dataset/raw-data/ARC-AGI-2/data --output-dir data/arc-2-aug-1000 # ARC-2 official, 1120 examples
8
9# Sudoku-Extreme
10python dataset/build_sudoku_dataset.py # Full version
11python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000 --subsample-size 1000 --num-aug 1000 # 1000 examples
12
13# Maze
14python dataset/build_maze_dataset.py # 1000 examplespuzzle_visualizer.html in your browser.data/....OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/arc-2-aug-1000OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/sudoku-extreme-1k-aug-1000 epochs=20000 eval_interval=2000 lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/maze-30x30-hard-1k epochs=20000 eval_interval=2000 lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/sudoku-hard-full epochs=100 eval_interval=10 lr_min_ratio=0.1 global_batch_size=2304 lr=3e-4 puzzle_emb_lr=3e-4 weight_decay=0.1 puzzle_emb_weight_decay=0.1 arch.loss.loss_type=softmax_cross_entropy arch.L_cycles=8 arch.halt_max_steps=8 arch.pos_encodings=learnedeval/exact_accuracy in W&B.OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 evaluate.py checkpoint=<CHECKPOINT_PATH>arc_eval.ipynb notebook to finalize and inspect your results.1@misc{wang2025hierarchicalreasoningmodel,
2 title={Hierarchical Reasoning Model},
3 author={Guan Wang and Jin Li and Yuhao Sun and Xing Chen and Changling Liu and Yue Wu and Meng Lu and Sen Song and Yasin Abbasi Yadkori},
4 year={2025},
5 eprint={2506.21734},
6 archivePrefix={arXiv},
7 primaryClass={cs.AI},
8 url={https://arxiv.org/abs/2506.21734},
9}