| # Training Guide: HRM SRAM/DRAM Learning Performance |
|
|
| Use this guide to train the memory-tiered HRM on **real datasets** and verify its reasoning accuracy. |
|
|
| --- |
|
|
| ### 1. Dataset Generation |
| The model needs to learn how to solve puzzles. First, generate a training/test set (e.g., Sudoku). |
|
|
| ```bash |
| # Activate environment |
| source venv/bin/activate |
| |
| # Build a small Sudoku dataset (1000 examples) |
| python dataset/build_sudoku_dataset.py \ |
| --output-dir data/sudoku-1k \ |
| --subsample-size 1000 \ |
| --num-aug 10 |
| ``` |
|
|
| --- |
|
|
| ### 2. Start Training |
| Point the trainer to the tiered architecture configuration. |
|
|
| #### Train the Tiered Model (New) |
| ```bash |
| python pretrain.py \ |
| arch=hrm_tiered \ |
| data_path=data/sudoku-1k \ |
| epochs=1000 \ |
| global_batch_size=384 |
| ``` |
|
|
| #### Train the Baseline Model (Comparison) |
| ```bash |
| python pretrain.py \ |
| arch=hrm_v1 \ |
| data_path=data/sudoku-1k \ |
| epochs=1000 \ |
| global_batch_size=384 |
| ``` |
|
|
| --- |
|
|
| ### 3. Verify Reasoning Accuracy |
| Once training is complete (or during training), check the accuracy metrics in your W&B dashboard or via the evaluation script: |
|
|
| ```bash |
| # Replace with the path to your generated checkpoint |
| python evaluate.py checkpoint=checkpoints/Sudoku_ACT-torch/YOUR_RUN/step_1000 |
| ``` |
|
|
| **What to look for:** |
| - **`eval/exact_accuracy`**: Fraction of puzzles solved perfectly. |
| - **`eval/steps`**: Average number of reasoning steps taken by the ACT (Adaptive Computation Time) module. |
| - **`eval/q_halt_accuracy`**: How well the model learns when to stop thinking. |
| |
| --- |
| |
| ### 4. Training on ARC (Artificial General Intelligence Benchmark) |
| To train on the more complex ARC-AGI-2 dataset: |
| |
| ```bash |
| python dataset/build_arc_dataset.py --output-dir data/arc-2 |
| python pretrain.py arch=hrm_tiered data_path=data/arc-2 |
| ``` |
| |