DOCSEARCHbyAlgolia
Submit Docs

Hierarchical Reasoning Model (HRM)

repository·main·Indexed 12 days ago

https://github.com/sapientinc/hrm

A recurrent architecture for complex, goal-oriented reasoning tasks using a two-module system of high-level planning and low-level computation. Designed for efficiency and minimal training data, it includes implementations for solving Sudoku, Maze, and the Abstraction and Reasoning Corpus (ARC) datasets. Requires CUDA 12.6, PyTorch, and FlashAttention.

Tokens
2.7K
Snippets
10
Records
11
Agent score
46%

What's inside HRM

  1. Evaluate trained models

    main

    To evaluate your models, you can check metrics in Weights & Biases or run the evaluation script.

    1. Check Metrics: Look for eval/exact_accuracy in your W&B dashboard.
    2. Run Evaluation Script: For ARC-AGI, use the following command:
    OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 evaluate.py checkpoint=<CHECKPOINT_PATH>
    1. Final Inspection: Use the provided arc_eval.ipynb notebook to finalize and inspect your results.
    OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 evaluate.py checkpoint=<CHECKPOINT_PATH>

  2. Install the Hierarchical Reasoning Model (HRM) environment

    main

    To use HRM, you must set up a CUDA 12.6 environment with PyTorch and FlashAttention. The repository requires CUDA extensions to be built.

    1. Install CUDA 12.6

    CUDA_URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run

wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer.run $CUDA_URL
sudo sh cuda_installer.run --silent --toolkit --override

export CUDA_HOME=/usr/local/cuda-12.6

    2. Install PyTorch with CUDA 12.6

    PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu126
pip3 install torch torchvision torchaudio --index-url $PYTORCH_INDEX_URL

    3. Install Build Dependencies

    pip3 install packaging ninja wheel setuptools setuptools-scm

    4. Install FlashAttention

    Depending on your GPU architecture:

    • For Hopper GPUs:
      git clone git@github.com:Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install
    • For Ampere or earlier GPUs:
      pip3 install flash-attn

    5. Install Python Dependencies

    pip install -r requirements.txt
    # Install CUDA 12.6
CUDA_URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run

wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer.run $CUDA_URL
sudo sh cuda_installer.run --silent --toolkit --override

export CUDA_HOME=/usr/local/cuda-12.6

  3. Prepare datasets for full-scale experiments

    main

    Before running large-scale experiments, initialize submodules and build the required datasets using the provided scripts.

    Initialize submodules:

    git submodule update --init --recursive

    Build ARC Datasets:

    • ARC-1 (Official + ConceptARC): python dataset/build_arc_dataset.py
    • ARC-2 (Official): python dataset/build_arc_dataset.py --dataset-dirs dataset/raw-data/ARC-AGI-2/data --output-dir data/arc-2-aug-1000

    Build Sudoku/Maze Datasets:

    • Sudoku-Extreme (Full): python dataset/build_sudoku_dataset.py
    • Sudoku-Extreme (1k samples): python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000 --subsample-size 1000 --num-aug 1000
    • Maze: python dataset/build_maze_dataset.py
    git submodule update --init --recursive

# ARC-1
python dataset/build_arc_dataset.py

# ARC-2
python dataset/build_arc_dataset.py --dataset-dirs dataset/raw-data/ARC-AGI-2/data --output-dir data/arc-2-aug-1000

# Sudoku-Extreme
python dataset/build_sudoku_dataset.py

# Maze
python dataset/build_maze_dataset.py

  4. Launch full-scale training experiments

    main

    Experiments below assume an 8-GPU setup using torchrun.

    Small-sample (1K) Training

    • ARC-1: OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py (~24h)
    • ARC-2: OMP_NUM_THREADS=8 torchrun --nproc-per-node 8 pretrain.py data_path=data/arc-2-aug-1000 (~24h)
    • Sudoku Extreme (1k): OMP_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.0 (~10m)
    • Maze 30x30 Hard (1k): OMP_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.0 (~1h)

    Full Sudoku-Hard Training

    OMP_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=learned

    Runtime: ~2 hours*

    OMP_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=learned

  5. Visualize datasets

    main

    You can explore generated puzzles visually by using the built-in visualizer:

    1. Locate the generated dataset folder in data/....
    2. Open puzzle_visualizer.html in your web browser.
    3. Upload the dataset folder to the visualizer.

  6. Visualize ARC predictions

    main

    When calling test(visualize=True), the system generates a Matplotlib figure displaying:

    1. Input: The original puzzle input grid.
    2. Answer: The ground truth answer grid.
    3. Trials: The top $K$ predicted grids (based on the 2-vote logic).

    Grids are rendered using ARC_COLOR_MAP, which maps symbol values (0-9) to specific colors (e.g., 0: black, 1: blue, 2: red).

    # To generate visual plots of inputs, answers, and top predictions
test(visualize=True)

  7. Evaluate ARC model predictions

    main

    Use the test function to evaluate the performance of a Hierarchical Reasoning Model on the ARC (Abstraction and Reasoning Corpus) dataset. The function loads puzzle identifiers and model predictions, processes augmented data back to its original form, and calculates accuracy based on different 'K-shot' voting strategies (where the top $K$ predicted answers are considered).

    Key Workflow:

    1. Data Loading: Loads identifiers.json from the dataset path and model prediction files (matching the pattern {CHECKPOINT_PATH}_all_preds.*) from the checkpoint path.
    2. Preprocessing: Removes padding using PAD_PUZZLE_IDENTIFIER and crops grids using the crop function to remove EOS tokens.
    3. De-augmentation: Uses inverse_aug to revert augmented grids (rotations/reflections) back to their original orientation for comparison with ground truth.
    4. Voting Logic: For each puzzle, it aggregates predictions and calculates a confidence score using q_halt_logits. Accuracy is reported for $K ext{-shot}$ scenarios (e.g., 1-shot, 2-shot, 10-shot).
    # Example usage for evaluation
DATASET_PATH = "data/arc-aug-1000"
CHECKPOINT_PATH = "checkpoints/Arc-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 amphibian-turaco/step_414456"

# Run evaluation without visualization
test(visualize=False)

  8. Run a Quick Demo: Sudoku Solver

    main

    You can train a master-level Sudoku AI on a modern laptop GPU. This process involves downloading/building the dataset and then running the training script.

    Estimated Runtime: ~10 hours on an RTX 4070 laptop GPU.

    1. Download and build Sudoku dataset:
    python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000
    1. Start training (single GPU, smaller batch size):
    OMP_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.0
    python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000

OMP_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.0

  9. Revert augmented grids with `inverse_aug`

    main

    If the model was trained on augmented data (rotations, reflections, permutations), use inverse_aug to transform a grid back to its original state based on the augmentation metadata in the puzzle name.

    Parameters:

    • name (str): The augmentation name (e.g., name_t1_p[0,1,2]). If no underscore is present, the grid is returned unchanged.
    • grid (np.ndarray): The augmented grid.

    Logic: It extracts the transformation ID and permutation from the name, applies the inverse_dihedral_transform, and then applies the inverse permutation.

    # name contains transformation and permutation info
original_grid = inverse_aug(name, augmented_grid)

  10. Load identifiers and predictions with `load_identifiers_and_preds`

    main

    The load_identifiers_and_preds function synchronizes dataset metadata with model output files.

    Parameters:

    • dataset_path (str): Path to the directory containing identifiers.json.
    • checkpoint_path (str): Prefix for the prediction files (e.g., path/to/checkpoint will look for path/to/checkpoint_all_preds.*).

    Returns:

    • identifier_map (dict): Mapping from identifiers to puzzle names.
    • all_preds (dict): A dictionary containing concatenated PyTorch tensors for keys like puzzle_identifiers, inputs, labels, logits, and q_halt_logits.
    identifier_map, all_preds = load_identifiers_and_preds(DATASET_PATH, CHECKPOINT_PATH)

  11. Crop grids to remove padding with `crop`

    main

    The crop function uses Numba-accelerated logic to find the largest rectangle within a grid that does not contain any EOS (End Of Sequence) tokens.

    Parameters:

    • grid (np.ndarray): The input grid (expected to be reshaped to 30x30 internally).

    Returns:

    • np.ndarray: The cropped grid with EOS tokens removed (the function subtracts 2 from the values to handle the EOS token representation).
    @njit
def crop(grid: np.ndarray):
    # Returns the maximum-sized rectangle without EOS tokens
    ...