HRM-checkpoint-sudoku-full

HRM Checkpoint — Sudoku Full

Checkpoint from training the Hierarchical Reasoning Model (HRM) on the full Sudoku Extreme dataset, following similar setup to sapientinc/HRM-checkpoint-sudoku-extreme.

python3 ./pretrain.py \
  data_path=data/sudoku-extreme-full \
  epochs=100 \
  eval_interval=100 \
  lr_min_ratio=0.1 \
  global_batch_size=1152 \
  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

I tried to mimic the file structure in sapientinc/HRM-checkpoint-sudoku-extreme, but I figured I'd add some extra stats:

This has the output evaluate.py typically produces across several max_steps settings, in an easier to read JSON format: evaluate-Sudoku-extreme-full.json.

I also ran it in a loop where I whittled down the set to only ones that were unsolved. You can see my method in run_subset.py.. It produces stats.json. That's what I'm graphing below.

And here's a graph of that data, somewhat like Figure 5c in the Hierarchical Reasoning Model paper:

(I should say that even though the graph shows at M_max=1024 exact accuracy being at 100%, it's not really 100%. It's 99.9605%: Which corresponds to 422,619 correct of 422,786 total sudokus. Or 167 unsolved sudokus.)

Perhaps it would be useful to see the results as a table.

👉 Browse the Sudoku puzzles solved by HRM (grouped by required inference steps) here: Sudoku Puzzle Collection

Usage

You should be able to run it as

HRM_LOCATION="/tmp/hrm" # Or wherever
CHECKPOINT_LOCATION="/tmp/HRM-checkpoint-sudoku-full" # Or wherever, of course.

git clone https://github.com/sapientinc/HRM "${HRM_LOCATION}"
# Running this, requires a bunch of configuration. Obviously Sapient has their
# own README.md, etc. But I've made a docker image that you might be able to
# use as a guide as well. I'll link it below.

git clone https://huggingface.co/bnsh/HRM-checkpoint-sudoku-full/ "${CHECKPOINT_LOCATION}"

cd "${HRM_LOCATION}"
python3 ./evaluate.py checkpoint="${CHECKPOINT_LOCATION}/checkpoint" data_path=data/sudoku-extreme-full/

And, here's that Docker image I mentioned: bnsh/hrm-docker (setup and usage guide).

Training Details

• Hardware: NVIDIA A10g
• Runtime: ≈ 9 days, 3 hours, 34 minutes, 25 seconds (13174m 24.845s)
• Parameters: ~27.3M

Final Metrics

Run History (ASCII plots)

num_params ▁
train/accuracy ▂▁▂▁▁▃▄▄▄▄▅▅▅▆▅▆▇▅▇▆▇▆▇▆▇▇▇▇▇▇██████████
train/count ▁███████████████████████████████████████
train/exact_accuracy ▁▁▂▂▃▄▄▅▅▅▅▆▆▆▆▇▇▇▇▇▇▇▇▇▇█▇█▇███████████
train/lm_loss ██▇▇▇▇▆▆▆▆▅▅▅▅▅▅▅▅▅▄▄▄▅▄▄▄▃▄▄▄▃▃▃▂▁▂▂▁▁▁
train/lr ██████████▇▇▇▆▆▆▆▆▆▆▅▄▄▄▄▄▃▃▃▃▂▂▁▁▁▁▁▁▁▁
train/q_continue_loss ▁▄▃█▃▅▅▆▅▅▅▆▆▆▅▆▅▅▆▆▅▅▅▄▅▅▄▅▅▄▄▄▄▃▃▄▃▃▃▂
train/q_halt_accuracy █▂██▁▄█▅███████████▆████████████████████
train/q_halt_loss ▁▃▇▁▄▆▄▆▂▄▅▄▇▄▇▄▄▃▆▅▇▃▂▆█▇▆█▅▄▄▆▆▅▄▆▇▅▇▆
train/steps █▇▇█▆▅▅▅▄▇▄▃▃▃▃▃▃▂▃▂▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁█▁▁▁

Reference

Reference: Hierarchical Reasoning Model (HRM), Arxiv:2506.21734